示例：创建ddp分布式训练（pytorch gpu）-九游平台

创建训练作业

• 方式一：使用pytorch预置框架功能，通过mp.spawn命令启动训练作业。

  创建训练作业的关键参数如表1所示。

  表1 创建训练作业（预置框架）

  参数名称	说明
  创建方式	选择“自定义算法”。
  启动方式	选择“预置框架”，引擎选择“pytorch”，pytorch版本根据训练要求选择。
  代码目录	选择obs桶中训练code文件夹所在路径，例如“obs://test-modelarts/code/”。
  启动文件	选择代码目录中训练作业的python启动脚本。例如“obs://test-modelarts/code/main.py”。
  超参	当资源规格为单机多卡时，需要指定超参world_size和rank。 当资源规格为多机时（即实例数大于 1），无需设置超参world_size和rank，超参会由平台自动注入。

• 方式二：使用自定义镜像功能，通过torch.distributed.launch命令启动训练作业。

  创建训练作业的关键参数如表2所示。

  表2 创建训练作业（自定义镜像 torch.distributed.launch命令）

  参数名称	说明
  创建方式	选择“自定义算法”。
  启动方式	选择“自定义”。
  镜像	选择用于训练的pytorch镜像。
  代码目录	选择obs桶中训练code文件夹所在路径，例如“obs://test-modelarts/code/”。
  启动命令	输入镜像的python启动命令，例如： bash ${ma_job_dir}/code/torchlaunch.sh

• 方式三：使用自定义镜像功能，通过torch.distributed.run命令启动训练作业。

  创建训练作业的关键参数如表3所示。

  表3 创建训练作业（自定义镜像 torch.distributed.run命令）

  参数名称	说明
  创建方式	选择“自定义算法”。
  启动方式	选择“自定义”。
  镜像	选择用于训练的pytorch镜像。
  代码目录	选择obs桶中训练code文件夹所在路径，例如“obs://test-modelarts/code/”。
  启动命令	输入镜像的python启动命令，例如： bash ${ma_job_dir}/code/torchrun.sh

代码示例

文件目录结构如下所示，将以下文件上传至obs桶中：

code                             # 代码根目录
 └─torch_ddp.py                # pytorch ddp训练代码文件
 └─main.py                     # 使用pytorch预置框架功能，通过mp.spawn命令启动训练的启动文件
 └─torchlaunch.sh              # 使用自定义镜像功能，通过torch.distributed.launch命令启动训练的启动文件
 └─torchrun.sh                 # 使用自定义镜像功能，通过torch.distributed.run命令启动训练的启动文件

torch_ddp.py内容如下：

import os
import torch
import torch.distributed as dist
import torch.nn as nn
import torch.optim as optim
from torch.nn.parallel import distributeddataparallel as ddp
# 用于通过 mp.spawn 启动
def init_from_arg(local_rank, base_rank, world_size, init_method):
    rank = base_rank   local_rank
    dist.init_process_group("nccl", rank=rank, init_method=init_method, world_size=world_size)
    ddp_train(local_rank)
# 用于通过 torch.distributed.launch 或 torch.distributed.run 启动
def init_from_env():
    dist.init_process_group(backend='nccl', init_method='env://')
    local_rank=int(os.environ["local_rank"])
    ddp_train(local_rank)
def cleanup():
    dist.destroy_process_group()
class toymodel(nn.module):
    def __init__(self):
        super(toymodel, self).__init__()
        self.net1 = nn.linear(10, 10)
        self.relu = nn.relu()
        self.net2 = nn.linear(10, 5)
    def forward(self, x):
        return self.net2(self.relu(self.net1(x)))
def ddp_train(device_id):
    # create model and move it to gpu with id rank
    model = toymodel().to(device_id)
    ddp_model = ddp(model, device_ids=[device_id])
    loss_fn = nn.mseloss()
    optimizer = optim.sgd(ddp_model.parameters(), lr=0.001)
    optimizer.zero_grad()
    outputs = ddp_model(torch.randn(20, 10))
    labels = torch.randn(20, 5).to(device_id)
    loss_fn(outputs, labels).backward()
    optimizer.step()
    cleanup()
if __name__ == "__main__":
    init_from_env()

main.py内容如下：

import argparse
import torch
import torch.multiprocessing as mp
parser = argparse.argumentparser(description='ddp demo args')
parser.add_argument('--world_size', type=int, required=true)
parser.add_argument('--rank', type=int, required=true)
parser.add_argument('--init_method', type=str, required=true)
args, unknown = parser.parse_known_args()
if __name__ == "__main__":
    n_gpus = torch.cuda.device_count()
    world_size = n_gpus * args.world_size
    base_rank = n_gpus * args.rank
    # 调用 ddp 示例代码中的启动函数
    from torch_ddp import init_from_arg
    mp.spawn(init_from_arg,
        args=(base_rank, world_size, args.init_method),
        nprocs=n_gpus,
        join=true)

#!/bin/bash
# 系统默认环境变量，不建议修改
master_host="$vc_worker_hosts"
master_addr="${vc_worker_hosts%%,*}"
master_port="6060"
job_id="1234"
nnodes="$ma_num_hosts"
node_rank="$vc_task_index"
ngpus_per_node="$ma_num_gpus"
# 自定义环境变量，指定python脚本和参数
python_script=${ma_job_dir}/code/torch_ddp.py
python_args=""
cmd="python -m torch.distributed.launch \
    --nnodes=$nnodes \
    --node_rank=$node_rank \
    --nproc_per_node=$ngpus_per_node \
    --master_addr $master_addr \
    --master_port=$master_port \
    --use_env \
    $python_script \
    $python_args
"
echo $cmd
$cmd

torchrun.sh内容如下：

pytorch 2.1版本需要将“rdzv_backend”参数设置为“static：--rdzv_backend=static”。

#!/bin/bash
# 系统默认环境变量，不建议修改
master_host="$vc_worker_hosts"
master_addr="${vc_worker_hosts%%,*}"
master_port="6060"
job_id="1234"
nnodes="$ma_num_hosts"
node_rank="$vc_task_index"
ngpus_per_node="$ma_num_gpus"
# 自定义环境变量，指定python脚本和参数
python_script=${ma_job_dir}/code/torch_ddp.py
python_args=""
if [[ $node_rank == 0 ]]; then
    ext_args="--rdzv_conf=is_host=1"
else
    ext_args=""
fi
cmd="python -m torch.distributed.run \
    --nnodes=$nnodes \
    --node_rank=$node_rank \
    $ext_args \
    --nproc_per_node=$ngpus_per_node \
    --rdzv_id=$job_id \
    --rdzv_backend=c10d \
    --rdzv_endpoint=$master_addr:$master_port \
    $python_script \
    $python_args
    "
echo $cmd
$cmd