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

我就是我 2023-09-29 13:34 28阅读 0赞

![在这里插入图片描述][af5968da566e4bef90ab9988349ff017.png_pic_center]

## 一、概述 ##

我们知道 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 较慢，不推荐使用**。

![在这里插入图片描述][2c92a3ca633e42a2ba4096177462cf5f.png_pic_center]

*  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 n 个进程，分别在  n 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<

[af5968da566e4bef90ab9988349ff017.png_pic_center]: https://img-blog.csdnimg.cn/af5968da566e4bef90ab9988349ff017.png#pic_center
[2c92a3ca633e42a2ba4096177462cf5f.png_pic_center]: https://img-blog.csdnimg.cn/2c92a3ca633e42a2ba4096177462cf5f.png#pic_center