PyTorch下的多卡训练

nn.DataParallel

DataParallel只能实现单个主机多GPU的训练，下图是由HuggingFace 制作的DataParallel的原理图。

DataParallel

mini-batch在 GPU:0 上进行分组，并将分好组的mini-batch分发到各个GPU上
把整个model复制到各个GPU上
各个GPU上进行各自的forward pass
把loss收集到 GPU:0 上，计算gradient 1, 2, 3, 4
把gradient加起来，用optimizer更新GPU:0上的模型权重。
最后回到步骤一，分数据，复制新的模型...

相当于，DataParallel的“并行计算”只存在于forward pass中，梯度回传以及模型参数的更新是在一个GPU上进行的。

PyTorch例子：

import torch
import torch.nn as nn

# 假设我们有一个简单的CNN 叫做 ConvNet
# 训练时使用多卡只需要使用：
def train(gpu, args):
    model = ConvNet()
    model = nn.DataParallel(model)
    torch.cuda.set_device(gpu)
    model.cuda(gpu)
    # 以下省略

Q：如果loss计算和back propagation只在主GPU上进行，那么是不是主显存很容易爆？

因为 loss = criterion(prediction, target) 这个语句，当需要预测的样本非常多的时候，实际上是非常占用显存的。

解决多GPU负载不均匀问题

在这篇blog 里面提到上述问题。解决方案是，让每个GPU单独计算loss，再将他们合并在主GPU上计算gradient。
所以可以修改model的forward函数，本来forward是返回模型的预测结果，但是如果我们在forward里面就把loss计算也写进去，这样最后汇总到主GPU的时候，就只会汇总loss，大大减少了主GPU的显存的消耗。以下code来自于：link

# 修改model里面的forward函数
def forward(self, x, target, args):
    features = self.extract_features(x)
    if target is None:
        return features
    classification_res = self.classifier(features)
    """
    forward通常到这里就结束了，然后 return classification_res，
    如果使用DataParallel，返回了output之后，会在主GPU上合并成一个list，
    每个GPU返回batch_size / n_gpus 个样本。样本量很大时非常占用空间。
    """
    
    """
    这时候我们可以直接在forward函数里面先计算loss，然后直接返回loss，
    或者是其他任何你想要返回的变量。最后它们都会汇总于主GPU。
    """
    # compute loss
    criterion = nn.CrossEntropyLoss()
    loss = criterion(classification_res, target)
    return loss

nn.DistributedDataParallel

每个GPU只有一个进程，并且一个进程控制一个model。这些不同的GPUs可以在同一个机器或者不同的机器。它们之间的信息交流只有gradient。

训练过程中，每个进程都从硬盘中加载自己的mini-batch并将其传入GPU。

每个GPU都独立完成forward pass，使用自己加载的mini-batch来计算损失函数的梯度。
通过GPU间的通信计算梯度的平均值
back propagation，更新模型

注意，上述过程的第二第三步是同时进行的，也就是说，每到模型的一层layer，每个GPU先单独计算自己的gradient，然后所有GPU之间进行一次gradient平均值，然后大家一起回传同样的gradient到下一层layer。这样让每个GPU上的模型的权重更新保持一致。

其中第二步中，用到了All Reduce算法来计算多GPU之间的梯度均值。更多阅读资料在文末的Reference中

模型并行 vs 数据并行

模型并行：不同的GPU输入相同的数据，运行模型的不同部分，比如多层网络的不同层。
数据并行：不同的GPU输入不同的数据，运行相同的完整的模型 model = nn.DataParallel(model)

当模型非常非常大，一张GPU已经存不下的时候，可以使用模型并行，把模型的不同部分交给不同的机器负责。但是缺点：

会带来很大的通信开销
模型并行的各个部分存在一定的依赖，规模伸缩性差。

同步更新 vs 异步更新

对于数据并行来说，由于每个GPU都负责一部分数据，那就涉及到更新参数的方式：

同步更新：每个batch所有GPU计算完成之后，再统一计算新的权值，然后所有GPU同步新值后，再进行下一轮计算。
异步更新：每个GPU计算玩梯度后，无需等待其他更新，立即更新整体权重并同步
一般使用数据并行+同步更新