常见的大模型训练方式有数据并行 / 模型并行 / 流水线并行等,不同的并行方式实际是将参与训练的实例进行分组。同一组内的实例间通信频率(流量)远高于组间的实例进行通信的频率(流量),因此在任务调度时应该尽量地把同一组内的实例尽可能调度到网络较亲和的节点上。本文将介绍如何减少在大规模训练下跨交换机通信的频率。
机器学习平台会在任务创建时向对应的容器中注入一些环境变量,这些环境变量记录了交换机的信息。
/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 ...
/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 <启动脚本的路径>