【PyTorch教程】PyTorch分布式并行模块DistributedDataParallel(DDP)详解

• 本章的重点是学习如何使用 PyTorch 中的 Distributed Data Parallel (DDP) 库进行高效的分布式并行训练。以提高模型的训练速度。

DDP简介

• DDP 是PyTorch 中的 Distributed Data Parallel ，用于多卡加速模型训练。

1. 单卡训练回顾

• 在单卡训练中，你有一个模型在一块 GPU:0 上，DataLoader 产生一个输入批量 (InputBatch) 输入到模型 (model) 中。

  

• 然后模型执行前向传播来计算损失函数，接着执行反向传播来计算模型参数的梯度，然后优化器 (Optimizer) 使用这些梯度更新模型参数。

2. 与DataParallel比较

• PyTorch 中的 DistributedDataParallel 和 DataParallel 有什么区别呢？

1）DataParallel

• 优点：

  • 易于使用。对原代码的变动很少，只需要在模型实例化后，再添加一句代码，对模型对象 model 再封装一层即可：

  model = nn.DataParallel(model)
  

• 缺点：

  • 并不能提供最好的性能，在每个前向传播中复制模型；
  • DataParallel 是基于单进程多线程 (single-process multi-thread) 来进行并发计算。这样会受到 Python 的全局解释器锁 (GIL) 的限制，因而无法发挥最好的并发性能。
  • 如果为了获得更好的性能，可以考虑本节介绍的 DistributedDataParallel 。

2）DistributedDataParallel

• 优点：
  • 并发性能好。DistributedDataParallel 是基于多进程 (multi-process) 来进行并行计算的，因此各个模型副本不会受到 Python 全局解释器锁的限制。
  • 在 DDP 创建时就进行模型复制，而不是像 DataParallel 那样在前向传播阶段再复制模型，从而能加速训练。
• 缺点：
  • 比 DataParallel 要多一个代码步骤。就是要初始化多进程组 (init_process_group) 。

3. 多卡DDP训练

• 当我们把上述的训练划分到多块 (假设是 4 块) 不同的 GPUs 上时，DDP 在每个 GPU 上启动一个进程，每个进程上都有一份模型的本地副本。顾名思义，模型的所有副本和优化器 (Optimizer) 都彼此相同。不仅模型的初始参数相同，而且 Optimizer 所使用随机种子也相同。如下图所示：

  

• 在整个训练过程中，DDP 会在内部自动保持多卡之间模型和 Optimizer 的同步。

• 现在，每个 GPU 都拥有相同的模型。下面我们将要介绍 DDP 中的数据加载机制。

• 如前所述，我们从 DataLoader 中获得输入批量 (InputBatch) 。但这次，我们还使用了分布式采样器 (DistributedSampler) 。分布式采样器确保每个 GPU 接收到不同的输入，这就是我们在 DDP 中所强调的 “数据并行” (Data Parallel) 。如下图所示：

  

• 与单卡训练相比，我们现在能高效地同时处理四倍的数据！

• 在每个 GPU 中，模型都接收到不同的输入，各自进行前向传播和反向传播。并且，由于输入的不同，现在计算出的梯度也不同了。经过 Optimizer 优化器计算后，在 4 块 GPU 上会得到不同的参数。最终，我们得到四个不同的模型。

• 紧接着，DDP 会启动同步。通过 bucketed Ring-AllReduce 算法 将 4 块 GPU 上的梯度聚合起来。这个算法最酷的地方在于：它将梯度计算和通信结合起来。如下图所示：

  

• 在同步阶段中，不会等到所有梯度计算完成才开始同步。具体地说，当反向传播仍在进行时，它就会沿着环开始通信。这确保 GPU 一直保持运行而不会空载。4 个模型之间同步完成后的示意图如下所示：

  

• 同步完成后，每个 GPU 中的模型梯度都相同了。

• 然后，进行 Optimizer 步骤时会把所有 GPU 上的模型参数都更新成相同的值。所有 GPU 上的模型都是相同的，且继续保持同步。

• 这就是一个 DDP 训练的所有步骤，并且 4 个相同的模型再进入下一次迭代。不断重复上述过程。