发起 PyTorchDDP 分布式训练

PyTorch DistributedDataParallel（DDP）是一种数据并行的分布式训练方法。通过 DDP 创建多个进程进行模型训练，通过 ring-all-reduce 的方法做进程通讯，完成梯度的交换及参数更新。

1. 用户在【自定义训练】模块创建一个训练任务时选择实例配置为 PyTorch DDP，按需配置各种训练角色并提交任务表单进入任务创建环节。PyTorch DDP 仅包含 worker 这一种角色用于训练模型，其中编号为 0 的 worker（worker0）额外承担保存 checkpoint 或日志的任务。
2. 任务提交后，机器学习平台将为用户创建对应的实例并向所有实例注入相关的环境变量，通过环境变量用户代码得知集群的信息以及当前实例对应的训练角色，从而完成对应角色的本职任务直到训练结束。
   • 训练之前，平台将等待所有实例之间的网络通畅。
   • 训练过程中，任意实例失败（退出码非 0）则训练任务失败。
   • 所有实例训练完成（退出码为 0）则训练任务成功。

平台为用户注入了有可能在训练代码或者入口命令需要用到的环境变量，通过 PyTorch 官方提供的 torch.distributed.launch 启动命令结合环境变量启动训练任务。

环境变量

• MLP_WORKER_0_HOST：worker0 （RANK=0 的 worker）的地址。
• MLP_WORKER_0_PORT：worker0 （RANK=0 的 worker）的端口。
• MLP_ROLE_INDEX：训练实例（机器）的编号，不同实例中该环境变量的值不同。
• MLP_WORKER_NUM：训练实例的数量。
• MLP_WORKER_GPU：单个 woker 包含的 GPU 数量。

如需使用【自定义训练】模块的通用环境变量详见通用环境变量列表。

关键参数

• nproc_per_node：单个实例（机器）上运行的进程数，使用 GPU 时通常为每台机器上的 GPU 数量。
• nnodes：对应环境变量 MLP_WORKER_NUM 的值。
• node_rank：对应环境变量 MLP_ROLE_INDEX 的值。
• master_addr：对应环境变量 MLP_WORKER_0_HOST 的值。
• master_port：对应环境变量 MLP_WORKER_0_PORT 的值。

使用如下命令启动 DDP 分布式训练任务：

python -m torch.distributed.launch --nproc_per_node <单个实例上的进程数> --master_addr $MLP_WORKER_0_HOST --node_rank $MLP_ROLE_INDEX --master_port $MLP_WORKER_0_PORT --nnodes $MLP_WORKER_NUM <代码文件的绝对路径>

# 选择 CPU 实例规格时，nproc_per_node = 1，则代表每个实例上仅执行 1 个训练进程。
python -m torch.distributed.launch --nproc_per_node 1 --master_addr $MLP_WORKER_0_HOST --node_rank $MLP_ROLE_INDEX --master_port $MLP_WORKER_0_PORT --nnodes $MLP_WORKER_NUM <代码文件的绝对路径>

# 选择 GPU 实例规格时，通过环境变量配置每个实例上的训练进程数量与 GPU 数量一致。
python -m torch.distributed.launch --nproc_per_node $MLP_WORKER_GPU --master_addr $MLP_WORKER_0_HOST --node_rank $MLP_ROLE_INDEX --master_port $MLP_WORKER_0_PORT --nnodes $MLP_WORKER_NUM <代码文件的绝对路径>

具体的示例代码详见 GitHub。