示例:创建ddp分布式训练(pytorch gpu)-j9九游会登录
本文介绍三种使用训练作业来启动pytorch ddp训练的方法及对应代码示例。
- 使用pytorch预置框架功能,通过mp.spawn命令启动
- 使用自定义镜像功能
- 通过torch.distributed.launch命令启动
- 通过torch.distributed.run命令启动
创建训练作业
- 方式一:使用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
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
相关文档
意见反馈
文档内容是否对您有帮助?
如您有其它疑问,您也可以通过华为云社区问答频道来与我们联系探讨