在芯片性能提升有限的今天，分布式训练成为处理超大规模数据集和模型的主要方式。本文将为大家介绍流行的深度学习框架PyTorch最新版本（v1.5）的分布式数据并行包的设计、实现和评估。论文地址：https://arxiv.org/pdf/2006.15704.pdfPyTorch是一个广泛应用于深度学习研究和应用的科学计算包。深度学习的最新进展证明了大型数据集和大型模型的价值，这需要扩展模型的能力以使用更多计算资源进行训练。同时，由于其简洁的原理和广泛的适用性，数据并行已经成为分布式训练的热门解决方案。通常，分布式数据并行会在每个计算资源上复制模型以独立生成梯度，然后在每次迭代中传递这些梯度以保持模型副本的一致性。尽管该技术在概念上很简单，但计算和通信之间微妙的依赖关系使得优化分布式训练的效率变得非常重要。因此，在本文中，来自FacebookAI和华沙大学的研究人员展示了PyTorch中分布式数据并行模型的设计、实现和评估。从v1.5开始，PyTorch本身提供了几种加速分布式数据并行的技术，包括分桶梯度、重叠计算与通信和跳过梯度同步。相关评估结果表明，在正确的配置下，PyTorch分布式数据并行模型可以实现256个GPU的近线性可扩展性。接下来我们看一下PyTorch分布式数据并行训练的模型设计、实现和效果评估。系统设计PyTorch提供了数据分布式并行（DistributedDataParalle，DDP）模型，帮助实现多进程、多机器的并行训练。在分布式训练期间，每个模型都有自己的模型本地副本和本地优化器。在纠错方面，分布式数据并行训练和本地训练在数学上必须是等价的。下面的图1描述了DDP构建块的组成，它由PythonAPI前端和C++梯度下降核心算法组成，使用c10d聚合通信库。PythonAPI前端在设计API时，研究人员制定了以下两个设计目标，以实现必要的功能：非侵入式：提供给应用程序的API必须是非侵入式的；拦截：API需要允许拦截各种信号并立即触发相应的算法。分布式数据并行化旨在使用更多的计算资源来加速训练。根据以上需求，研究人员使用nn.Module实现分布式数据并行。nn.Module将本地模型作为构造函数的参数，并透明地同步反向传播中的梯度。下面的代码是使用DDP模型的示例：梯度下降研究人员展示了几种用于PyTorch分布式数据并行训练的梯度下降技术。DDP中的梯度下降算法有了新的改进。为了介绍当前实现的结构，研究人员从一个简单的原始解决方案开始，逐渐引入更复杂的版本，最后在PyTorchv1.5.0上使用当前版本。初始方案DDP首先修正所有的训练过程，保证每个过程：从相同的模型状态开始；每次迭代花费相同数量的梯度。为了完成第二点，初始方案在执行局部反向传播之后和更新局部参数之前插入梯度同步链接。幸运的是，PyTorch的autograd引擎接受自定义的向后挂钩。DDP可以注册autogradhooks以在每次反向传播后触发计算。然后它使用AllReduce聚合通信来调用所有进程中每个参数的平均梯度，并将结果写回梯度张量。初始方案足以实现预期目标，但存在两个性能缺陷。聚合通信在小张量上表现不佳，这种缺陷在具有大量小参数的大模型上尤为明显。由于两者之间存在边界，因此分别执行梯度计算和同步会导致错过通信重叠计算的机会。梯度分桶（bucketing）梯度分桶的思想是受到聚合通信在大张量上更高效这一事实的启发。下面图2(a)和(b)给出的定量视图展示了AllReduce60Mtorch的float32参数在每个AllReduce中的参数个数不同时的完整执行时间：这些实验表明不需要等待对于每个梯度张量通过在两者都可用时启动AllReduce，DDP可以实现更高的吞吐量和更低的延迟，同时等待更短的时间并将多个梯度存储到一个AllReduce操作中。CommunicationOverlapComputing使用分桶，DDP只需等待同一个桶中的所有内容，然后再发起通信。在这样的设置中，在反向传播结束时触发AllReduce是不够的。因此，需要响应更频繁的信号，更快地启动AllReduce。因此，DDP为每个梯度累加器注册了autogradhooks。在下面图3(a)的示例中，两个垂直轴代表时间，虚线代表梯度准备就绪的时间。在过程1中，4个梯度是按顺序计算的。在过程2中，g_2是在g_3和g_4之后计算的；在图3(b)的例子中，在一次迭代中跳过了梯度g_3对应的参数，导致g_3的ready信号不存在。为了解决这个问题，DDP在forwardpass的输出张量中遍历autograd图，找到所有涉及的参数。涉及张量的就绪状态足以表明反向传播已完成。下面的算法1给出了DDP的伪代码：下图4显示了DDP在前向和反向传播过程中如何与局部模型交互：通过多次迭代。因此，研究人员需要为这个用例引入另一个接口（即不同步）。下面是一个示例代码片段：AggregateCommunicationDDP是在AggregateCommunicationLibrary的基础上构建的，包括3个选项NCCL、Gloo和MPI。DDP从这三个库中获取API，并将它们包装到同一个ProcessGroupAPI中。由于所有通信都是聚合操作，因此对所有ProcessGroup实例的后续操作必须匹配它们的类型并遵循相同的顺序。对所有库使用相同的ProcessGroupAPI允许研究人员在相同的DDP实现上试验不同的通信算法。如果单个NCCL、Gloo或MPIProcessGroup无法使链路容量饱和，DDP可以通过使用round-robinProcessGroups来实现更高的带宽利用率。具体实现DDP的实现在之前的版本中进行了多次改进。研究人员介绍了PyTorchv1.5.0的当前状态。DDP在Python和C++中都有实现，Python公开API并组成非性能关键组件，而C++提供核心梯度下降算法。PythonAPI通过Pybind11API调用C++内核。Python前端Python前端中的实现细节决定了DDP的行为。可配置旋钮在DDP构造函数API中公开。具体包括：分组处理，找出DDP中运行AllReduce的进程组实例，有助于避免与默认进程组混淆；bucket_cap_mb控制AllReduce的bucketsize，里面的application应该调整这个旋钮来优化训练速度；通过遍历autograd图，发现没有使用参数来验证DDP应该检测未使用的参数。本地模型中的ModelDeviceAffinity也可以控制DDP的行为，尤其是当模型非常大以至于需要跨多个设备运行时。对于大型模型，模型的每一层都可以放在不同的设备上，使用Tensor.to(device)API将中间输出从一个设备传输到另一个设备。DDP还可以在多个模型上运行。当层（例如BatchNorm）需要跟踪状态（例如运行方差和均值）时，模型缓冲区是必需的。DDP通过对rank为0的进程授权来支持modelbuffers。核心梯度下降开发过程中的主要工作是梯度下降，这也是DDP决定性能的关键步骤。reducer.cpp中的这个实现有4个主要组件：构建参数到桶的映射、安装autograd挂钩、启动桶AllReduce以及检测全局未使用的参数。Parameter-to-BucketMapping已经对DDP的速度产生了相当大的影响。在每次反向传播中，将张量从所有参数梯度复制到桶中，并在AllReduce后将平均梯度复制回张量。AutogradHook是DDP反向传播的入口点。在构造过程中，DDP会遍历模型中的所有参数，为每个参数找到梯度累加器，并为每个梯度累加器安装相同的post-hook函数。梯度累加器会在对应的梯度就绪时posthook，当整个bucket准备好开始AllReduce操作时DDP一定会开启。BucketAllreduce是DDP中通信开销的主要来源。默认情况下，存储桶大小为25MB。实验评估研究人员展示了PyTorchDDP在使用专用32-GPU集群和共享权限时的评估结果，其中GPU部署在4台服务器上，并通过MellanoxMT27700ConnectX-4100GB/s网卡连接。每台服务器配备8个NVIDIATeslaV100GPU。下图5显示了服务器上8个GPU的互连：接下来，研究人员使用两种流行的模型ResNet50和BERT来测量PyTorchDDP在每次迭代时的延迟和可扩展性，并且大多数实验使用随机生成的合成输入和标签，这足以比较每次迭代的延迟。延迟故障为了验证通信重叠计算的有效性，下图6分别显示了使用NCCL和Gloobackpass的ResNet50和BERT模型的延迟故障。所有实现都在4台服务器上使用了32个GPU。结果表明，反向传递是PyTorchDDP训练过程中最长的一步，因为AllReduce通信（即梯度同步）是在此过程中完成的。BucketsizeBucketsize是一个重要的配置选项。根据经验，bucket_cap_mb的默认值是25MB，用于最大努力估计。研究人员在两台机器上使用16个GPU来比较不同桶大小的每次迭代的延迟。另一个极端是在短时间内通过全梯度，结果如下图7所示。下图8显示了相同设置下在32个GPU上的实验结果。在这种情况下，异常值的跨度更大，这并不奇怪。因为参与者越多，同步的时间必然越长，扼杀者的影响也越明显。可扩展性为了解DDP的可扩展性，研究人员使用NCCL和Gloo后端在多达256个GPU上测量了ResNet50和BERT每次迭代的训练延迟。结果如下图9所示。下面的图10给出了梯度下降在第1、2、4和8次迭代时每次迭代的平均延迟。除了每次迭代延迟之外，测量收敛速度以验证加速不会因较慢的收敛而消除也很重要。实验使用MNIST数据集训练ResNet。学习率设置为0.02，batchsize为8，结果如下图11(a)所示；图11(b)是batchsize设置为256，学习率设置为0.06时的测量结果。Round-Robin进程组PyTorch分发包支持将Round-Robin进程组与多个NCCL或Gloo进程组组合，从而以Robin-Robin顺序将聚合通信分配给各个进程组实例。下面的图12显示了使用1、3和5个NCCL或Gloo进程组的循环进程组每次迭代的延迟。最显着的加速是使用NCCL后端的BERT模型。
