Torch对接MindSpeed框架

前提条件

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

操作步骤

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

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

  2. 按照MindSpeed框架的指导说明准备好MindSpeed框架环境后,进入Megatron-LM安装目录。

    cd Megatron-LM/

  3. 修改pretrain_gpt.py文件。

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

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

    此处以编辑“tests_extend/model_tests/perf_model/llama2/pretrain_llama2_70b_4k.sh”脚本为例。

    1. 打开“tests_extend/model_tests/perf_model/llama2/pretrain_llama2_70b_4k.sh”脚本。
      vim tests_extend/model_tests/perf_model/llama2/pretrain_llama2_70b_4k.sh
    2. 按“i”进入编辑模式,开启周期性CheckPoint加速功能需要增加脚本中加粗部分。
      #!/bin/bash
      
      export CUDA_DEVICE_MAX_CONNECTIONS=1
      source "tests_extend/system_tests/env_npu.sh"
      
      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
      
      # Change for multinode config
      NPUS_PER_NODE=16
      MASTER_ADDR=<master_ip_address>
      MASTER_PORT=6000
      NNODES=8
      NODE_RANK=<local_rank>
      WORLD_SIZE=$(($NPUS_PER_NODE*$NNODES))
      
      CKPT_DIR=./ckpt_llama
      DATA_PATH="/home/dataset/llama2/alpaca_text_document"
      TOKENIZER_MODEL="/home/dataset/model/llama-2-7b-hf/tokenizer.model"
      
      TP=8
      PP=4
      CP=1
      
      DISTRIBUTED_ARGS="
          --nproc_per_node $NPUS_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} \
          --use-fused-rotary-pos-emb \
          --use-fused-swiglu \
          --use-fused-rmsnorm \
          --log-throughput \
          --overlap-grad-reduce \
          --overlap-param-gather \
          --use-ascend-mc2 \
          --num-layers-per-virtual-pipeline-stage 2 \
          --sequence-parallel \
          --use-distributed-optimizer \
          --num-layers 80 \
          --hidden-size 8192 \
          --ffn-hidden-size 28672 \
          --num-attention-heads 64 \
          --tokenizer-type Llama2Tokenizer \
          --tokenizer-model ${TOKENIZER_MODEL} \
          --seq-length 4096 \
          --max-position-embeddings 4096 \
          --micro-batch-size 1 \
          --global-batch-size 32 \
          --make-vocab-size-divisible-by 1 \
          --lr 1.0e-6 \
          --train-iters 5000 \
          --lr-decay-style cosine \
          --untie-embeddings-and-output-weights \
          --attention-dropout 0.0 \
          --init-method-std 0.01 \
          --hidden-dropout 0.0 \
          --position-embedding-type rope \
          --normalization RMSNorm \
          --swiglu \
          --use-flash-attn \
          --no-masked-softmax-fusion \
          --attention-softmax-in-fp32 \
          --min-lr 1.0e-7 \
          --weight-decay 0.1 \
          --clip-grad 1.0 \
          --adam-beta1 0.9 \
          --initial-loss-scale 4096.0 \
          --adam-beta2 0.95 \
          --adam-eps 1e-5 \
          --no-gradient-accumulation-fusion \
          --disable-bias-linear \
          --group-query-attention \
          --num-query-groups 8 \
          --lr-warmup-fraction 0.01 \
          --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 \
          --save $CKPT_DIR \
      
      set +x

      周期性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即可使用MindIO ACP加速周期性CheckPoint特性。