一、PyTorch分布式 DP与DDP
1.1 PyTorch分布式支持
![](https://i-blog.csdnimg.cn/direct/903677ccd71449bea565f627dec83d34.png)
![](https://i-blog.csdnimg.cn/direct/8125dde03fea4a70adbf0f4561ed5094.png)
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的分布式支持。