PyTorch如何加速数据并行训练？分布式秘籍大揭秘

在芯片性能提升有限的今天，分布式训练成为了应对超大规模数据集和模型的主要方法。本文将向你介绍流行深度学习框架 PyTorch 最新版本（ v1.5）的分布式数据并行包的设计、实现和评估。

论文地址：https://arxiv.org/pdf/2006.15704.pdf

PyTorch 是深度学习研究和应用中广泛使用的科学计算包。深度学习的最新进展证明了大型数据集和大型模型的价值，这需要扩展模型训练更多计算资源的能力。

同时，由于简明的原理和广泛的适用性，数据并行已经成为了分布式训练的一种热门方案。通常，分布式数据并行技术会在每个计算资源上复制模型以独立生成梯度，然后在每次迭代时传递这些梯度以保持模型副本的一致性。尽管该技术在概念上很简单，但计算和通信之间的细微依赖关系使优化分布式训练的效率变得不简单。

因此，在这篇论文中，来自 Facebook AI 和华沙大学的研究者介绍了 PyTorch 分布式数据并行模型的设计、实现以及评估。

从 v1.5 开始，PyTorch 自身提供了几种加速分布数据并行的技术，包括分桶梯度（bucketing gradients）、通信重叠计算（overlapping computation with communication）以及跳过梯度同步（skipping gradient synchronization）。相关评估结果显示，在配置正确的情况下，PyTorch 分布式数据并行模型可以用 256 个 GPU 达到接近线性的可扩展性。

接下来，我们来看 PyTorch 分布式数据并行训练的模型设计、具体实现和效果评估。

系统设计

PyTorch 提供了一个数据分布式并行（DistributedDataParalle, DDP）模型来帮助实现在多个进程和机器的并行训练。在分布训练期间，每个模型都有自己的本地模型副本和本地优化器。就纠错而言，分布式数据并行训练和本地训练在数学上必须是等价的。

下图 1 描述了 DDP 构造块的组成，其中包含一个 Python API 前端和 C++ 梯度下降核心算法，并采用了 c10d 聚合通信库。

Python API 前端

在设计 API 时，研究者制定了以下两个设计目标来达到必要的功能：

非侵入式：对应用提供的 API 必须是非侵入式的；

拦截式：API 需要允许拦截各种信号并立即触发适当的算法。

分布式数据并行化旨在使用更多的计算资源来加速训练。

根据以上需求，研究者用 nn.Module 实现了分布式数据并行。nn.Module 采用本地模型作为构造函数的参数，并在反向传播中透明地同步梯度。下面的代码是使用 DDP 模型的示例：

梯度下降

研究者阐述了在 PyTorch 上进行分布式数据并行训练的几种梯度降低技术。DDP 中的梯度下降算法已经有了新的改进。为了介绍当前实现的结构，研究者从一个简单的初始方案（naive solution）开始，逐步介绍更多复杂的版本，最终在 PyTorch v1.5.0 上使用当前版本。

初始方案

DDP 首先校正了所有的训练进程，以保证各个进程：

从相同的模型状态开始； 每次迭代花费同样多的梯度。

为了完成第二点，初始方案在进行本地反向传播之后、更新本地参数之前插入了一个梯度同步环节。幸运的是，PyTorch 的 autograd 引擎能够接受定制的 backward 钩子（hook）。DDP 可以注册 autograd 钩子来触发每次反向传播之后的计算。然后，它会使用 AllReduce 聚合通信来号召计算所有进程中每个参数的平均梯度，并且把结果写回梯度 tensor。

初始方案足以完成想要的目标，但存在两项性能缺陷。聚合通信在小型 tensor 上性能表现很差，这种缺陷在带有大量小参数的大型模型上尤为突出。由于两者之间存在界限，分别进行梯度计算和同步化会造成通信重叠计算机会的缺失。

梯度分桶（bucketing ）

梯度分桶的观点是受聚合通信在大型 tensor 上更加高效的启发而提出的。

下图 2(a)和 (b) 给出的定量视图展示了在每个 AllReduce 中参数数目不同的情况下，AllReduce 60M torch 的 float32 参数的完整执行时间：

这些实验表明，不用等到每个梯度 tensor 都可用时再启动 AllReduce，DDP 在等待较短的时间并将多个梯度存储到一个 AllReduce 操作中时，就可以实现更高的吞吐量和更短的延迟。

通信重叠计算

在使用分桶的情况下，DDP 只需在启动通信之前在同一个 bucket 中等待所有的内容。在这样的设置下，在反向传播的最后触发 AllReduce 就显得不足了。因此需要对更加频繁的信号做出相应，并且更加迅速地启动 AllReduce。因此，DDP 为每个梯度累加器都注册了 autograd 钩子。

下图 3(a)的示例中，两个竖直轴表示时间，虚线代表梯度准备就绪的时间。进程 1 中，4 个梯度按顺序计算。进程 2 中，g_2 在 g_3 和 g_4 之后计算；图 3(b)的示例中，梯度 g_3 对应的参数在一次迭代中被跳过了，导致 g_3 的就绪信号缺失。

为了解决这个问题，DDP 遍历了前向传播的输出 tensor 中的 autograd 图以找到涉及到的所有参数。涉及到 tensor 的就绪状态足以充当反向传播完成的信号。

以下算法 1 给出了 DDP 的伪代码：

下图 4 展示了 DDP 在前向传播和反向传播过程中如何与本地模型交互：

梯度累加

此外，DDP 无法分辨应用程序是计划在反向传播之后立即调用 optimizer.step()还是通过多次迭代累加梯度。因此，研究者需要为这个用例再引入一个接口（即 no sync）。以下是样例代码片段：

聚合通信

DDP 是在集合通信库基础上建立的，包括 3 个选项 NCCL、Gloo 和 MPI。DDP 采用了来自这三个库的 API，并将它们封装进同一个 ProcessGroup API 中。

由于所有的通信都是聚合操作，因此所有的 ProcessGroup 实例上的后续操作必须和其类型匹配并遵循相同的顺序。对所有的库使用同一个 ProcessGroup API 允许研究者在相同的 DDP 实现上试验不同的通信算法。

如果单一 NCCL、Gloo 或 MPI 的 ProcessGroup 无法使链路容量达到饱和，通过使用循环的 ProcessGroups，DDP 可以获得更高的带宽利用率。

具体实现

DDP 的实现在之前的几个版本中已经改进了多次。研究者介绍了当前 PyTorch v1.5.0 的状态。DDP 同时在 Python 和 C++ 上都可以实现，Python 开放了 API 并组成了非性能关键因素组件，而 C++ 提供了核心梯度下降算法。Python API 通过 Pybind11 的 API 调用了 C++ 内核。

Python 前端

Python 前端中的实现细节决定了 DDP 的行为。可配置的 Knobs 在 DDP 构造函数 API 中开放。具体包括：

分组处理以找出 DDP 中运行 AllReduce 的进程组实例，它能够帮助避免与默认进程组混淆； bucket_cap_mb 控制 AllReduce 的 bucket 大小，其中的应用应调整 knob 来优化训练速度； 找出没有用到的参数以验证 DDP 是否应该通过遍历 autograd 图来检测未用到的参数。

本地模型中的 Model Device Affinity 也能控制 DDP 的行为，尤其是当模型因为太大而需要跨越多个设备运行时，更是如此。对于大型模型，模型的每一层可以放在不同的设备上，使用 Tensor.to(device) API 可以将中间输出从一个设备转移到另一个上。DDP 也可以在多个模型上运行。

当层（例如 BatchNorm）需要跟踪状态，例如运行方差和均值时，模型缓冲器（buffer）是非常必要的。DDP 通过让 rank 为 0 的进程获得授权来支持模型缓冲器。

核心梯度下降

开发过程中的主要工作就是梯度降低，它也是 DDP 中决定性能的关键步骤。这个在 reducer.cpp 中的实现有 4 个主要的组成部分：构建 parameter-to-bucket map、安装 autograd 钩子，启动 bucket AllReduce 以及检测全局未用过的参数。

Parameter-to-Bucket Mapping 已经对 DDP 的速度有了相当大的影响。在每次反向传播中，tensor 从全部的参数梯度到 bucket 被复制，平均梯度在 AllReduce 之后又被复制回 tensor。

Autograd Hook 是 DDP 反向传播的进入点。在构造期间，DDP 遍历模型中的所有参数，找出每个参数的梯度累加器，并且为每个梯度累加器安装相同的 post-hook 函数。当相应的梯度准备就绪时，梯度累加器会启用 post hook，并且当整个 bucket 准备好启动 AllReduce 操作时，DDP 会确定启用。

Bucket Allreduce 是 DDP 中通信开销的主要来源。默认情况下，bucket 的大小是 25MB。

实验评估

研究者展示了使用专属 32-GPU 集群和共享权限时 PyTorch DDP 的评估结果，其中 GPU 部署在 4 台服务器，并通过迈络思 MT27700 ConnectX-4 100GB/s 的网卡连接。每台服务器配有 8 个英伟达 Tesla V100 GPU。

下图 5 展示了一台服务器上 8 个 GPU 的互连方式：

接下来，研究者利用 ResNet50 和 BERT 这两个流行的模型度量了 PyTorch DDP 在每次迭代时的延迟和可扩展性，并且大多数实验使用随机生成的合成输入和标签，这对于比较每次迭代时的延迟来说足够了。

延迟故障

为了验证通信重叠计算的有效性，下图 6 展示了 ResNet50 和 BERT 模型分别使用 NCCL 和 Gloo 反向传递时的延迟故障。所有实现都用到了 4 台服务器上的 32 个 GPU。

结果显示，在 PyTorch DDP 训练时，反向传递是耗时最长的步骤，这是因为 AllReduce 通信（即是梯度同步）在这一过程中完成。

Bucket 大小

bucket 大小是一个重要的配置选项。根据经验，出于最大努力估计，bucket_cap_mb 的默认值是 25MB。研究者使用两台机器上的 16 个 GPU 比较不同 bucket 大小下每次迭代的延迟。另一个极端是在短时间内传递全部的梯度，结果如下图 7 所示。

下图 8 给出了相同设置下、32 个 GPU 上的实验结果。在这种情况下，离群值（outlier）的跨度更大，这并不意外。因为在有更多参与者的情况下，同步必然要花费更长的时间，并且 strangler 的影响更明显。

可扩展性

为了理解 DDP 的可扩展性，研究者用多达 256 个 GPU 上的 NCCL 和 Gloo 后端来度量 ResNet50 和 BERT 每次迭代的训练延迟。结果如下图 9 所示。

下图 10 给出了每 1、2、4 和 8 次迭代进行梯度下降时每次迭代的平均延迟。

除了每次迭代延迟，测量收敛速度以验证加速度是否会因收敛放缓而被消除也非常关键。实验采用 MNIST 数据集来训练 ResNet。学习率设置为 0.02，批处理大小是 8。结果如下图 11(a)所示；图 11(b)是将批处理大小设为 256，学习率设为 0.06 的测量结果。

循环分配（Round-Robin）进程组

PyTorch 分布式包支持将 Round-Robin 进程组和多个 NCCL 或 Gloo 进程组组合在一起，从而按照 Robin-Robin 顺序向各个进程组实例分配聚合通信。

下图 12 展示了使用 1、3 和 5 个 NCCL 或 Gloo 进程组的 Round-Robin 进程组每次迭代的延迟。最显著的加速是使用 NCCL 后端的 BERT 模型。