昇腾社区首页
中文
注册

对接ModelLink框架

前提条件

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

本次发布包配套以上链接的1.0.RC3分支,以发布时Commit ID:31100602dc2b420e9c323963fbaaf2e212d5eaf0为基准使用,环境及代码准备请参考ModelLink仓库的README.md。

操作步骤

  1. 编辑“torchrun”文件。
    1. 查找环境中的“torchrun”文件。
      which torchrun
    2. 打开以上命令显示路径下的“torchrun”文件。
      vim {torchrun文件路径}/torchrun
    3. 按“i”进入编辑模式,修改以下内容。
      # 增加加粗内容
      import re
      import sys
      import mindio_ttp.framework_ttp
      from torch.distributed.run import main as torch_main
    4. 按“Esc”键,输入:wq!,按“Enter”保存并退出编辑。
  2. 编辑预训练脚本(仅供参考)。

    此处以编辑“examples/llama2/pretrain_llama2_7b_ptd.sh”脚本为例。

    1. 打开“examples/llama2/pretrain_llama2_7b_ptd.sh”脚本。
      vim examples/llama2/pretrain_llama2_7b_ptd.sh
    2. 按“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,因此推荐在大集群中使用。根据显存情况选择开启与否,调节参数。

      表1 原生优化器与开启故障快速恢复特性后优化器参数的理论数值变化

      优化器

      原生

      开启故障快速恢复特性

      说明

      fp16/bf16

      20

      20

      -

      fp32

      16

      16

      -

      fp16/bf16 Distributed

      4+16/d

      4+16*N/d

      • d:DP Group Size
      • N:副本数,N<d
    3. 按“Esc”键,输入:wq!,按“Enter”保存并退出编辑。