配置弹性训练

前提条件

• 在相应节点上完成以下组件的安装：Ascend Docker Runtime、Ascend Operator、ClusterD、Ascend Device Plugin和Volcano（以上MindCluster组件版本均需与TaskD配套）
• 在容器内安装torch_npu（7.1.RC1及以上版本）、Kernels（8.2.RC1及以上版本）、Toolkit（8.2.RC1及以上版本）、TaskD和MindIO（7.2.RC1及以上版本）

操作步骤

1. 在分布式环境初始化完成，能够获取到全局rank之后，修改训练脚本，在训练脚本中拉起TaskD Manager。
   1. 创建manager.py文件，放在调用训练脚本时的当前目录下。manager.py文件内容如下所示。

      from taskd.api import init_taskd_manager, start_taskd_manager
      import os
      
      job_id=os.getenv("MINDX_TASK_ID")
      node_nums=XX         # 用户填入任务节点总数，int类型
      proc_per_node=XX     # 用户填入任务每个节点的训练进程数量，int类型
      
      init_taskd_manager({"job_id":job_id, "node_nums": node_nums, "proc_per_node": proc_per_node})
      start_taskd_manager()

   2. 在训练脚本中增加以下代码，拉起TaskD Manager。

      sed -i '/import os/i import taskd.python.adaptor.patch' $(pip3 show torch | grep Location | awk -F ' ' '{print $2}')/torch/distributed/run.py
      export TASKD_PROCESS_ENABLE="on"
      if [[ "${RANK}" == 0 ]]; then
          export MASTER_ADDR=${POD_IP}
          python manager.py &           # 具体执行路径由当前路径决定
      fi
          
      torchrun ...

2. 修改任务YAML。
   1. 在任务YAML中修改容器暴露端口，在所有的Pod下增加TaskD通信使用的端口9601。

      ports:                         
         - containerPort: 9601             
           name: taskd-port 

   2. 在任务YAML中新增以下加粗字段。

      ...  
         labels:  
           ...  
           process-recover-enable: "on"  
           fault-scheduling: "force"
       ... 
      ...  
         annotations:  
           ...
           wait-reschedule-timeout: "270" # 进程级恢复等待故障节点重调度的超时时间，默认为270秒，取值范围为30~270。进程级恢复和弹性训练均开启时，等待此时间后若故障节点调度成功，则进行进程级恢复，否则触发弹性训练
           recover-strategy: "elastic-training"    # 任务可用恢复策略，取值为elastic-training，表示开启弹性训练
       ... 
      ...
      spec:
        replicaSpecs:
          Master:
            template:
              spec:
                containers:
                - name: ascend # do not modify
                  env:
                    - name: MINDIO_WAIT_MINDX_TIME         # 未开启进程级恢复，开启弹性训练场景下建议配置60以上
                      value: "60"
                  args:
                    - | 
                      ...
                      bash scripts/train_start.sh /job/code /job/output pretrain_gpt.py \
                        ...
                        --enable-high-availability \
                        --enable-elastic-training \
                        ...
          Worker:
            template:
              spec:
                containers:
                - name: ascend # do not modify
                  env:
                    - name: MINDIO_WAIT_MINDX_TIME         # 未开启进程级恢复，开启弹性训练场景下建议配置60以上
                      value: "60"
                  args:
                    - |
                      ...
                      bash scripts/train_start.sh /job/code /job/output pretrain_gpt.py \
                        ...
                        --enable-high-availability \
                        --enable-elastic-training \
                        ...
      ...

3. 修改训练框架代码。
   进入“mindxdl-deploy”仓库，根据mindxdl-deploy开源仓版本说明进入版本对应分支，获取“samples/train/resumable-training/fault-tolerance/without-ranktable/pytorch/llama2”目录下的train_start.sh文件，在管理节点构造成如下的目录结构。

   root@ubuntu:/data/atlas_dls/public/code/LLAMA2_for_PyTorch_2.7_code/scripts#
   scripts/
   └── train_start.sh