#!/usr/bin/env bash
set -euo pipefail

# Container-side launcher. It applies the Megatron inference server patches
# documented in pretrain_kaiyuan2b before starting the server, because docker
# --rm containers do not retain in-place edits across runs.

bash scripts/laoyao_megatron_inference_setup.sh

MEGATRON_PATH=${MEGATRON_PATH:-/opt/Megatron-Bridge/3rdparty/Megatron-LM}
CHECKPOINT=${CHECKPOINT:-/work/runs/pretrain_8192_8gpu_dp8_mbs14_full_recompute_weighted_heldoutval_resume10000/checkpoints}
TOKENIZER_PATH=${TOKENIZER_PATH:-/work/tokenizer/glm5.2}
SERVER_PORT=${SERVER_PORT:-5000}
MASTER_PORT=${MASTER_PORT:-6000}
INFERENCE_MAX_SEQ_LENGTH=${INFERENCE_MAX_SEQ_LENGTH:-8192}

DISTRIBUTED_ARGS=(
  --nproc_per_node 1
  --nnodes 1
  --node_rank 0
  --master_addr localhost
  --master_port "${MASTER_PORT}"
)

export CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS:-1}

torchrun "${DISTRIBUTED_ARGS[@]}" "${MEGATRON_PATH}/tools/run_text_generation_server.py" \
  --load "${CHECKPOINT}" \
  --tensor-model-parallel-size 1 \
  --pipeline-model-parallel-size 1 \
  --num-layers 12 \
  --hidden-size 1536 \
  --ffn-hidden-size 4608 \
  --num-attention-heads 24 \
  --num-query-groups 4 \
  --group-query-attention \
  --seq-length 8192 \
  --max-position-embeddings 8192 \
  --position-embedding-type rope \
  --rotary-base 10000 \
  --swiglu \
  --disable-bias-linear \
  --normalization RMSNorm \
  --make-vocab-size-divisible-by 1 \
  --vocab-size 154856 \
  --tokenizer-type HuggingFaceTokenizer \
  --tokenizer-model "${TOKENIZER_PATH}" \
  --bf16 \
  --micro-batch-size 1 \
  --num-experts 12 \
  --moe-ffn-hidden-size 6144 \
  --moe-layer-freq '[0,0,1,0,1,0,1,0,1,0,1,0]' \
  --moe-router-topk 4 \
  --moe-router-load-balancing-type aux_loss \
  --moe-aux-loss-coeff 0.02 \
  --moe-z-loss-coeff 0.001 \
  --moe-token-dispatcher-type alltoall \
  --moe-expert-capacity-factor 1.25 \
  --moe-grouped-gemm \
  --no-load-rng \
  --no-load-optim \
  --exit-on-missing-checkpoint \
  --inference-max-requests 1 \
  --inference-max-seq-length "${INFERENCE_MAX_SEQ_LENGTH}" \
  --port "${SERVER_PORT}"
