PyTorch场景（基于MindSpeed-LLM）

本章节将指导用户了解配置在线压测的关键步骤。在线压测的特性介绍、使用约束、支持的产品型号等请参见在线压测。

前提条件

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

操作步骤

在分布式环境初始化完成，能够获取到全局rank之后，修改训练脚本。在训练脚本中拉起TaskD Manager，并在训练进程内部拉起TaskD Worker。

拉起TaskD Manager。

创建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         # 用户填入任务节点总数
proc_per_node=XX     # 用户填入任务每个节点的训练进程数量

init_taskd_manager({"job_id":job_id, "node_nums": node_nums, "proc_per_node": proc_per_node})
start_taskd_manager()

manager.py文件中的参数详细说明请参见def init_taskd_manager(config:dict) -> bool:。

在训练脚本中增加以下代码，拉起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 ...

拉起TaskD Worker。

修改QWEN3_for_PyTorch_2.7_code/mindspeed_llm/training/training.py文件，在代码中增加如下加粗字段。

def pretrain(train_valid_test_dataset_provider,
             model_provider,
             model_type,
             forward_step_func,
             process_non_loss_data_func=None,
             extra_args_provider=None,
             args_defaults={}):
    print_rank_0('time to initialize megatron (seconds): {:.3f}'.format(
        time.time() - _TRAIN_START_TIME))
    print_datetime('after megatron is initialized')
    import torch.distributed as dist
    if dist.is_initialized():
       rank = dist.get_rank()
       from taskd.api.taskd_worker_api import init_taskd_worker
       from taskd.api.taskd_worker_api import start_taskd_worker
       init_taskd_worker(rank,5000,"pt")
       start_taskd_worker()
    app_metrics['app_model_init_finish_time'] = one_logger_utils.get_timestamp_in_ms()
    one_logger_utils.on_pretrain_start()

如果训练中出现报错“the libtaskd.so has not been loaded”，则需在训练脚本中导入LD_PRELOAD环境变量。该环境变量允许系统提前加载指定的so文件。示例如下。

export LD_PRELOAD=/usr/local/Ascend/cann/lib64/libmspti.so:/usr/local/lib/python3.10/dist-packages/taskd/python/cython_api/libs/libtaskd.so

libmspti.so：该so由MindStudio提供，集成在CANN包内。当使用默认安装路径时，路径为：/usr/local/Ascend/cann/lib64/libmspti.so。

libtaskd.so：该so由TaskD组件提供，安装该whl包后，路径为： TaskD所在路径/taskd/python/cython_api/libs/libtaskd.so。
TaskD所在路径可通过以下命令进行查询。回显中的Location字段即为TaskD所在路径。
```
pip show taskd
```

修改任务YAML。

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

ports:                         
   - containerPort: 9601             
     name: taskd-port

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

...  
   labels:  
     ...  
     fault-scheduling: "force"
 ... 
...  
   annotations:  
     ...  
     recover-strategy: "recover"    # 任务可用恢复策略，取值为recover，表示开启进程级别重调度
 ... 
...
spec:
  replicaSpecs:
    Master:
      template:
        spec:
          containers:
          - name: ascend # do not modify
            ...
            args:
              - | 
                cd /job/code; 
                chmod +x scripts/train_start.sh; 
                bash scripts/train_start.sh
    Worker:
      template:
        spec:
          containers:
          - name: ascend # do not modify
            ...
            args:
              - |
                cd /job/code; 
                chmod +x scripts/train_start.sh; 
                bash scripts/train_start.sh
...

父主题： 配置在线压测