在芯片性能提升有限的本日，分布式演习成为了应对超大规模数据集和模型的紧张方法。本文将向你先容盛行深度学习框架 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 模型。
