66 lines
2.1 KiB
Bash
Executable File
66 lines
2.1 KiB
Bash
Executable File
#!/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}"
|