Torch对接MindSpeed-LLM框架

前提条件

本次发布包配套MindSpeed-LLM的1.0.0分支,环境、代码、数据集准备请用户参考MindSpeed-LLM仓库的相关指导说明,并确保其安全性。

操作步骤

  1. 使用业务用户登录到计算节点。

    业务用户非{MindIO-install-user}、HwHiAiUser、hwMindX用户,由用户根据实际情况决定。

  2. 进入MindSpeed-LLM安装目录。

    cd MindSpeed-LLM/

  3. 修改pretrain_gpt.py文件。

    1. 打开pretrain_gpt.py文件。
      vim pretrain_gpt.py
    2. 按“i”进入编辑模式,在文件头部找到from mindspeed_llm import megatron_adaptor,另起一行加入加粗内容。
      from mindspeed_llm import megatron_adaptor
      import torch_mindio
    3. 按“Esc”键,输入:wq!,按“Enter”保存并退出编辑。

  4. 编辑预训练脚本(仅供参考)。

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

    1. 打开“examples/legacy/llama2/pretrain_llama2_7b_ptd.sh”脚本。
      vim examples/legacy/llama2/pretrain_llama2_7b_ptd.sh
    2. 按“i”进入编辑模式,开启周期性CheckPoint加速功能需要增加脚本中加粗部分。
      #!/bin/bash
      
      export CUDA_DEVICE_MAX_CONNECTIONS=1
      export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
      
      export MINDIO_AUTO_PATCH_MEGATRON=true
      export GLOO_SOCKET_IFNAME=enp189s0f0
      export LD_LIBRARY_PATH=/usr/local/Ascend/driver/lib64/driver:$LD_LIBRARY_PATH
      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
      "
      
      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

      周期性CheckPoint加速功能相关参数说明如下:

      • MINDIO_AUTO_PATCH_MEGATRONMindIO ACP框架自动patch Megatron的源码,用于开启加速周期性CheckPoint特性。
      • GLOO_SOCKET_IFNAME:根据主节点高速网卡实际情况进行配置。
      • LD_LIBRARY_PATH:CANN包驱动的so库地址,请根据CANN实际的安装路径进行修改。
      • set_env.sh文件路径:请根据CANN实际的安装路径进行修改。
    3. 按“Esc”键,输入:wq!,按“Enter”保存并退出编辑。

  5. 完成步骤步骤 3~步骤 4的.py文件修改,MindSpeed-LLM即可使用MindIO ACP加速周期性CheckPoint特性。