PyTorch Distributed Data Parallel使用详解

DDP

Distributed Data Parallel 简称 DDP，是 PyTorch 框架下一种适用于单机多卡、多机多卡任务的数据并行方式。由于其良好的执行效率及广泛的显卡支持，熟练掌握 DDP 已经成为深度学习从业者所必备的技能之一。本文结合具体代码，详细地说明了 DDP 在项目中的使用方式。读者按照本文所给的范例，只需稍经调试，即可实现 DDP 的整套流程。

概念辨析

具体讲解 DDP 之前，我们先了解了解它和 Data Parallel (DP) 之间的区别。DP 同样是 PyTorch 常见的多 GPU 并行方式之一，且它的实现非常简洁：

# 函数定义
torch.nn.DataParallel(module, device_ids=None, output_device=None, dim=0)
'''
module : 模型
device_ids : 参与训练的 GPU 列表
output_device : 指定输出的 GPU, 通常省略, 即默认使用索引为 0 的显卡
'''
# 程序模板
device_ids = [0, 1]
net = torch.nn.DataParallel(net, device_ids=device_ids)

基本原理及固有缺陷：在 Data Parallel 模式下，数据会被自动切分，加载到 GPU。同时，模型也将拷贝至各个 GPU 进行正向传播。在多个进程之间，会有一个进程充当 master 节点，负责收集各张显卡积累的梯度，并据此更新参数，再统一发送至其他显卡。因此整体而言，master 节点承担了更多的计算与通信任务，容易造成网络堵塞，影响训练速度。

常见问题及解决方案：Data Parallel 要求模型必须在 device_ids[0] 拥有参数及缓冲区，因此当卡 0 被占用时，可以在 nn.DataParallel 之前添加如下代码：

# 按照 PIC_BUS_ID 顺序自 0 开始排列 GPU 设备
os.environ['CUDA_DEVICE_ORDER'] = 'PIC_BUS_ID'
# 设置当前使用的 GPU 为 2、3 号设备
os.environ['CUDA_VISIBLE_DEVICES'] = '2, 3'

如此，device_ids[0] 将被默认为 2 号卡，device_ids[1] 则对应 3 号卡

相较于 DP, Distributed Data Parallel 的实现要复杂得多，但是它的优势也非常明显：

• DDP 速度更快，可以达到略低于显卡数量的加速比；

• DDP 可以实现负载的均匀分配，克服了 DP 需要一个进程充当 master 节点的固有缺陷；

• 采用 DDP 通常可以支持更大的 batch size，不会像 DP 那样出现其他显卡尚有余力，而卡 0 直接 out of memory 的情况；

• 另外，在 DDP 模式下，输入到 data loader 的 bacth size 不再代表总数，而是每块 GPU 各自负责的 sample 数量。比方说，batch_size = 30，有两块 GPU。在 DP 模式下，每块 GPU 会负责 15 个样本。而在 DDP 模式下，每块 GPU 会各自负责 30 个样本；

• DDP 基本原理：倘若我们拥有 N 张显卡，则在 Distributed Data Parallel 模式下，就会启动 N 个进程。每个进程在各自的卡上加载模型，且模型的参数完全相同。训练过程中，各个进程通过一种名为 Ring-Reduce 的方式与其他进程通信，交换彼此的梯度，从而获得所有的梯度信息。随后，各个进程利用梯度的平均值更新参数。由于初始值和更新量完全相同，所以各个进程更新后的参数仍保持一致。

常用术语

• rank

• 进程号

• 多进程上下文中，通常假定 rank = 0 为主进程或第一个进程

• node

• 物理节点，表示一个容器或一台机器

• 节点内部可以包含多个 GPU

• local_rank

• 一个 node 中，进程的相对序号

• local_rank 在 node 之间独立

• world_size

• 全局进程数

• 一个分布式任务中 rank 的数量

• group

• 进程组

• 一个分布式任务就对应一个进程组

• 只有当用户创立多个进程组时，才会用到

代码实现

Distributed Data Parallel 可以通过 Python 的 torch.distributed.launch 启动器，在命令行分布式地执行 Python 文件。执行过程中，启动器会将当前进程（其实就是 GPU）的 index 通过参数传递给 Python，而我们可以利用如下方式获取这个 index：

import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', default=-1, type=int,
                   metavar='N', help='Local process rank.')
args = parser.parse_args()
# print(args.local_rank)
# local_rank 表示本地进程序号

随后，初始化进程组。对于在 GPU 执行的任务，建议选择 nccl (由 NVIDIA 推出) 作为通信后端。对于在 CPU 执行的任务，建议选择 gloo (由 Facebook 推出) 作为通信后端。倘若不传入 init_method，则默认为 env://，表示自环境变量读取分布式信息

dist.init_process_group(backend='nccl', init_method='env://')
# 初始化进程组之后, 通常会执行这两行代码
torch.cuda.set_device(args.local_rank)
device = torch.device('cuda', args.local_rank)
# 后续的 model = model.to(device), tensor.cuda(device)
# 对应的都是这里由 args.local_rank 初始化得到的 device

数据部分，使用 Distributed Sampler 划分数据集，并将 sampler 传入 data loader。需要注意的是，此时在 data loader 中不能指定 shuffle 为 True，否则会报错 (sampler 已具备随机打乱功能)

dev_sampler = data.DistributedSampler(dev_data_set)
train_sampler = data.DistributedSampler(train_data_set)
dev_loader = data.DataLoader(dev_data_set, batch_size=dev_batch_size,
                            shuffle=False, sampler=dev_sampler)
train_loader = data.DataLoader(train_data_set, batch_size=train_batch_size,
                              shuffle=False, sampler=train_sampler)

模型部分，首先将将模型送至 device，即对应的 GPU 上，再使用 Distributed Data Parallel 包装模型（顺序颠倒会报错）

model = model.to(device)
model = nn.parallel.DistributedDataParallel(
   model, device_ids=[args.local_rank], output_device=args.local_rank
)

Distributed Data Parallel 模式下，保存模型应使用 net.module.state_dict()，而非 net.state_dict()。且无论是保存模型，还是 LOGGER 打印，只对 local_rank 为 0 的进程操作即可，因此代码中会有很多 args.local_rank == 0 的判断

if args.local_rank == 0:
   LOGGER.info(f'saving latest model: {output_path}')
   torch.save({'model': model.module.state_dict(),
               'optimizer': None, 'epoch': epoch, 'best-f1': best_f1},
              open(os.path.join(output_path, 'latest_model_{}.pth'.format(fold)), 'wb'))

利用 torch.load 加载模型时，设置 map_location=device，否则卡 0 会承担更多的开销

load_model = torch.load(best_path, map_location=device)
model.load_state_dict(load_model['model'])

• dist.barrier() 可用于同步多个进程，建议只在必要的位置使用，如初始化 DDP 模型之前、权重更新之后、开启新一轮 epoch 之前

• 计算 accuracy 时，可以使用 dist.all_reduce(score, op=dist.ReduceOp.SUM)，将各个进程计算的准确率求平均

• 计算 f1-score 时，可以使用 dist.all_gather(all_prediction_list, prediction_list)，将各个进程获得的预测值和真实值汇总到 all_list，再统一代入公式

启动方式

torch.distributed.launch

# 此处 --nproc_per_node 4 的含义是 server 有 4 张显卡
python torch.distributed.launch --nproc_per_node 4 train.py
# 倘若使用 nohup, 则注意输入命令后 exit 当前终端
python torch.distributed.launch --nproc_per_node 4 train.py

• torchrun，推荐使用这种方式，因为 torch.distributed.launch 即将弃用

代码中，只需将 Argument Parser 相关的部分替换为

local_rank = int(os.environ['LOCAL_RANK'])

然后将 args.local_rank 全部改为 local_rank 即可

启动命令

# 单机多卡训练时, 可以不指定 nnodes
torchrun --nnodes=1 --nproc_per_node=4 train.py
# 倘若使用 nohup, 则注意输入命令后 exit 当前终端
nohup torchrun --nnodes=1 --nproc_per_node=4 train.py > nohup.out &

来源：https://juejin.cn/post/7211775098310950973