PyTorch Distributed Data Parallel使用详解

2023-03-19 18:03:08 pytorch distributed 详解

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 &

以上就是PyTorch Distributed Data Parallel使用详解的详细内容，更多关于PyTorch Distributed Data Parallel的资料请关注其它相关文章！

相关文章