模型demo验证
更新时间：2025-05-15
背景信息
用户在训练过程中，需要定期对模型进行推理效果验证，我们提供了基于AIAK的推理效果验证方案，便于客户快速进行效果验证。下文以LLaMA模型为例：
注意： 您应自觉遵守“LLaMA”等第三方模型的用户协议、使用规范和相关法律法规，并就使用第三方模型的合法性、合规性自行承担相关责任。百度云不对第三方模型的合法性、安全性、准确性进行任何保证，百度云不对由此引发的任何损害承担责任。
使用方法
镜像源版本
Plain Text
1[someone@host]docker pull registry.baidubce.com/cce-ai-native/aiak-megatron:0612使用方法
使用镜像启动docker，准备推理测试环境：
Plain Text
1# 使用镜像启动容器，并挂载模型CKPT目录
2[someone@host]docker run -itd --name inference-test -v /path/to/model:/mnt/model -v /path/to/tokenizer:/mnt/tokenizer registry.baidubce.com/cce-ai-native/aiak-megatron:0612
3
4# 进入容器测试推理
5[someone@host]docker exec -it inference-test bash推理脚本适配（以llama-7b为例）
Plain Text
1[someone@host]cd Megatron-LM/
2[someone@host]vim examples/run_text_generation_server_llama_7b.sh
3# 配置1：CHECKPOINT位置(示例)：
4CHECKPOINT_PATH=/mnt/model/megatron_llama_7b_checkpoint_tp2_pp2
5
6# 配置2：TOKENIZER位置(示例)：
7TOKENIZER_PATH=/mnt/tokenizer/tokenizer.model
8
9# 配置3：llama-7b模型的参数设置：
10GPUS_PER_NODE=4 #请根据ckpt的rank切分数设置
11
12MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
13MASTER_PORT=${MASTER_PORT:-9999}
14NNODES=${1:-1}
15NODE_RANK=${2:-0}
16
17DISTRIBUTED_ARGS="--nproc_per_node $GPUS_PER_NODE
18 --nnodes $NNODES
19 --node_rank $NODE_RANK
20 --master_addr $MASTER_ADDR
21 --master_port $MASTER_PORT"
22
23CUDA_DEVICE_MAX_CONNECTIONS=1 torchrun $DISTRIBUTED_ARGS $ROOT_DIR/../tools/run_llama_text_generation_server.py \
24 --num-layers 32 \
25 --hidden-size 4096 \
26 --ffn-hidden-size 11008 \
27 --disable-bias-linear \
28 --num-attention-heads 32 \
29 --max-position-embeddings 2048 \
30 --no-position-embedding \
31 --tokenizer-type LLaMASentencePieceTokenizer \
32 --tokenizer-model ${TOKENIZER_PATH} \
33 --fp16 \
34 --micro-batch-size 1 \
35 --seq-length 1024 \
36 --out-seq-length 1024 \
37 --temperature 0.1 \
38 --top_p 0.9 \
39 --tensor-model-parallel-size 2 \
40 --pipeline-model-parallel-size 2 \
41 --load ${CHECKPOINT_PATH} \
42 --swiglu \
43 --use-rotary-position-embeddings \
44 --untie-embeddings-and-output-weights \
45 --attention-dropout 0 \
46 --hidden-dropout 0 \
47 --fused-rmsnorm \
48 --rmsnorm-epsilon 1e-6 \
49 --no-query-key-layer-scaling \
50 --seed 42启动模型的推理服务：
Plain Text
1#启动推理服务：
2[someone@host]nohup bash examples/run_text_generation_server_llama_7b.sh &> server.log &
3
4使用Client测试推理效果：
5[someone@host] python tools/text_generation_cli.py 127.0.0.1:5000
6Enter prompt: 1+1=2,2+2=
7Enter number of tokens to generate (default = 30):
8Enter temperature (default = 0.1):
9Megatron Response:
101+1=2,2+2=4,3+3=6,4+4=8,5+5=10,6+6=12其他模型推理适配
为了支持其他模型的推理服务，只需修改run_text_generation_server_llama_7b.sh脚本，具体如下：
- llama-13b模型
Plain Text
1[someone@host]cd Megatron-LM/
2[someone@host]vim examples/run_text_generation_server_llama_7b.sh
3# 配置1：CHECKPOINT位置(示例)：
4CHECKPOINT_PATH=/mnt/model/megatron_llama_13b_checkpoint_tp2_pp4
5
6# 配置2：TOKENIZER位置(示例)：
7TOKENIZER_PATH=/mnt/tokenizer/tokenizer.model
8
9# 配置3：llama-13b模型的参数设置：
10GPUS_PER_NODE=8 #请根据ckpt的rank切分数设置
11
12MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"}
13MASTER_PORT=${MASTER_PORT:-9999}
14NNODES=${1:-1}
15NODE_RANK=${2:-0}
16
17DISTRIBUTED_ARGS="--nproc_per_node $GPUS_PER_NODE
18 --nnodes $NNODES
19 --node_rank $NODE_RANK
20 --master_addr $MASTER_ADDR
21 --master_port $MASTER_PORT"
22
23CUDA_DEVICE_MAX_CONNECTIONS=1 torchrun $DISTRIBUTED_ARGS $ROOT_DIR/../tools/run_llama_text_generation_server.py \
24 --num-layers 40 \
25 --hidden-size 5120 \
26 --ffn-hidden-size 13824 \
27 --disable-bias-linear \
28 --num-attention-heads 40 \
29 --max-position-embeddings 2048 \
30 --no-position-embedding \
31 --tokenizer-type LLaMASentencePieceTokenizer \
32 --tokenizer-model ${TOKENIZER_PATH} \
33 --fp16 \
34 --micro-batch-size 1 \
35 --seq-length 2048 \
36 --out-seq-length 2048 \
37 --temperature 0.1 \
38 --top_p 0.9 \
39 --tensor-model-parallel-size 2 \
40 --pipeline-model-parallel-size 4 \
41 --load ${CHECKPOINT_PATH} \
42 --swiglu \
43 --use-rotary-position-embeddings \
44 --untie-embeddings-and-output-weights \
45 --attention-dropout 0 \
46 --hidden-dropout 0 \
47 --fused-rmsnorm \
48 --rmsnorm-epsilon 1e-6 \
49 --no-query-key-layer-scaling \
50 --seed 42