分布式训练 - 单机多卡（DP和DDP）

起初为调用大规模的模型训练，单卡GPU是不够使用的，需要借用服务器的多GPU使用。就会涉及到单机多卡，多机多卡的使用。在这里记录一下使用的方式和踩过的一些坑。文中若有不足，请多多指正。

由于分布式的内容较多，笔者准备分几篇来讲一次下深度学习的分布式训练，深度学习的框架使用的是Pytorch框架。

----1.分布式训练的理论基础

----2.GPU训练

----3.单机多卡的使用 DP,DDP

----4.多机多卡的使用 DDP

在GPU训练文章中我们已经了解到了多GPU的训练，最简单的是单机多卡操作torch.nn.DataParallel。只需要几步就可以实现多GPU的使用，加快训练的速度，但是很遗憾的是，这个操作还不是很优秀，于是今天我们来聊一聊比较优秀的操作~~

整个分布式文章系列的链接都在上边啦，有需要的小伙伴点击链接就可以看到其他的知识啦！

在进入文章主题之前，先问大家两个问题，或许你也会有这样的疑惑：

问1：为什么非要单机多卡？ why？

答：加速神经网络模型的训练最简单直接的办法是使用 GPU，往往一块 GPU 的显存容量是有限的，如果遇到了比较大的模型或者是参数量巨大的情况下会直接爆显存，最简单粗暴的办法就是加 GPU，一块不够我就使用两块，两块还不够我就再加两块（感觉经费在燃烧）。这会有一个核心的问题，为了使用多块 GPU 进行训练，我们必须想一个办法在多个 GPU 上进行分发数据和模型，并且协调训练的过程。

问2：上一篇讲得单机多卡操作torch.nn.DataParallel，哪里不好？

答：要回答这个问题我们得先简单回顾一下torch.nn.DataParallel，要使用这玩意，我们将模型和数据加载到多个 GPU 中，控制数据在 GPU 之间的流动，协同不同 GPU 上的模型进行并行训练。具体怎么操作？

简单回顾一下上篇说到的怎么快捷使用 nn.DataParallel：（只需要几步，让你的模型如虎添翼，速度加倍，快乐加倍）

使用 torch.nn.DataParrallel() 进行并行数据处理的流程步骤是：
构造模型net： net= FooNet(neural_num=3, layers=3)
并行化：net = nn.DataParallel(net)
将模型转移到GPU上：net.to(device)
数据迁移：inputs, labels = inputs.to(device), labels.to(device)
总之，流程是：构建模型 --> 并行化 --> 数据，模型传到GPU上面去。

今天，我们来详细说说这个包装模型：

1.DataParallel(DP)

DataParallel是基于Parameter server的算法，负载不均衡的问题比较严重，有时在模型较大的时候（比如bert-large），reducer的那张卡会多出3-4g的显存占用。

model = nn.DataParallel(model.to(device), device_ids=gpus, output_device=gpus[0])

DataParallel 需要设置的参数有：gpus=[0,1,2,3]

• 参与训练的 GPU 有哪些，device_ids=gpus。 如果不设置将默认使用所有的 GPU。
• 用于汇总梯度的 GPU 是哪个，output_device=gpus[0] 。如果不设置将默认输出物理设备序号第一个。

DataParallel 会自动帮我们将数据切分 load 到相应 GPU，将模型复制到相应 GPU，进行正向传播计算梯度并汇总（是不是非常的人性化）。值得注意的是，模型和数据都需要先导入进 GPU 中，DataParallel 的 module 才能对其进行处理，否则会报错：

# 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 = 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.to(device)将模型迁移到GPU里面，images.cuda，target.cuda把数据迁移到GPU里面。
nn.DataParallel(model.to(device), device_ids=gpus, output_device=gpus[0])包装模型

缺点：

• 在每个训练批次（batch）中，因为模型的权重都是在 一个进程上先算出来 然后再把他们分发到每个GPU上，所以网络通信就成为了一个瓶颈，而GPU使用率也通常很低。核心是PS算法的问题（详细理论问题请看第一篇文章或点击）。
• 除此之外，nn.DataParallel 需要所有的GPU都在一个节点（一台机器）上，且并不支持 Apex 的混合精度训练。

一句话，一个进程算权重使通信成为瓶颈，nn.DataParallel慢而且不支持混合精度训练。

实例：

单机多卡实验（DP）

Single Node Multi-GPU Cards Training (with DataParallel)，源码见 snmc_dp.py, 与上一篇文章说到使用 GPU 训练模型 snsc.py 对比一下可知，DP只需要花费最小的代价，既可以使用多卡进行训练 (其实就一行)，但是因为GIL锁的限制，DP的性能是低于DDP（后边会说到）的。

"""
(SNMC) Single Node Multi-GPU Cards Training (with DataParallel)
Try to compare with smsc.py and find out the differences.
"""
import torch
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
import os

BATCH_SIZE = 256
EPOCHS = 5
# 设置了一下GPU 的使用
os.environ['CUDA_VISIBLE_DEVICES']='1,2,3'

if __name__ == "__main__":

    # 1. define network
    device = "cuda"
    net = torchvision.models.resnet18(pretrained=False, num_classes=10)
    net = net.to(device=device)
    # Use single