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

DDP简介

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

1. 单卡训练回顾

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

  

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

2. 多卡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 个相同的模型再进入下一次迭代。不断重复上述过程。