从事分布式深度学习相关工作的同学，应该都频繁地用到了AllReduce（规约）操作。

但是对于训练框架中集成的AllReduce相关操作，其背后实现的原理是什么？

除了最近几年名声大噪的Ring AllReduce是否还有其他的AllReduce算法？

他们各自的性能开销如何？如何取舍？

本文尝试从一个较为全面的角度来展现AllReduce算法的前世今生，既分析经典算法，也介绍发展中的新秀。

MPI中的AllReduce算法

其实说到AllReduce，很多人脑海里的第一反应都是MPI_AllReduce。作为集合通信中的元老，和高性能计算领域的通信标准，在MPI_AllReduce这个通信原语背后，MPI中实现了多种AllReduce算法。

以openMPI源码为例，里面实现了多种allreduce的算法。具体的算法选择在

ompi/mca/coll/tuned/coll_tuned_decision_fixed.c

reduce+broadcast

这是最朴素最直观的一种算法（也是经典的PS架构对应的算法）。如图3所示，parameter server作为中心节点，先全局reduce接收所有其他节点的数据，经过本地计算后，再broadcast回所有其他节点。其中实线表示真实发生了的通信；而虚线表示只是示意，并没有实际通信。

对于耗时的估计，本文都按照公式α + S / B + S * C来计算。

其中α表示2个通信节点（这里节点不指机器节点，而是指集合通信中的一个node，比如1个mpi process即可认为是一个通信节点）间的latency， S(ize)表示要allreduce的数据块大小， B(andwidth)表示2个通信节点间的带宽， C(omputation)表示每字节数据的计算耗时。另外以N(umber)表示节点个数。

通信只需要2步，一步reduce，一步broadcast。它们的通信耗时都是α + S / B。（abc与d之间的通信是同时进行的，耗时只算1份）

在parameter server节点上的计算耗时是N*S*C。

总体耗时是2*(α + S/B) + N*S*C。

该算法最大的缺点就是parameter server节点的带宽会成为瓶颈。

也许有人会质疑说可以用分布式PS，PS集群的带宽就足够大了。

分两种情况，如果是数据并行，那分布式PS之间的权重协调一致将是不得不考虑的问题；如果是模型并行，那对于某一块特定的参数来说，PS：worker都是1：N的关系，当worker数量较大时，还是不可避免遇到PS的带宽瓶颈。

recursive halving and doubling

可以视为经典的树形算法，过程也非常直白。

如果节点数是2的幂，所需通信步数是2*log2N。

下图是4个节点abcd进行recursive halving and doubling的allreduce示意图。其中实线表示真实发生了的通信；而虚线表示只是示意，并没有实际通信。

相比reduce+broadcast，最大的改进是规避了单节点的带宽瓶颈。

图4 recursive算法示意图，节点数为2的幂

如果节点数不为2的幂，则会先调整至符合2的幂后，再进行上述halving and doubling的操作。总的通信步数是2*( log2K+1)

在halving和doubling的每一步，通信耗时都是α + S/B，计算耗时都是S*C。

步数约等于log2N，因此整体耗时是2*log2N*(α + S/B + S*C )。

Butterfly

Recursive算法中一个明显的不足是，在halving阶段有一半的节点没有进行send发送操作，只是“傻傻”等待接收数据。比如图4中的第一步，a->b，c->d发送数据的时候，b和d节点的发送带宽没有被利用起来。

Butterfly算法则弥补了这一点。通信的每步中，所有节点的send和recv带宽都被利用起来了。

这个算法在openMPI中对应mca/coll/base/coll_base_allreduce.c中的ompi_coll_base_allreduce_intra_recursivedoubling函数（对，别被函数名字骗了，它实现的不是上面介绍的recursive算法）。

如果节点数是2的幂，所需通信步数只要log2N。

图6 Butterfly算法示意图，节点数是2的幂(是不是有点蝴蝶的形状？)

如果节点数不为2的幂，也是先调整为符合2的幂后，再进行上述操作。总的通信步数是log2K+1。

通信步数大概是log2N， 每步的通信耗时是α + S/B，计算耗时是S*C。

整体耗时大概是log2N*(α + S/B + S*C )。

Ring AllReduce

Butterfly已经在每步中把每个节点的send/recv带宽都利用起来了，那是不是完美无缺了？答案是否定的。其潜在的问题是如果数据块过大（S过大），每次都完整send/recv一个S的数据块，并不容易把带宽跑满，且容易出现延时抖动。

Ring算法默认把每个节点的数据切分成N份。当然，这要求数据块中的元素个数count =S/sizeof(element)大于N，否则要退化为使用其他算法。

图8 Ring Allreduce算法示意图

第一阶段通过(N-1)步，让每个节点都得到1/N的完整数据块。每一步的通信耗时是α+S/(NB)，计算耗时是(S/N)*C。 这一阶段也可视为scatter-reduce。

第二阶段通过(N-1)步，让所有节点的每个1/N数据块都变得完整。每一步的通信耗时也是α+S/(NB)，没有计算。这一阶段也可视为allgather。

整体耗时大概是2*(N-1)*[α+S/(NB)] + (N-1)*[(S/N)*C]

Segmented Ring算法

主要是针对较大数据量的情况下，直接按照进程数来切分数据，可能导致单次要传输的数据依然过大，影响性能。所以改成分段执行ring。

其所需步数以及耗时和选定的切分size大小有关。切分后的过程和ring类似，此处不再赘述。

MPI中选择AllReduce算法的逻辑

我们何时该选择哪种算法呢？不妨看下openMPI里是怎么帮用户做选择的。

对照图2中的代码片段，其算法选择逻辑如下

2015年NCCL开始实现AllReduce

上述openMPI的算法最晚在2009年就都已经成熟并开源了，而英伟达在2015年下半年首次公开发布NCCL。

既然openmpi已经实现了这么多AllReduce算法，为什么英伟达还要开发NCCL？

从openMPI的源码里我们能看到，其完全没有考虑过深度学习的场景，基本没有考虑过GPU系统架构。很明显的一点，MPI中各个工作节点基本视为等同，并没有考虑节点间latency和带宽的不同，所以并不能充分发挥异构场景下的硬件性能。

而NCCL的优势就在于完全贴合英伟达自己的硬件，能充分发挥性能。但是基本的算法原理其实相比openmpi里实现的ring算法是没有变化的。

NCCL1.x只能在单机内部进行通信，NCCL2.0开始支持多节点（2017年Q2）。所以在NCCL2之前大家还会依赖MPI来进行集合通信。

2016年百度在深度学习中引入Ring AllReduce

openMPI代码中2007年就有ring算法了，为什么会有Baidu在2016年提出Ring Allreduce的说法？

其实在baidu的论文题目里就说得很清楚了，他们是“Bringing HPC Techniques to Deep Learning”，ring算法是早就有了，但是应用到深度学习领域确实是他们首创的（所以跨领域跨学科的复合型人才是多么重要^_^）。

Baidu还开源了他们基于TensorFlow修改的源码，把TF里原来进行梯度规约的地方替换成了mpi实现的ring allreduce。

具体代码在tensorflow/contrib/mpi_collectives/ring.h中

可以看到实现的是常规ring，而不是segmented ring。并且里面使用MPI_Sendrecv MPI_Irecv MPI_Send这些mpi通信原语来实现，和具体mpi库无关（无论是openmpi还是MPICH2）。也没有直接用MPI_AllReduce原语，因为按照openMPI的实现它很可能跑去用其它非ring算法了。

TensorFlow里的AllReduce

在tf早期版本中，分布式训练只有PS架构。

在2017年后，开始逐步支持多种allreduce算法，其中的ring-allreduce实现正是baidu贡献的。

NCCL2.0之后，TensorFlow/Baidu里的allreduce算法集成了NCCL来做GPU间通信，而不是依赖MPI了。

MPI和NCCL的关系

是不是从此我们只要NCCL，不再需要MPI了呢？NO

Nvidia的策略还是比较聪明，不和MPI竞争，只结合硬件做MPI没做好的通信性能优化。在多机多卡分布式训练中，MPI还是广泛用来做节点管理。当红炸子鸡Horovod也是这么做的，NCCL只做实际的规约通信。

Ring算法的问题

Ring算法是不是就完美无缺了呢？显然不是，因为虽然它可以充分利用带宽，但非常遗憾的是随着节点数的增多，单个ring越来越庞大，延迟将不可接受。

这也带来了各种改良的Ring算法。

改良Ring算法

2018年下半年机智团队提出分层Ring AllReduce

也是2D的一种形式，组内reduce->组间allreduce->组内broadcast，这种方法意在充分利用组内的高带宽的同时，弱化组间的低网络带宽带来的影响。

第一阶段reduce的通信步数是1，通信耗时是α + S / B，计算耗时n*S*C。其中n是我们分的每组内的节点数（这里是GPU卡数），当时设置的是每组16卡。

第二阶段allreduce的整体耗时套用前面的公式，是2*(m-1)*[α+S/(mB)] + (m-1)*[(S/m)*C]。其中m是我们第二层的节点数（也就是我们分的组数）。

第三阶段broadcast的通信步数是1，通信耗时是α + S / B，没有计算。

整体耗时是2*(α+S/B)+n*S*C+2*(m-1)*[α+S/(mB)] + (m-1)*[(S/m)*C]。

2018年11月索尼公司提出2D-Torus算法

该算法的论文直到2019年5月一直在修改中，其主要思想也是分层，是组内scatter-reduce->组间allreduce->组内allgather。

第一阶段scatter-reduce通信步数n-1，每一步的通信耗时是α+S/(nB)，计算耗时是(S/n)*C。 第二阶段allreduce整体耗时套用前面的公式，是2*(m-1)*[α+S/(mB)] + (m-1)*[(S/m)*C]。其中m是第二层的节点数（也就是分的组数）。

第三阶段allgather通信步数n-1，每一步通信耗时也是α+S/(nB)，没有计算。

整体耗时大概是

2*(n-1)*[α+S/(nB)] + (n-1)*[(S/n)*C] + 2*(m-1)*[α+S/(mB)] + (m-1)*[(S/m)*C]

2018年12月谷歌提出2D-Mesh算法

主要思想还是分层，并且和索尼2D-Torus很类似，都是水平和垂直两个方向，但是其步骤更简单。

第一步的水平和垂直两个方向的ring，并没有做reduce-scatter，而是最简单的转2圈来完成。通信步数分别是n-1和m-1（水平和垂直方向同时执行，m是行数，n是列数），每一步的通信耗时是α+S/B，计算耗时是S*C。总体耗时是(m+n-2)*( α+S/B+S*C)。

第二步则换方向，在垂直和水平两个方向，继续转2圈来完成。耗时同样是(m+n-2)*( α+S/B+S*C)。

整个算法耗时是2*(m+n-2)*( α+S/B+S*C)。

为什么谷歌的2D-Mesh可以比上述2D-Torus更简洁和高效？关键原因在于使用的TPU节点可以同时进行2路send和2路recv，而我们普通的服务器都是只有一张网卡，只能同时进行1路send和1路recv。

这也是谷歌论文里可以说自己的网络通信次数从传统Ring的O(N2)降低到了O(N)。

2018年7月IBM提出3D-Torus算法

IBM提出的名为BlueConnect的算法，比上述2D算法更进一步，形成了3个维度上的分解，因此可以归类为3D-Torus算法。

其主要思想是考虑了节点间不同switch的带宽不同，从而做出不同的分解。

以图15为例，里面同色的节点间以S0.0/S0.1/S0.2/S0.3这些switch相连，带宽是w0；

S0.X的switch和上一层switch S1.0/S1.1之间带宽是w1；

S1.X的switch和顶层的switch S2.0之间带宽是w2。

Step1：在节点内(同一个S0.X的switch下，带宽w0)同时执行四个reduce-scatter操作（A0-A1-A2, B0-B1-B2,C0-C1-C2, D0-D1-D2）。

Step2：在节点间(同一个S1.X的switch下，带宽w1)同时执行六个reduce-scatter操作（A0-B0, A1-B1, A2-B2,C0- D0, C1-D1, C2-D2）。 A {0,1,2}→B {0,1,2}在S1.0上运行，C {0,1,2}→D {0,1,2}在S1.1上运行，所有这些都同时发生，共享S1.x带宽。

Step3：在节点间同时执行六个reduce-scatter操作（A0-C0, A1-C1, A2-C2,B0- D0, B1-D1, B2-D2）,同时在S2.0上共享S1.x带宽。Step3后每个GPU上有S(ize)/12的数据是完成所有ReduceScatter后完整的数据，此时完成所有reduce-scatter操作。

Step4：all-gather将以完全相同但相反的顺序执行，最终完成all-reduce。

此种方法考虑到了一般机房内的3层拓扑结构（机器内->机器间交换机->上层交换机/路由器），不过其性能理论上和w0/w1/w2以及节点数N息息相关，不同取值情况下其与上述2D算法的性能优劣对比有很大不同，本文不再详细分析阐述。

2019年上半年NCCL2.4提出double binary tree算法

其主要思想是利用二叉树中大约一半节点是叶子节点的特性，通过将叶子节点变换为非叶子节点，得到两颗二叉树，每个节点在其中一颗二叉树上是叶子节点，在另一颗二叉树上是非叶子节点。

这种方法理论上是能够提供比ring算法更低的延迟（log2N < N），但实际效果需要测试。后续机智团队在资源条件合适的情况下会安排进行对比。

各种AllReduce算法对比

以上对各种经典的和新提出的AllReduce算法进行了一些介绍和分析，如有纰漏欢迎交流和批评指正。

值得一提的是，上述介绍的经典ring-allreduce、分层ringallreduce都是机智团队实际应用到业务中并取得了较好效果的算法。欢迎有兴趣的同学加入，一起探索。

参考资料

[1]：https://github.com/open-mpi/ompi/issues/4067

[2]: https://devblogs.nvidia.com/massively-scale-deep-learning-training-nccl-2-4

[3]: https://www.researchgate.net/publication/304101213_Hierarchical_redesign_of_classic_MPI_reduction_algorithms

[4]: http://andrew.gibiansky.com/blog/machine-learning/baidu-allreduce/

[5]: https://nnabla.org/paper/imagenet_in_224sec.pdf

[6]: http://www2.cs.uh.edu/~gabriel/courses/cosc6374_f07/ParCo_18_PerformanceModeling_2.pdf

[7]:Highly Scalable Deep Learning Training System with Mixed-Precision: Training ImageNet in Four Minutes

[8]:Massively Distributed SGD: ImageNet/ResNet-50 Training in a Flash

[9]:Image Classification at Supercomputer Scale

[10]:公司内部分资料