基于交换机信息的分布式通信优化

常见的大模型训练方式有数据并行 / 模型并行 / 流水线并行等，不同的并行方式实际是将参与训练的实例进行分组。同一组内的实例间通信频率（流量）远高于组间的实例进行通信的频率（流量），因此在任务调度时应该尽量地把同一组内的实例尽可能调度到网络较亲和的节点上。本文将介绍如何减少在大规模训练下跨交换机通信的频率。

机器学习平台会在任务创建时向对应的容器中注入一些环境变量，这些环境变量记录了交换机的信息。

• MLP_RACK_TOPO_FILE：交换机信息文件在容器中的路径，默认路径为 /ml_platform/rack_topo，示例如下：

  # ${MLP_TASK_INSTANCE_ID}          ${MLP_PRIMARY_HOST}    ${MLP_WORKER_RACK_SW_ID}
  t-20230615154118-9mlks-worker-0    10.252.133.193         febd30c7302dda1fae1399fe0266930f
  t-20230615154118-9mlks-worker-1    10.252.132.1           febd30c7302dda1fae1399fe0266930f
  ...
  

• MLP_${MLP_ROLE}_RACK_SW_ID：当前任务实例所在节点的交换机 hash_id，仅对 HPC 的容器才会注入该值。
• MLP_${MLP_ROLE}_RACK_RANK_INDEX：PyTorchDDP、MPI 框架的任务会注入按照交换机 hash_id 排序后的 worker index 环境变量，同时保证 worker0 的 index=0。训练中可使用该环境变量来指定 node_rank 即可减少 allreduce 等场景跨交换机通信的频率。
  • MPI 框架的任务会按照交换机 hash_id 进行排序，并在原 /root/mpi_hostfile 之外，额外生成 hostfile 文件，路径为 /root/mpi_rack_hostfile，同时保证 worker_0 的位于首位。除此之外，平台会注入环境变量 MLP_MPI_RACK_HOSTS，其格式与 MLP_MPI_HOSTS 相同，但其中 IP 的顺序是按照交换机 hash_id 排序后的。

下文是一个 PyTorchDDP 的启动示例（注意将 NODE_RANK 指定为 MLP_WORKER_RACK_RANK_INDEX）。对于 MPI 的用户直接通过 /root/mpi_rack_hostfile 启动任务即可，无需额外的环境变量配置。

export GPUS_PER_NODE=${MLP_WORKER_GPU:-1}
export NNODES=${MLP_WORKER_NUM:-1}
export NODE_RANK=${MLP_WORKER_RACK_RANK_INDEX:-0}
export PORT=${MLP_WORKER_0_PORT:-1234}
export MASTER_ADDR=${MLP_WORKER_0_HOST:-127.0.0.1}
export MASTER_PORT=${MLP_WORKER_0_PORT:-1234}

DISTRIBUTED_ARGS="--nproc_per_node $GPUS_PER_NODE --nnodes $NNODES --node_rank $NODE_RANK --master_addr $MASTER_ADDR --master_port $MASTER_PORT"
python -m torch.distributed.launch $DISTRIBUTED_ARGS <启动脚本的路径>