PyTorch训练（三）：DDP（DistributedDataParallel）【“单机多卡”、“多机多卡”分布式训练模式】【支持混合精度（fp16/fp32）】【只将数据并行，模型大于显卡则不行】

一、概述

我们知道 PyTorch 本身对于单机多卡提供了两种实现方式

• DataParallel（DP）：Parameter Server模式，一张卡位reducer，实现也超级简单，一行代码。
• DistributedDataParallel（DDP）：All-Reduce模式，本意是用来分布式训练，但是也可用于单机多卡。

DataParallel（DP）是基于Parameter server的算法，实现比较简单，只需在原单机单卡代码的基础上增加一行：

gpu_ids = [0, 2, 3]
model = nn.DataParallel(model, device_ids=gpu_id)

但是DP模式负载不均衡的问题比较严重，有时在模型较大的时候（比如bert-large），reducer的那张卡会多出3-4g的显存占用，并且速度也比较慢。

在 pytorch 1.0 之后，官方终于对分布式的常用方法进行了封装，支持 all-reduce，broadcast，send 和 receive 等等。通过 MPI 实现 CPU 通信，通过 NCCL 实现 GPU 通信。官方也曾经提到用 DistributedDataParallel 解决 DataParallel 速度慢，GPU 负载不均衡的问题，目前已经很成熟了～

官方建议用新的DDP（DistributedDataParallel），采用all-reduce算法，本来设计主要是为了多机多卡使用，但是单机上也能用。

使用 4 块 Tesla V100-PICE 在 ImageNet 进行了运行时间的测试，测试结果发现 Apex 的加速效果最好，但与 Horovod/Distributed 差别不大，平时可以直接使用内置的 Distributed。Dataparallel 较慢，不推荐使用。

• torch.nn.DataParallel 效果不好的原因主要是其全程维护一个 optimizer，对各 GPU 上梯度进行求和，而在主 GPU 进行参数更新，之后再将模型参数 broadcast 到其他 GPU。
• DistributedDataParallel 在每次迭代中，每个进程具有自己的 optimizer，并独立完成所有的优化步骤。在各进程梯度计算完成之后，各进程需要将梯度进行汇总平均，然后再由 rank=0 的进程，将其 broadcast 到所有进程。
• 相较于 DataParallel，torch.distributed 传输的数据量更少，因此速度更快，效率更高。

同时 pytorch 官方文档也建议使用 DistributedDataParallel（DDP）替换 DataParallel（DP），因此本文着重介绍 DistributedDataParallel（DDP）的相关内容。

DistributedDataParallel 多卡可以分为：

• 数据并行（常用方式，每张卡具有相同的模型和参数，训练时将 batch 数据拆分输入不同的模型中），
• 模型并行（将模型拆分，不同部分放置在不同的 GPU 上，并行计算），Workload Partitioning（将模型拆分，不同部分放置在不同的 GPU 上，串行计算）。

由于第一种数据并行方式是最常用，并且是官方实践性能最佳的方式，因此本文相关内容主要针对数据并行方式给出。

二、基础知识

GPU 并行计算基础处理流程：设置可见 GPU 并配置并行化参数，创建模型，通过 API 将模型并行化，将模型和数据搬到 GPU，进行前向和后向传播。

在前向传播中，会自动将 batch_size 切分后分配到可见的 GPU 上并行计算。

结束后，会有一台或其他方式收集前向传播计算结果并根据 loss 更新每块 GPU 上模型参数。

因此在多 GPU 训练时，可以：

• 提高整体 batch_size，
• 同时增加 learning_rate（一般为 batch_size 增加倍数的一半）。

1、概念

rank

• 多机多卡：代表某一台机器
• 单机多卡：代表某一块GPU

world_size【world_size = int(os.getenv(“WORLD_SIZE”, ‘1’)) 】

• 多机多卡：代表有几台机器
• 单机多卡：代表有几块GPU

local_rank

• 多机多卡：代表某一块GPU的编号
• 单机多卡：代表某一块GPU的编号

2、管理方式

管理方式

• group：进程组，默认情况下，只有一个组，一个 job 即为一个组，也即一个 world。
• world size：表示全局进程个数，一般和 GPU 数相同（单进程单GPU情况）。
• rank：表示进程序号，用于进程间通讯，表征进程优先级，序号一般从 0 到 world_size - 1。rank = 0 的主机为 master 节点。
• local_rank：进程内 GPU 编号，非显式参数，一般为一台主机内的 GPU 序号（从 0 到该机 GPU 数减一），由 torch.distributed.launch 内部指定。

3、常用方法

torch.cuda

• torch.cuda.is_available()：判断 GPU 是否可用
• torch.cuda.device_count()：计算当前可见可用的 GPU 数
• torch.cuda.get_device_name()：获取 GPU 型号，如 Tesla K80
• torch.cuda.manual_seed()：为当前 GPU 设置随机种子
• torch.cuda.manual_seed_all()：为所有可见可用 GPU 设置随机种子
• torch.cuda.current_device()：返回当前设备索引
• os.environ.setdefault(“CUDA_VISIBLE_DEVICES”, “2,3”)：设置实际可见的 GPU，等价于 os.environ[‘CUDA_VISIBLE_DEVICES’] = ‘2,3’
• torch.cuda.set_device()：作用与 os.environ.setdefault() 类似，官方建议使用 os.environ.setdefault()。但实际这个更好使 =_=

torch.distributed

• torch.distributed.get_world_size()：获取全局并行数
• torch.distributed.new_group()：使用 world 的子集，创建新组，用于集体通信等
• torch.distributed.get_rank()：获取当前进程的序号，用于进程间通讯。
• torch.distributed.local_rank()：获取本台机器上的进程的序号

其他

• torch.device()：创建 device 对象，如 torch.device(‘cpu’)，torch.device(‘cuda:1’)
• tensor.to()，module.to()：将 tensor 转换类型或者搬到 GPU（会重新创建一个新的 tensor），将 module 搬到 GPU （会复用之前的 module）

三、DDP的使用【torch.distributed】

官方建议用新的DDP，采用all-reduce算法，本来设计主要是为了多机多卡使用，但是单机上也能用。

DDP 多卡运行的本质还是通过创建多个进程来实现并行，但由于多个进程使用的是同一份代码，因此需要在代码中增加相关逻辑来指定进程与 GPU 硬件之间的关联关系。同时为了实现数据的并行化，需要为不同的进程（不同的 GPU）加载不同的数据，因此需要一个特殊的 data sampler 来实现，这个 DDP 通过 torch.utils.data.distributed.DistributedSampler 来实现。

与 DataParallel 的单进程控制多 GPU 不同，在 distributed 的帮助下，我们只需要编写一份代码，torch 就会自动将其分配给 $n$ 个进程，分别在 $n$ 个 GPU 上运行。

Pytorch 中分布式的基本使用流程如下：

1. 导入必要的包

   import torch.distributed as dist  # 分布式训练依赖
   from torch.utils.data.distributed import DistributedSampler  # 分布式训练依赖
   from torch.nn.parallel import DistributedDataParallel # 分布式训练依赖
   

2. 在使用 distributed 包的任何其他函数之前，需要使用 init_process_group 初始化进程组，同时初始化 distributed 包。

   # 分布式训练-新增: DDP backend初始化
   dist.init_process_group(backend='nccl')
   torch.cuda.set_device(args.local_rank) # 根据local_rank来设定当前使用哪块GPU
   args.world_size = int(os.getenv("WORLD_SIZE", '1')) # 单机多卡：代表有几块GPU
   args.global_rank = dist.get_rank() # 获取当前进程的序号，用于进程间通讯
   

3. 如果需要进行小组内集体通信，用 new_group 创建子分组

4. 创建分布式并行模型 DDP(model, device_ids=device_ids)，find_unused_parameters参数设置为True

   # 分布式训练-新增：定义并把模型放置到哪些GPU上，需要在调用`model=DDP(model)`前做
   device = torch.device("cuda", args.local_rank)
   model = import_module("moels." + model_name).Model().to(device) # 动态载入模型
   model = DistributedDataParallel(model, device_ids=[args.local_rank],find_unused_parameters=True)  # 前提是model已经.cuda() 了
   

5. 为数据集创建 Sampler

   # 加载数据集
   train_dataset = SampleDataset(config.train_path)
   val_dataset = SampleDataset(config.test_path)
   test_dataset = SampleDataset(config.test_path)
   
   # 分布式训练-新增：为数据集创建 Sampler
   train_sampler = DistributedSampler(train_dataset)
   val_sampler = DistributedSampler(val_dataset)
   test_sampler = DistributedSampler(test_dataset)
   

6. DataLoader使用sample参数，并且shuffle一定要设置为False

   # 数据迭代器；分布式训练-新增： sampler参数
   train_dataloader = DataLoader(dataset=train_dataset, batch_size=config.batch_size, shuffle=False, collate_fn=collate_fn, drop_last=True, sampler=train_sampler) 
   val_dataloader = DataLoader(dataset=val_dataset, batch_size=config.batch_size, shuffle=False, collate_fn=collate_fn, drop_last=True, sampler=val_sampler)
   test_dataloader = DataLoader(dataset=test_dataset, batch_size=config.batch_size, shuffle=False, collate_fn=collate_fn, drop_last=True, sampler=test_sampler)
   

7. 在 API 层面，pytorch 为我们提供了 torch.distributed.launch 启动器，用于在命令行分布式地执行 python 文件。使用启动工具 torch.distributed.launch 在每个主机上执行一次脚本，开始训练

   CUDA_VISIBLE_DEVICES=0,1,2,3 python -m torch.distributed.launch --nproc_per_node=4 main.py
   

   (pytorch) ninjia@ailian-Super-Server:~$ nvidia-smi
   Sat May 21 13:42:18 2022       
   +-----------------------------------------------------------------------------+
   | NVIDIA-SMI 455.23.05    Driver Version: 455.23.05    CUDA Version: 11.1     |
   |-------------------------------+----------------------+----------------------+
   | GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
   | Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
   |                               |                      |               MIG M. |
   |===============================+======================+======================|
   |   0  GeForce RTX 208...  On   | 00000000:18:00.0 Off |                  N/A |
   | 28%   55C    P2   225W / 250W |   9929MiB / 11019MiB |     90%      Default |
   |                               |                      |                  N/A |
   +-------------------------------+----------------------+----------------------+
   |   1  GeForce RTX 208...  On   | 00000000:3B:00.0 Off |                  N/A |
   | 30%   58C    P2   214W / 250W |   9307MiB / 11016MiB |     88%      Default |
   |                               |                      |                  N/A |
   +-------------------------------+----------------------+----------------------+
   |   2  GeForce RTX 208...  On   | 00000000:86:00.0 Off |                  N/A |
   | 29%   55C    P2   166W / 250W |  10123MiB / 11019MiB |     86%      Default |
   |                               |                      |                  N/A |
   +-------------------------------+----------------------+----------------------+
   |   3  GeForce RTX 208...  On   | 00000000:AF:00.0 Off |                  N/A |
   | 30%   57C    P2   123W / 250W |   9643MiB / 11019MiB |     88%      Default |
   |                               |                      |                  N/A |
   +-------------------------------+----------------------+----------------------+
                                                                                  
   +-----------------------------------------------------------------------------+
   | Processes:                                                                  |
   |  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
   |        ID   ID                                                   Usage      |
   |=============================================================================|
   |    0   N/A  N/A     94633      G   /usr/lib/xorg/Xorg                  4MiB |
   |    0   N/A  N/A     94707      G   /usr/bin/gnome-shell                0MiB |
   |    0   N/A  N/A    869045      C   ...3/envs/pytorch/bin/python     9921MiB |
   |    0   N/A  N/A    869046      C   ...3/envs/pytorch/bin/python        0MiB |
   |    0   N/A  N/A    869047      C   ...3/envs/pytorch/bin/python        0MiB |
   |    0   N/A  N/A    869048      C   ...3/envs/pytorch/bin/python        0MiB |
   |    1   N/A  N/A     94633      G   /usr/lib/xorg/Xorg                  9MiB |
   |    1   N/A  N/A     94707      G   /usr/bin/gnome-shell                3MiB |
   |    1   N/A  N/A    869045      C   ...3/envs/pytorch/bin/python        0MiB |
   |    1   N/A  N/A    869046      C   ...3/envs/pytorch/bin/python     9289MiB |
   |    1   N/A  N/A    869047      C   ...3/envs/pytorch/bin/python        0MiB |
   |    1   N/A  N/A    869048      C   ...3/envs/pytorch/bin/python        0MiB |
   |    2   N/A  N/A     94633      G   /usr/lib/xorg/Xorg                  4MiB |
   |    2   N/A  N/A     94707      G   /usr/bin/gnome-shell                0MiB |
   |    2   N/A  N/A    869045      C   ...3/envs/pytorch/bin/python        0MiB |
   |    2   N/A  N