【大模型推理】all-reduce

https://andrew.gibiansky.com/blog/machine-learning/baidu-allreduce/#ref-4

1. ALL reduce , reduce, broadcast 概念

Introduction

在过去的几年中，神经网络已经被证明是解决各种问题的令人难以置信的有效工具，并且在规模和计算需求上都迅速增长。2012年，用于图像识别的SuperVision卷积网络在物体识别方面取得了巨大的进展，使用了两个GPU一周时间和6000万个参数1。2016年，研究人员在语言建模方面取得了突破性进展，该网络具有超过10亿个参数，在32个GPU上训练了三周2。在百度研究院的硅谷AI Lab内部，2014年我们的Deep Voice语音识别系统的第一次迭代大约有1100万个参数5，而一年后的下一次迭代已经增长到了1亿个参数3。

随着神经网络的参数数量和计算需求的增长，在许多节点和许多GPU上高效地并行训练神经网络变得越来越重要，因为等待几个月来训练大型网络会减慢实验速度并限制进一步发展。在这篇博文中，我们介绍了一种来自高性能计算（HPC）领域的技术，并演示了如何将其应用于深度学习，以在神经网络训练中实现显著的性能提升。

The Communication Problem

在将神经网络的训练并行化到许多GPU上时，您必须选择如何将不同的操作分布到可供您使用的不同GPU上。在这里，我们重点介绍一种称为数据并行随机梯度下降（SGD）的技术。与标准SGD中一样，梯度下降是对数据的子集（minibatches）进行的，需要多次迭代才能遍历整个数据集。然而，在数据并行训练中，每个GPU都有整个神经网络模型的完整副本，并且对于每次迭代，只分配minibatch中样本的一个子集。对于每次迭代，每个GPU在其数据上运行网络的前向传播，然后是误差反向传播，以计算相对于网络参数的损失梯度。最后，GPU相互通信，对不同GPU计算的梯度进行平均，将平均后的梯度应用于权重以获得新的权重。所有GPU都以锁步方式进行迭代，一旦一个GPU完成了它的迭代，它必须等待所有其他GPU完成它们的迭代，以便权重可以正确更新。这相当于在单个GPU上执行SGD，但我们通过将数据分布在多个GPU上并并行执行计算来获得加速比。

当您只有两个GPU并且参数以兆字节数据测量时，这些GPU如何通信可能并不重要。但是，当你的模型有数十亿个参数时，梯度可能会占用GB的空间（因为每个参数都有一个梯度值），而且你要协调几十个GPU，通信机制就变得至关重要了。

例如，考虑可能的最直接的通信机制。每个GPU在其minibatch的子集上计算梯度。然后，每个GPU将其梯度发送到单个GPU，该GPU获取所有梯度的平均值，并将平均值发送回所有其他GPU。
在这里插入图片描述

Data transfer to and from a single reducer GPU

需要发送的数据越多，发送的时间就越长；每个通信信道都有一个最大吞吐量（带宽）。例如，一个好的互联网连接可以提供每秒15兆的带宽，而千兆以太网连接可以提供每秒125兆的带宽。HPC群集中的专用网络硬件（如Infiniband）可在节点之间提供每秒数GB的带宽。

在数据从单个GPU发送和接收的直接机制中，单个GPU必须接收来自所有GPU的所有参数，并将所有参数发送到所有GPU。系统中的GPU越多，通信成本就越大。

让我们来评估一下这种沟通策略在一个真实模型上是如何工作的，例如一个以百度的深度语音23建模的语音识别网络，具有3亿个可训练参数。3亿个参数，每个参数4个字节，大约是1.2G的数据。让我们假设您的系统上的网络硬件可以支持每秒1GB的带宽；在这种情况下，如上所述将您的系统并行到两个GPU上，将使每次迭代速度减慢1.2秒。将你的训练并行化到10个GPU上，每次迭代都会减慢10.8秒；随着GPU数量的增加，每次迭代所需的时间都会线性增长。即使每次迭代需要几秒钟，这种通信成本的线性增长很快使进一步的并行化变得不切实际，并降低了训练的效率。

一种替代方案是放弃训练算法的同步性质，并消除所有GPU在梯度下降的迭代中同步前进的约束。然而，虽然这可以使您的模型更容易并行化，但消除此约束的算法（异步SGD的变体）可能难以调试，并且某些模型可能会收敛到次标准结果，因此出于本博文的目的，我们不考虑它们。

相反，我们可以从高性能计算领域使用分布式归约算法，并利用带宽最优环allreduce来解决通信问题。

这个算法分两步进行：首先是一次散列归约（scatter-reduce），然后是全聚集（allgather）。在散列归约步骤中，GPU们会交换数据，使得每个GPU最终都得到结果的一部分。在全聚集步骤中，GPU们会交换这些部分，使得所有GPU都得到完整的最终结果。

1. Scatter-Reduce

为了简化问题，假设目标是元素级别地求所有单个大浮点数数组的和；系统中有N个GPU，每个GPU都有同样大小的数组，最终每个GPU都应该有同样大小的数组，包含原数组中数字的和。

首先，GPU们将数组划分成N个较小的块（其中N是环形的GPU数量）。
在这里插入图片描述

接下来，GPU们将进行N-1次散列归约。在每次迭代中，GPU会将一个块发送给它的右邻居，并从左邻居接收一个块并累加到该块。每个GPU在每次迭代中发送和接收的块是不同的；第n个GPU从发送块n并接收块n - 1开始，然后从那里逆序进行，每次迭代发送它在前一次迭代中接收到的块。

例如，在上图中第一迭代中，五个GPU会发送和接收以下块：

GPU Send Receive
0 Chunk 0 Chunk 4
1 Chunk 1 Chunk 0
2 Chunk 2 Chunk 1
3 Chunk 3 Chunk 2
4 Chunk 4 Chunk 3

在这里插入图片描述
在完成了第一次发送和接收之后，每个GPU都会有一个块，该块由两个不同GPU上的该块值的和组成。例如，第二个GPU的第一个块将是第二个GPU和第一个GPU该块值的和。
第一次迭代

在接下来的迭代中，该过程继续进行，到最后，每个GPU都会有一个块，其中包含该块在所有GPU上的所有值的和。以下图像展示了所有数据传输和中间结果，从第一迭代开始，直到散列归约完成。

第一次迭代

第二次迭代
在这里插入图片描述

第三次迭代
在这里插入图片描述

第四次迭代
在这里插入图片描述

第五次迭代
在这里插入图片描述

2. Allgather

完成散列归约步骤后，每个GPU都有一个值的数组，其中一些值（每个GPU一个块）是最终值，它们包含来自所有GPU的贡献。为了完成allreduce，每个GPU必须交换这些块，以便所有GPU都有所有必要的值。

环形allgather与散列归约完全相同（有N-1次发送和接收的迭代），except相反的是，GPU们接收到的值，GPU们简单地覆盖块。第n个GPU从发送块n + 1并接收块n开始，然后在未来的迭代中始终发送它刚刚接收到的块。

例如，在我们的五个GPU设置中第一迭代中，GPU们会发送和接收以下块：
GPU Send Receive
0 Chunk 1 Chunk 0
1 Chunk 2 Chunk 1
2 Chunk 3 Chunk 2
3 Chunk 4 Chunk 3
4 Chunk 0 Chunk 4

在这里插入图片描述
完成了第一迭代后，每个GPU都会有两个最终数组的块。

在接下来的迭代中，该过程继续进行，到最后，每个GPU都会有整个数组的完全累积值。以下图像展示了所有数据传输和中间结果，从第一迭代开始，直到allgather完成。
Allgather data transfers (iteration 1)
在这里插入图片描述

Allgather data transfers (iteration 2)
在这里插入图片描述

Allgather data transfers (iteration 3)
在这里插入图片描述

Allgather data transfers (iteration 4)
在这里插入图片描述

Allgather data transfers (iteration 5)
在这里插入图片描述

Allreduce Communication Cost

回想一下，我们之前描述的简单通信算法，其通信成本随GPU数量成正比。allreduce的主要原因是这种情况不再成立。

在我们描述的系统中，每个中N个GPU会在散列归约和allgather中分别发送和接收值N-1次。每次发送时，GPU会发送K / N个值，其中K是不同GPU上正在求和的数组中的总值。因此，每个GPU传输和接收的总数据量为

Data Transferred=2(N−1)K/N

关键在于，这个值与GPU数量无关。

由于所有传输在离散迭代中同步发生，allreduce的速度受到环中相邻GPU之间最慢（最低带宽）的连接的限制。给定每个GPU的邻居选择正确，该算法是带宽最优的，并且是执行allreduce的最快可能算法（假设延迟成本与带宽相比可以忽略不计）4。一般来说，如果一个节点上的所有GPU在环中彼此相邻，则算法的运行效果最佳；这可以最大限度地减少网络争用的数量，否则可能会显著降低GPU-GPU连接的有效带宽。

Since all of the transfers happen synchronously in discrete iterations, the speed of the allreduce is limited by the slowest (lowest bandwidth) connection between adjacent GPUs in the ring. Given the right choice of neighbors for every GPU, this algorithm is bandwidth-optimal and is the fastest possible algorithm for doing an allreduce (assuming that latency cost is negligible compared to bandwidth)4. In general, the algorithm functions best if all GPUs on a node are next to each other in the ring; this minimizes the amount of network contention, which could otherwise significantly reduce the effective bandwidth of the GPU-GPU connections.

Allreduce在深度学习中的应用

ring allreduce是高性能计算领域中一个著名的算法，但在深度学习中的使用相对较少。在我们的实验室中，我们设法使用这个工具作为我们所有数据并行训练的基础，使我们能够有效地将训练扩展到几十个GPU。

为了最小化通信开销，我们可以利用神经网络的结构。在每次迭代中，每个GPU运行前向传播以计算错误，然后运行反向传播以计算神经网络的每个参数的梯度。反向传播从输出层开始计算梯度，并在put层中移动，这意味着输出层参数的梯度明显早于前面层的梯度。由于allreduce一次可以对网络的一个参数子集进行操作，因此我们可以在其他梯度仍在计算的时候对输出层参数启动allreduce。这样做可以将通信与反向传播步骤中的其余计算重叠，因此可以减少每个GPU最终等待通信完成的总时间。

例如，考虑一个类似于2中的语言模型，但具有大约3亿个可学习参数（因此总梯度大小为1.2 GB）。使用allreduce，每个GPU必须发送和接收大约2.4GB的数据。使用CUDA感知的MPI实现（例如OpenMPI），我们可以使用GPUDirect RDMA在GPU之间传输数据，带宽大约为每秒10 GB；但是，我们集群中的节点之间的连接速度较慢，Infiniband提供的带宽大约为每秒6 GB。由于限制因素是Infiniband连接，因此单次迭代需要大约

由于网络更深的层首先具有可用的梯度，因此我们可以在整个反向传播传递完成之前开始进行数据传输，因此真正的开销可能小于400毫秒；通信和计算之间的重叠可能会根据正在优化的神经网络的性质而有所不同。

我们实现了上述语言模型，并测试了当我们从单个GPU（没有通信开销）扩展到40个GPU时，每次迭代所花费的时间。这40个GPU被排列成5个节点，每个节点有8个GPU，通过Infiniband连接。我们以32的批处理大小运行了300次迭代的语言模型，并计算了每秒处理的样本数。
在这里插入图片描述

正如你所看到的，整个系统的吞吐量随着GPU的数量线性扩展，经过一定的操作后，添加更多的GPU并不会导致每次迭代的显著放缓。在40个GPU上运行模型每次迭代大约需要650 – 700毫秒，而在单个GPU上大约需要370毫秒。由于根据我们的估计，通信需要400毫秒，因此通过将反向传播与数据传输重叠，我们在每次迭代中额外节省了70 – 120毫秒。

Ring Allreduce是一种来自高性能计算领域的技术，它允许我们在许多设备和许多节点上高效地平均神经网络中的梯度。通过在训练期间使用这种带宽最优的算法，您可以大幅减少通信开销并扩展到更多设备，同时仍然保留同步随机梯度下降的确定性和可预测的收敛特性。该算法是网络架构和深度学习框架不可知的，可以为数据并行训练的效率提供切实和立竿见影的好处，同时也相当直截了当，易于实现。

为了让你更容易地利用这些技术，今天我们发布了baidu-allreduce，这是一个演示allreduce算法的C库，你可以将它嵌入到任何支持MPI的应用程序中。此外，Uber的优秀Horovod库实现了我们在这里首创的技术。

我们希望其他深度学习框架将在适当的情况下利用类似的技术，并且使用这些工具，您将能够轻松高效地将您的神经网络模型扩展到许多机器，而不依赖于您选择的框架。