分布式训练 - 单机多卡（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-machine multi-GPU DataParallel,
    # you would like to speed up training with the minimum code change.
    net = nn.DataParallel(net)

    # 2. define dataloader
    trainset = torchvision.datasets.CIFAR10(
        root="./data",
        train=True,
        download=True,
        transform=transforms.Compose(
            [
                transforms.RandomCrop(32, padding=4),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize(
                    (0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)
                ),
            ]
        ),
    )
    train_loader = torch.utils.data.DataLoader(
        trainset,
        batch_size=BATCH_SIZE,
        shuffle=True,
        num_workers=4,
        pin_memory=True,
    )

    # 3. define loss and optimizer
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(
        net.parameters(),
        lr=0.01,
        momentum=0.9,
        weight_decay=0.0001,
        nesterov=True,
    )

    print("            =======  Training  ======= \n")

    # 4. start to train
    net.train()
    for ep in range(1, EPOCHS + 1):
        train_loss = correct = total = 0

        for idx, (inputs, targets) in enumerate(train_loader):
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = net(inputs)

            loss = criterion(outputs, targets)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            train_loss += loss.item()
            total += targets.size(0)
            correct += torch.eq(outputs.argmax(dim=1), targets).sum().item()

            if (idx + 1) % 50 == 0 or (idx + 1) == len(train_loader):
                print(
                    "   == step: [{:3}/{}] [{}/{}] | loss: {:.3f} | acc: {:6.3f}%".format(
                        idx + 1,
                        len(train_loader),
                        ep,
                        EPOCHS,
                        train_loss / (idx + 1),
                        100.0 * correct / total,
                    )
                )

    print("\n            =======  Training Finished  ======= \n")

结果：

吐槽一下，本来想用四块GPU给训练看一下，谁知道哪个同学把 GPU0 吃的这么死，如果不添加修改使用GPU的话，直接会爆显存。各位童鞋，如果使用公共服务器的话千万不要这样干，很容易被别人喷，小心别人提着大刀来 gan 你 。。
画红色圈圈的是我的实验开启的进程号和所占内存，很明显看出来 GPU 1 使用的显存多，这也印证了有一块GPU会肩负着其他重要任务（进行参数汇总，梯度更新等等）。这也就是DP模式的缺点所在，如果模型很大，这个差别会特别明显，造成很严重的资源浪费。

既然 DataParallel 有缺点，那么有没有相对于 DataParallel 好一些的方法呢？ 有的，那就是DistributedDataParallel，下面我们来详细说一下DistributedDataParallel方法。

2.DistributedDataParallel(DDP)

DataParallel：单进程控制多 GPU。
DistributedDataParallel：多进程控制多 GPU，一起训练模型。

2.1 介绍

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

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

和单进程训练不同的是，多进程训练需要注意以下事项：

• 在喂数据的时候，一个batch被分到了好几个进程，每个进程在取数据的时候要确保拿到的是不同的数据（DistributedSampler）；
• 要告诉每个进程自己是谁，使用哪块GPU（args.local_rank）；
• 在做BatchNormalization的时候要注意同步数据。

2.2 使用方式

很多东西先感知它，详细的解释和底层原理请看我写的其他文章。

2.2.1 启动方式的改变

（这个很重要，也是很多小白搞不懂的地方，听我慢慢讲）
在我们使用 pycharm 或者其他 IDE 写完程序单机右键点击运行的时候，系统会启动一个进程进行处理。而我们 DistributedDataParallel 方式是多进程控制多 GPU 的形式，所以你再直接单机右键点击运行必然会报错，报错如下：

恭喜你，如果是这种报错是你的启动方式不对，需要使用命令行进行启动。

启动方式为：

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

CUDA_VISIBLE_DEVICES=0,1,2,3 选择要使用的 GPU （不设置的话默认使用所都的 GPU）

python -m 的使用可以点击查看博主的这篇文章,详细说明了Python -m 的使用案例和作用。

torch.distributed.launch：启动方式，想详细了解原理的同学可以点击博主写的另一篇文章，分布式训练启动工具—Launch utility

– nproc_per_node 参数用于指定为当前主机创建的进程数，由于我们现在说的是单机多卡，所以这里的node数量为1，-- nproc_per_node设置为所使用的 GPU 的数量。（博主使用了4块GPU，所以设置成了4）

2.2.2 初始化

在启动器为我们启动python脚本后，在执行过程中，启动器会将当前进程的（其实就是 GPU的）index 通过参数传递给 python，我们可以这样获得当前进程的 index：即通过参数 local_rank 来告诉我们当前进程使用的是哪个GPU，用于我们在每个进程中指定不同的device：

def parse():
    parser = argparse.ArgumentParser()
    parser.add_argument('--local_rank', type=int, default=0，help='node rank for distributed training')
    args = parser.parse_args()
    return args

def main():
    args = parse()
    torch.cuda.set_device(args.local_rank)
    torch.distributed.init_process_group(
        backend='nccl'
        init_method='env://'
    )
    device = torch.device(f'cuda:{args.local_rank}')
    ...

对于上边代码中不熟悉argparse这个库使用的童鞋可以点击查看博主写Python中argparse模块的使用这篇文章，会快速带你熟悉argparse库的使用，在这里就不在详细说明了。

其中 torch.distributed.init_process_group 用于初始化GPU通信方式（NCCL）和参数的获取方式（env代表通过环境变量）。

[详细的 GPU 通信方式和参数获取方式会在后序文章中详细写到，对于单纯使用分布式知道有这么回事就可以，不宜深究]

使用 init_process_group 设置GPU之间通信使用的后端和端口，通过 NCCL 实现 GPU 通信。

2.2.3 DataLoader

在读取数据的时候，我们要保证一个batch里的数据被均摊到每个进程上，每个进程都能获取到不同的数据，但如果我们手动去告诉每个进程拿哪些数据的话太麻烦了，PyTorch也为我们封装好了这一方法。之后，使用 DistributedSampler 对数据集进行划分。如此前我们介绍的那样，它能帮助我们将每个 batch 划分成几个 partition，在当前进程中只需要获取和 rank 对应的那个 partition 进行训练。

所以我们在初始化 DataLoader 的时候需要使用到 torch.utils.data.distributed.DistributedSampler 这个特性：

train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler)

这样就能给每个进程一个不同的 sampler，告诉每个进程自己分别取哪些数据。

2.2.4 模型的初始化

和 torch.nn.DataParallel 的方式一样，我们对于模型的初始化也是简单的一句话就行了：

model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank])

使用 DistributedDataParallel 包装模型，它能帮助我们为不同 GPU 上求得的梯度进行 all reduce（即汇总不同 GPU 计算所得的梯度，并同步计算结果）。all reduce 后不同 GPU 中模型的梯度均为 all reduce 之前各 GPU 梯度的均值。

2.2.5 同步BN

为什么要同步BN？

现有的标准 Batch Normalization 因为使用数据并行（Data Parallel），是单卡的实现模式，只对单个卡上对样本进行归一化，相当于减小了批量大小（batch-size）（详见BN工作原理部分）。 对于比较消耗显存的训练任务时，往往单卡上的相对批量过小，影响模型的收敛效果。 跨卡同步 Batch Normalization 可以使用全局的样本进行归一化，这样相当于‘增大‘了批量大小，这样训练效果不再受到使用 GPU 数量的影响。 最近在图像分割、物体检测的论文中，使用跨卡BN也会显著地提高实验效果，所以跨卡 BN 已然成为竞赛刷分、发论文的必备神器。
现在在 PyTorch 中已经提供了实现同步 BN 的方法:

model= torch.nn.SyncBatchNorm.convert_sync_batchnorm(model)

详细的使用我没有验证过，看过有些文章说到使用PyTorch自带的同步 BN 有缺陷，有时间验证一下，看看怎么样，在 apex 中使用同步 BN 还是很方便的，在后边的文章中会详细说到。

总结

至此，我们就可以使用 torch.distributed 给我们带来的多进程训练的性能提升了，汇总代码结果如下：

# main.py
import torch
import argparse
import torch.distributed as dist

parser = argparse.ArgumentParser()
parser.add_argument('--local_rank', default=-1, type=int,
                    help='node rank for distributed training')
args = parser.parse_args()

dist.init_process_group(backend='nccl')
torch.cuda.set_device(args.local_rank)

train_dataset = ...
#每个进程一个sampler
train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)

train_loader = torch.utils.data.DataLoader(train_dataset, batch_size=..., sampler=train_sampler)

model = ...
model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.local_rank])

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()

在使用时，调用 torch.distributed.launch 启动器启动：

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

DDP基本用法（代码编写流程）

• 使用 torch.distributed.init_process_group 初始化进程组
• 使用 torch.utils.data.distributed.DistributedSampler 创建 DataLoader
• 使用 torch.nn.parallel.DistributedDataParallel 创建 分布式模型
• 调整其他必要的地方(tensor放到指定device上，S/L checkpoint，指标计算等)
• 使用 torch.distributed.launch / torch.multiprocessing 或 slurm 开始训练

实例：

单机多卡实验（DDP）

"""
(MNMC) Multiple Nodes Multi-GPU Cards Training
    with DistributedDataParallel and torch.distributed.launch
Try to compare with [snsc.py, snmc_dp.py & mnmc_ddp_mp.py] and find out the differences.
"""

import os

import torch
import torch.distributed as dist
import torch.nn as nn
import torchvision
import torchvision.transforms as transforms
from torch.nn.parallel import DistributedDataParallel as DDP

BATCH_SIZE = 256
EPOCHS = 5


if __name__ == "__main__":

    # 0. set up distributed device
    rank = int(os.environ["RANK"])
    local_rank = int(os.environ["LOCAL_RANK"])
    torch.cuda.set_device(rank % torch.cuda.device_count())
    dist.init_process_group(backend="nccl")
    device = torch.device("cuda", local_rank)

    print(f"[init] == local rank: {local_rank}, global rank: {rank} ==")

    # 1. define network
    net = torchvision.models.resnet18(pretrained=False, num_classes=10)
    net = net.to(device)
    # DistributedDataParallel
    net = DDP(net, device_ids=[local_rank], output_device=local_rank)

    # 2. define dataloader
    trainset = torchvision.datasets.CIFAR10(
        root="./data",
        train=True,
        download=False,
        transform=transforms.Compose(
            [
                transforms.RandomCrop(32, padding=4),
                transforms.RandomHorizontalFlip(),
                transforms.ToTensor(),
                transforms.Normalize(
                    (0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)
                ),
            ]
        ),
    )
    # DistributedSampler
    # we test single Machine with 2 GPUs
    # so the [batch size] for each process is 256 / 2 = 128
    train_sampler = torch.utils.data.distributed.DistributedSampler(
        trainset,
        shuffle=True,
    )
    train_loader = torch.utils.data.DataLoader(
        trainset,
        batch_size=BATCH_SIZE,
        num_workers=4,
        pin_memory=True,
        sampler=train_sampler,
    )

    # 3. define loss and optimizer
    criterion = nn.CrossEntropyLoss()
    optimizer = torch.optim.SGD(
        net.parameters(),
        lr=0.01 * 2,
        momentum=0.9,
        weight_decay=0.0001,
        nesterov=True,
    )

    if rank == 0:
        print("            =======  Training  ======= \n")

    # 4. start to train
    net.train()
    for ep in range(1, EPOCHS + 1):
        train_loss = correct = total = 0
        # set sampler
        train_loader.sampler.set_epoch(ep)

        for idx, (inputs, targets) in enumerate(train_loader):
            inputs, targets = inputs.to(device), targets.to(device)
            outputs = net(inputs)

            loss = criterion(outputs, targets)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

            train_loss += loss.item()
            total += targets.size(0)
            correct += torch.eq(outputs.argmax(dim=1), targets).sum().item()

            if rank == 0 and ((idx + 1) % 25 == 0 or (idx + 1) == len(train_loader)):
                print(
                    "   == step: [{:3}/{}] [{}/{}] | loss: {:.3f} | acc: {:6.3f}%".format(
                        idx + 1,
                        len(train_loader),
                        ep,
                        EPOCHS,
                        train_loss / (idx + 1),
                        100.0 * correct / total,
                    )
                )
    if rank == 0:
        print("\n            =======  Training Finished  ======= \n")

命令行启动代码：

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

我们通过查看显卡的使用情况可以明显看出，该程序使用了4块 GPU 启动了4个不同的进程，并且每个进程消耗的显存是一样的，解决了负载不均衡问题（PID=16630 是使用nn.Dataparell 方法进行数据并行处理，也是使用了4块GPU，但是很明显看出0号GPU上消耗显存巨大，说明别人训练的模型负载均衡很严重），通过上边的对比，很明显看出使用DDP方法对于分布式训练是有很多好处的，建议大家使用。

结果：

参考链接：

• https://zhuanlan.zhihu.com/p/158375055