从事分布式深度学习相关工作的同学,应该都频繁地用到了AllReduce(规约)操作。
![9c52c94c56a8d619a7aea191d19fc2d0.png](https://i-blog.csdnimg.cn/blog_migrate/c4e0a1b28a56593f087093c824ce9835.jpeg)
但是对于训练框架中集成的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
![335f5546009611a24d44ae958f92a4e5.png](https://i-blog.csdnimg.cn/blog_migrate/2171d614c850f08bcc6956061634c1e1.jpeg)
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)表示节点个数。
![0e600b53570563f0839d1b072aea4639.png](https://i-blog.csdnimg.cn/blog_migrate/fc842e27ccdfbada754508880aa79247.jpeg)
通信只需要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,最大的改进是规避了单节点的带宽瓶颈。