Usage Guidance

After installing the MindIO TFT SDK, you need to start the MindIO TFT module in your framework and synchronize the optimizer data update status to this module during training.

Integrating with the MindSpeed-LLM Framework

Prerequisites

• Before use, please understand the Constraints of MindIO TFT.
• For MindSpeed-LLM Framework preparation, see MindSpeed-LLM. The matching Megatron-LM version is core_v0.12.1.

Procedure

1. (Optional) Edit the "torchrun" file.

   1. Locate the "torchrun" file in the environment.

      which torchrun

   2. Open the "torchrun" file at the path shown by the command above.

      vim {torchrun file path}/torchrun

   3. Press i to enter insert mode, and add import mindio_ttp.framework_ttp at the corresponding position in the file.

      import re
      import sys
      import mindio_ttp.framework_ttp
      from torch.distributed.run import main as torch_main

   4. Press Esc, type :wq!, and press Enter to save and exit insert mode.

2. Edit the pre-training script (for reference only).

   The following uses editing 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 insert mode. To enable the high availability feature, add the following content to the script.

      export GLOO_SOCKET_IFNAME=enp189s0f0
      export TTP_ADDR="master node ip"
      source /usr/local/Ascend/cann/set_env.sh
      
      # After --bf16 in GPT_ARGS, add the following content
          \
          --enable-high-availability \
          --enable-hbmfault-repair \
          --enable-worker-reboot \
          --distributed-optimizer-no-replica \
      

      The modified pretrain_llama2_7b_ptd.sh script example is as follows:

      #!/bin/bash
      
      export CUDA_DEVICE_MAX_CONNECTIONS=1
      export PYTORCH_NPU_ALLOC_CONF=expandable_segments:True
      
      export GLOO_SOCKET_IFNAME=enp189s0f0
      export TTP_ADDR="master node ip"
      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 \
          --enable-high-availability \
          --enable-hbmfault-repair \
          --enable-worker-reboot \
          --distributed-optimizer-no-replica \
      "
      
      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 the high availability feature are described as follows:

      • GLOO_SOCKET_IFNAME: Configure this based on the actual high-speed NIC of the master node.

      • TTP_ADDR: The IP address of the cluster master node, which must conform to the standard IPv4 or IPv6 format. For details, see Environment Variables.

      • set_env.sh file path: Modify this based on the actual installation path of CANN.

      • enable-high-availability: The master switch for MindIO TFT, disabled by default. When configured, the dying gasp feature is enabled by default.

        After the MindIO TFT switch is enabled, the memory usage of various optimizers will change. For details about the changes, see Table 1.

        For distributed optimizers, the static memory increases due to the addition of optimizer replicas. However, as the cluster size increases, the DP Size becomes larger, and the average increase in memory per device is very small, which helps avoid OOM. Therefore, it is recommended to be used in large clusters. Choose whether to enable it and adjust parameters based on the memory situation.

      • enable-hbmfault-repair: switch for MindIO UCE, disabled by default. When configured, it performs fault detection on on-chip memory and completes online repair, achieving step-level recomputation functionality. This switch takes effect when enable-high-availability is enabled. This feature depends on the memory management mechanism of PyTorch and can only be used when the PyTorch environment variable PYTORCH_NO_NPU_MEMORY_CACHING is not configured, meaning the memory reuse mechanism is enabled. If export PYTORCH_NO_NPU_MEMORY_CACHING = 1, this feature cannot be used.

      • enable-worker-reboot: switch for MindIO ARF, disabled by default. When configured, it performs process-level restart repair to continue training when a general fault occurs. This switch takes effect when enable-high-availability is enabled.

      • distributed-optimizer-no-replica: After enabling the high availability feature, the distributed optimizer adds optimizer replicas by default, which increases on-chip memory usage. After enabling this switch, the distributed optimizer does not increase replica memory usage. In MindIO UCE and MindIO ARF scenarios, periodic checkpoints are directly used for online repair.

      Table 1 Theoretical numerical changes in optimizer parameters between the native optimizer and the optimizer when MindIO TFT enabled

      Optimizer	Native	MindIO TFT Enabled	Description
      fp16/bf16	20	20	-
      fp32	16	16	-
      fp16/bf16 Distributed	4 + 16/d	4 + 16 * N/d	d: DP Group Size N: Number of replicas, N < d

   3. Press Esc, type :wq!, and press Enter to save and exit insert mode.

Integrating with MindCluster

MindIO TFT provides services in the form of an SDK, with no resident processes. The service starts when the training process starts. When the training job ends, the service exits.

When integrating with MindCluster, MindCluster manages Kubernetes containers. The integration process within Kubernetes containers is consistent with bare-metal installation and deployment.

Procedure

• When the Python environment is not installed on shared storage, to facilitate use on large clusters, the MindIO TFT SDK can be integrated into the container image. When a pod is deployed from this image, the MindIO TFT SDK is already installed.

• The MindIO TFT Controller module and Processor module exchange heartbeat messages. When Kubernetes performs network isolation, you need to add the communication port to the yaml file configured when creating the pod.

  Modify the yaml file configured when creating the pod. Here, "pod.yaml" is used as an example.

  1. Open the "pod.yaml" file.

     vim pod.yaml

  2. Press i to enter insert mode and add the following content.

     ports:
       - containerPort: 8000 # Port for communication between the MindIO TFT service Controller and Processor
         name: ttp-port

  3. Press Esc, type :wq!, and press Enter to save and exit insert mode.

• To adapt to the Kubernetes network, make the following modifications based on the pre-training script in Step 2.

  # Comment out the following two lines. These environment variables are configured by MindCluster.
  # MASTER_ADDR=$(hostname -I | awk '{print $1}')
  # MASTER_PORT=XXXX
  
  # Obtain the MASTER_ADDR and MASTER_PORT environment variables (the service network IP address of Kubernetes) from Kubernetes.
  CONTROLLER_ADDR=$(hostname -I | awk '{print $1}')
  PROCESSOR_ADDR=${MASTER_ADDR}
  export CONTROLLER_ADDR
  export PROCESSOR_ADDR

Integrating with Non-MindSpeed-LLM Frameworks

Prerequisites

Before you proceed, please understand the Constraints of MindIO TFT.

Feature Reference

The functional adaptation points required for related features are shown in Table 1, and the code reference links corresponding to each functional adaptation point are shown in Table 2.

Table 1 Features and functional adaptation points

Table 2 Code reference links for related features