PyTorch 分布式使用方式及代码解析

一、PyTorch分布式 DP与DDP

1.1 PyTorch分布式支持

1.2 PyTorch分布式调用-DP 

1.3 PyTorch分布式调用-DDP 

1.4 PyTorch分布式-通信后端 

gloo:具有各种原语的集体通信库，用
于多机训练。Facebook为在Linux上
运行而构建的
NCCL(Nvidia Collective multi-GPU
Communication Library) : 是一个实现
多GPU的collective communication通信
（all-gather, reduce, broadcast）库，
Nvidia做了很多优化，以在PCIe、
NVLink、InfiniBand上实现较高的通信
速度
分布式后端适用场景：
GPU分布式训练——nccl
CPU分布式训练——gloo

分布式实现中的核心通信模块为 torch.distributed,该模块支持三种后端
• nccl
• gloo
• mpi
import torch.distributed as dist
...
dist.init_process_group(backend= args.dist_backend, init_method= args.dist_url, world_size= args.world_size,
rank= args.rank)
• backend (str or Backend) – 后端使用。根据构建时配置，有效值包括mpi、gloo、nccl。
• init_method (str, optional) – 指定如何初始化进程组的URL。 tcp、file、master方式
• world_size (int, optional) – 参与作业的进程数。
• rank (int, optional) – 当前流程的排名。
• timeout (timedelta, optional) – 针对进程组执行的操作超时，默认值等于30分钟，这仅适用于gloo后端。
• group_name (str, optional, deprecated) – 团队名字。
torch.distributed.is_nccl_available()
torch.distributed.is_mpi_available()

 二、实际代码解析

2.1 PyTorch分布式调用-DP

单进程控制多 GPU

# main.py
import torch
import torch.distributed as dist
gpus = [0, 1, 2, 3]
torch.cuda.set_device('cuda:{}'.format(gpus[0]))
train_dataset = ...
train_loader =
torch.utils.data.DataLoader(train_dataset,
batch_size=...)
model = ...
model = nn.DataParallel(model.to(device),
device_ids=gpus, output_device=gpus[0])
optimizer = optim.SGD(model.parameters())

主程序段：

for epoch in range(100):
for batch_idx, (data, target) in enumerate(train_loader):
# 这里要 images/target.cuda()
images = images.cuda(non_blocking=True)
target = target.cuda(non_blocking=True)
...
output = model(images)
loss = criterion(output, target)
...
optimizer.zero_grad()
loss.backward()
optimizer.step()

• model.cuda() 模型和数据均需要load进gpu
• device_ids 参与训练得gpu有哪些
• output_device 用于汇总梯度的gpu是哪个
batchsize设置成n倍的单卡batchsize

2.2 PyTorch分布式调用-DDP

Ø 在使用 distributed 包的任何其他函数之前，需要使用 init_process_group 初始化进程组，同时
初始化 distributed 包；
Ø 如果需要进行小组内集体通信，用 new_group 创建子分组；
Ø 创建分布式并行（DistributedDataParallel）模型 DDP(model, device_ids=device_ids)；
Ø 为数据集创建 Sampler；
Ø 使用启动工具 torch.distributed.launch 在每个主机上执行一次脚本，开始训练；

# main.py
import torch.distributed as dist
#进程同步
dist.init_process_group(backend='nccl’)
#把模型 数据加载到当前进程gpu
torch.cuda.set_device(args.local_rank)
#数据切分
train_dataset = ...
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler)
model = ...
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank])
optimizer = torch.optim.SGD(model.parameters())
主程序段

2.3 DP与DDP区别

Ø DP
• 使用简单，一行代码即可；
• 不支持多机多卡分布式并行；
• 负载不均衡
Ø DDP
• 负载均衡；
• 支持多机多卡并行；
• 相较于DP，需要传输的数据量小，效率更高；
结论：综上，推荐调用DDP完成PyTorch的分布式支持。