PyTorch中DistributedDataParallel使用笔记

瞻邈2023-09-25 17:52

1. 基本概念

在使用DistributedDataParallel时有一些概率必须掌握

多机多卡	含义
world_size	代表有几台机器，可以理解为几台服务器
rank	第几台机器，即第几个服务器
local_rank	某台机器中的第几块GPU

单机多卡	含义
world_size	代表机器一共有几块GPU
rank	第几块GPU
local_rank	第几块GPU，与rank相同

2. 使用方法

2.1. 修改主函数

在运行的时候，DistributedDataParallel会往你的程序中加入一个参数local_rank，所以要现在你的代码中解析这个参数，如：

python 复制代码

parser.add_argument("--local_rank", type=int, default=1, help="number of cpu threads to use during batch generation")

2.2. 初始化

python 复制代码

torch.distributed.init_process_group(backend="nccl")

os.environ["CUDA_VISIBLE_DEVICES"] = "0, 1, 2"  # 有几块GPU写多少

2.3. 设定device

python 复制代码

local_rank = torch.distributed.get_rank()
torch.cuda.set_device(local_rank)
global device
device = torch.device("cuda", local_rank)

我没用arg.local_rank，新定义了一个local_rank变量，是因为我更信任distributed.get_rank()这个函数

这里用torch.device来写，并且加了global，是因为后面模型和数据都要用到这个device，不会出错

2.4. 模型加载到多gpu

python 复制代码

model.to(device)  # 这句不能少，最好不要用model.cuda()
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[local_rank], output_device=local_rank, find_unused_parameters=True)  # 这句加载到多GPU上

2.5. 数据加载到gpu

python 复制代码

数据.to(device)

2.6. 启动

python 复制代码

torchrun --nproc_per_node=4 --rdzv_endpoint=localhost:12345 train_cylinder_asym.py

参考文献

Pytorch并行计算(二): DistributedDataParallel介绍_dist.barrier_harry_tea的博客-CSDN博客

DistributedDataParallel多GPU分布式训练全过程总结跟着做90%成功_BRiAq的博客-CSDN博客