Adapting Torch to MindSpeed-LLM

Prerequisites

• Learn about Constraints of MindIO ACP.
• Prepare the MindSpeed-LLM framework by referring to MindSpeed-LLM. Note that the matching Megatron-LM version is core_v0.12.1.

The release package is intended for use with the MindSpeed-LLM 2.3.0 branch. For detailed instructions on preparing the environment, code, and dataset, refer to the guidelines in the MindSpeed-LLM repository to ensure their security.

Procedure

1. Log in to the compute node as a service user.
   

   The service user is not {MindIO-install-user}, HwHiAiUser, or hwMindX, but it is determined based on the actual situation.

2. Go to the MindSpeed-LLM installation directory.

   cd MindSpeed-LLM/

3. Modify the pretrain_gpt.py file.
   1. Open the pretrain_gpt.py file.

      vim pretrain_gpt.py

   2. Press i to enter the insert mode, find from mindspeed_llm import megatron_adaptor in the file header, and add the information in bold in a new line.

      from mindspeed_llm import megatron_adaptor
      import mindio_acp

   3. Press Esc, type :wq!, and press Enter to save the changes and exit.
4. Edit the pre-training script (for reference only).

   The following uses the examples/mcore/llama2/pretrain_llama2_7b_ptd.sh script as an example.

   1. Open the examples/mcore/llama2/pretrain_llama2_7b_ptd.sh script.

      vim examples/mcore/llama2/pretrain_llama2_7b_ptd.sh

   2. Press i to enter the insert mode and add the following information in bold to the script to enable periodic checkpoint acceleration.

      #!/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/cann/set_env.sh
      
      NPUS_PER_NODE=8
      MASTER_ADDR=localhost
      MASTER_PORT=6000
      NNODES=1
      NODE_RANK=0
      WORLD_SIZE=$(($NPUS_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 $NPUS_PER_NODE \
          --nnodes $NNODES \
          --node_rank $NODE_RANK \
          --master_addr $MASTER_ADDR \
          --master_port $MASTER_PORT
      "
      
      GPT_ARGS="
          --use-mcore-models \
          --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 \
          --load $CKPT_LOAD_DIR \
          --save $CKPT_SAVE_DIR \
          | tee logs/train_llama2_7b.log

      The parameters related to periodic checkpoint acceleration are described as follows:

      • MINDIO_AUTO_PATCH_MEGATRON indicates the source code of the auto patch Megatron of MindIO ACP, which is used to enable the periodic checkpoint acceleration feature.
      • Set GLOO_SOCKET_IFNAME based on the high-speed NIC of the primary node.
      • LD_LIBRARY_PATH indicates the address of the driver .so library of the CANN package. Change it to the actual CANN installation path.
      • Change the path of the set_env.sh file to the actual CANN installation path.
   3. Press Esc, type :wq!, and press Enter to save the changes and exit.
5. After .py files from 3 to 4 are modified, MindSpeed-LLM can use the periodic checkpoint acceleration function of MindIO ACP.