对接ModelLink框架
前提条件
- 使用前请先了解MindIO TTP的约束限制章节。
- ModelLink框架准备参见ModelLink,匹配的Megatron-LM的版本为core_r0.6.0。

本次发布包配套以上链接的1.0.RC3分支,以发布时Commit ID:31100602dc2b420e9c323963fbaaf2e212d5eaf0为基准使用,环境及代码准备请参考ModelLink仓库的README.md。
操作步骤
- 编辑“torchrun”文件。
- 查找环境中的“torchrun”文件。
which torchrun
- 打开以上命令显示路径下的“torchrun”文件。
vim {torchrun文件路径}/torchrun
- 按“i”进入编辑模式,修改以下内容。
# 增加加粗内容 import re import sys import mindio_ttp.framework_ttp from torch.distributed.run import main as torch_main
- 按“Esc”键,输入:wq!,按“Enter”保存并退出编辑。
- 查找环境中的“torchrun”文件。
- 编辑预训练脚本(仅供参考)。
此处以编辑“examples/llama2/pretrain_llama2_7b_ptd.sh”脚本为例。
- 打开“examples/llama2/pretrain_llama2_7b_ptd.sh”脚本。
vim examples/llama2/pretrain_llama2_7b_ptd.sh
- 按“i”进入编辑模式,开启高可用功能需要增加脚本中加粗部分。
#!/bin/bash export CUDA_DEVICE_MAX_CONNECTIONS=1 export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True export GLOO_SOCKET_IFNAME=enp189s0f0 export TTP_OT=720 export TTP_ADDR="master node ip" source /usr/local/Ascend/ascend-toolkit/set_env.sh GPUS_PER_NODE=8 MASTER_ADDR=localhost MASTER_PORT=6000 NNODES=1 NODE_RANK=0 WORLD_SIZE=$(($GPUS_PER_NODE*$NNODES)) CKPT_SAVE_DIR="your model save ckpt path" DATA_PATH="your data path" TOKENIZER_MODEL="your tokenizer path" CKPT_LOAD_DIR="your model ckpt path" TP=1 PP=2 DISTRIBUTED_ARGS=" --nproc_per_node $GPUS_PER_NODE \ --nnodes $NNODES \ --node_rank $NODE_RANK \ --master_addr $MASTER_ADDR \ --master_port $MASTER_PORT " GPT_ARGS=" --tensor-model-parallel-size ${TP} \ --pipeline-model-parallel-size ${PP} \ --sequence-parallel \ --num-layers 32 \ --hidden-size 4096 \ --ffn-hidden-size 11008 \ --num-attention-heads 32 \ --tokenizer-type Llama2Tokenizer \ --tokenizer-model ${TOKENIZER_MODEL} \ --seq-length 4096 \ --max-position-embeddings 4096 \ --micro-batch-size 1 \ --global-batch-size 256 \ --make-vocab-size-divisible-by 1 \ --lr 1.25e-6 \ --train-iters 5000 \ --lr-decay-style cosine \ --untie-embeddings-and-output-weights \ --disable-bias-linear \ --attention-dropout 0.0 \ --init-method-std 0.01 \ --hidden-dropout 0.0 \ --position-embedding-type rope \ --normalization RMSNorm \ --use-fused-rmsnorm \ --swiglu \ --use-flash-attn \ --no-masked-softmax-fusion \ --attention-softmax-in-fp32 \ --min-lr 1.25e-7 \ --weight-decay 1e-1 \ --lr-warmup-fraction 0.01 \ --clip-grad 1.0 \ --adam-beta1 0.9 \ --initial-loss-scale 65536 \ --adam-beta2 0.95 \ --no-gradient-accumulation-fusion \ --no-load-optim \ --no-load-rng \ --use-distributed-optimizer \ --use-fused-swiglu \ --use-fused-rotary-pos-emb \ --overlap-grad-reduce \ --bf16 --enable-high-availability \ " DATA_ARGS=" --data-path $DATA_PATH \ --split 949,50,1 " OUTPUT_ARGS=" --log-interval 1 \ --save-interval 10000 \ --eval-interval 1000 \ --eval-iters 10 \ " torchrun $DISTRIBUTED_ARGS pretrain_gpt.py \ $GPT_ARGS \ $DATA_ARGS \ $OUTPUT_ARGS \ --distributed-backend nccl \ --jit-compile \ --load $CKPT_LOAD_DIR \ --save $CKPT_SAVE_DIR \ | tee logs/train_llama2_7b.log
高可用功能相关参数说明如下:
- GLOO_SOCKET_IFNAME:根据主节点高速网卡实际情况进行配置。
- TTP_OT:Torch框架等待CheckPoint保存时间,到时未保存完毕会强制退出进程。
- TTP_ADDR:集群主节点的IP地址。
- set_env.sh文件路径:请根据CANN实际的安装路径进行修改。
- enable-high-availability:故障快速恢复特性开关。默认关闭,配置后即开启临终遗言功能。
开启故障快速恢复特性开关后,各类优化器显存会发生变化,变化详情请参见表1。
对于分布式优化器而言,由于增加了优化器副本,导致静态内存有所增加。但是集群规模越大时,DP越大,平均到单卡的显存增加量很小,这样可以避免OOM,因此推荐在大集群中使用。根据显存情况选择开启与否,调节参数。
- 按“Esc”键,输入:wq!,按“Enter”保存并退出编辑。
- 打开“examples/llama2/pretrain_llama2_7b_ptd.sh”脚本。
父主题: 使用指导