Ring All-reduce: 分布式深度学习的巧妙同步

Ring All-reduce在高性能计算中经常被用到，百度首先将其引入到深度学习领域。

在 分布式深度学习之参数服务器架构 ,我们给大家介绍了以异步为主的参数服务器架构，参数服务器架构本身是支持同步训练的，但是没有专门为同步训练优化过。

Ring All-reduce弥补了这一空白，对同步式的分布式训练提供了通信优化。从 大批量SGD: 1小时训练ImageNet 一文开始，同步大批量训练越来越成为主流，因而Ring All-reduce方法在各个深度学习框架中都被支持了，百度在引入这个方法到深度学习中时，就为tensorflow提供了实现。

Naive同步

在看Ring All-reduce之前，先来看一下Naive版的同步。如下图，专门有一个GPU负责梯度的合并。然后每个worker都把计算好的梯度传递给负责聚合的GPU。这种方式就是典型的中心式，当模型很大，GPU很多的时候，中心节点的带宽就不够用了。

Ring All-reduce

而Ring All-reduce则是去中心化的，如下图,每个节点都处在一个逻辑环上，接受上家的数据，往下家传数据。

那么，那么这种方式是如何做同步的呢？这里，分为Scatter-reduce和Allgather两个阶段。我们一个一个的看。

Scatter-reduce

在Scatter阶段，把每个GPU上的要同步的模型梯度进行切分，切分成N份，N为系统中GPU的个数。如下图，5个GPU，所以每个GPU上的模型梯度被分为了5份。

注意，这里的前提是每个GPU上都是模型的一个全拷贝。即GPU上的各个梯度所代表的的含义都是相同的。对于一个GPU放不下的模型，需要采用 逻辑GPU 的概念，认为装载一个模型的几个GPU为逻辑上的一个GPU。

Scatter-reduce的通信分为N-1步，在每一步中，每个GPU都将其中的一块数据发给下家，并接受上家的数据，接受到的数据需要和本地数据做聚合。而至于每次发送哪一块数据给下家，我们先来看一个例子:

在第一步中，GPU0将块a传递给GPU1，同时接受GPU4的块e，因为GPU4在逻辑上是GPU0的上家。在接受了e 4 后，需要把e 0 更新为e 0 + e 4 。其他GPU以此类推。

在第二步中，每个GPU将上一步中刚刚更新过的参数发送出去。

第三步，第四步类似。

最后得到的结果如下，每个GPU都拥有了1/N的数据块，这1/N的数据块，是在全部GPU上这一部分数据的和。

这个操作是一个典型的并行通信操作，对于每一个块来说，在N-1步中，都会被遍历一遍。

Allgather

在得到的上面的结果后，就需要把每个GPU上所拥有的1/N结果广播到所有GPU上，从而实现每个GPU上都是全部数据的和。

Allgather和Scatter-reduce的操作很相似，只不过在GPU在收到数据后，需要override本地数据而非让本地数据去和传来的数据做加和。

数据传递流程如下，就是让每个GPU上的那1/N的计算好的和在所有GPU上遍历一遍。

第一步，GPU0接受GPU4的a结果，并将自己的b结果传出去。其他GPU也类似。

同理，第二步，GPU0接受GPU4的e结果，并将自己的a结果传出去。

以此类推

最后得到的结果如下:

优势在哪里？

通过上面的运行流程，Ring All-reduce有如下优势:

• 去中心化，每次的通信都发生在逻辑上相邻的GPU之间。

• 小数据传送，每一步和单个GPU相关的数据传送量是 2 * all / N，对于一个GPU来说，在两个阶段传送的数据总量是 2(N-1) * all / N，可以看到，每个GPU上传送的数据总量和集群GPU的个数是独立的。

在将Ring All-reduce集成到深度学习系统中时，除了它本身的优化，还可以和梯度的计算做流水化并行。因为梯度的计算是从神经网络的高层开始到底层的，因而，在高层的梯度计算完成之后，这部分数据就可以开始Ring All-reduce去同步了。

有了Ring All-reduce之后，分布式系统的吞吐量与GPU的关系如下：

思考

勤思考，多提问是Engineer的良好品德。

• Ring All-reduce方法现在扩展到40个GPU上能有线性的提速，如果GPU更多，比如1024，会遇到什么问题？

参考文献

• [1]. https://andrew.gibiansky.com/blog/machine-learning/baidu-allreduce/

欢迎加入AINLP技术交流群