2

Tensorflow上手5: 分布式计算中的Ring All-reduce算法

 2 years ago
source link: https://yaoyaowd.medium.com/tensorflow%E4%B8%8A%E6%89%8B5-%E5%88%86%E5%B8%83%E5%BC%8F%E8%AE%A1%E7%AE%97%E4%B8%AD%E7%9A%84ring-all-reduce%E7%AE%97%E6%B3%95-1da4f0a2e651
Go to the source link to view the article. You can view the picture content, updated content and better typesetting reading experience. If the link is broken, please click the button below to view the snapshot at that time.
neoserver,ios ssh client

Tensorflow上手5: 分布式计算中的Ring All-reduce算法

在进行大规模机器学习的时候,一个常见的问题就是如何从多个GPU上收集gradient,然后对参数进行更新.这一步骤通常被各种机器学习平台用一个All-reduce操作完成,这个操作通常能够对大数组的每一个位置进行求和,求乘机,最大值,最小值等操作.

举一个例子,下图每个worker上有一个长度为4的数组:

https://towardsdatascience.com/visual-intuition-on-ring-allreduce-for-distributed-deep-learning-d1f34b4911da

一次求和的All-reduce操作后,每个worker上的数字如下图所示:

https://towardsdatascience.com/visual-intuition-on-ring-allreduce-for-distributed-deep-learning-d1f34b4911da

All-reduce看起来很简单,但是怎么做最有效,就成了一道非常经典的面试题目.

最粗暴的方法就是每一个Worker告诉其他所有Worker自己的数据,然后进行求和.这样一共要把n个数组进行n(n-1)/2次传输,对于4个节点就是6次数据传输,对于8个节点就要28次数据传输.

https://towardsdatascience.com/visual-intuition-on-ring-allreduce-for-distributed-deep-learning-d1f34b4911da

比这种粗暴方法稍好一点的是选择一个Master节点,所有的数据都传到Master,在Master节点进行求和后传回每一个Worker.通过这个简单的优化,我们只需要进行(n-1)*2次传输就可以了.虽然对4个节点我们还是需要6次传输,但是对于8个节点我们只要进行14次传输就可以了.但是这样的传输需要主要通过Master节点进行,这使得Master节点的网络传输成为瓶颈.

https://towardsdatascience.com/visual-intuition-on-ring-allreduce-for-distributed-deep-learning-d1f34b4911da

看过我之前介绍CUDA编程的朋友也许能想到另一种理论上非常优秀的解法,分治.假设我们有8个Worker,我们可以先相隔4交换,0与4,1与5,2与6,3与7.进行求和后再将上面4组数据两两交换,0与2,1与3,4与6,5与7.最后我们将Worker两两进行分治交换,也就是0与1,2与3,4与5,6与7.这样一共进行3轮,共12轮传输后,每个节点都能得到所有节点的数据和.假设有P个Worker,每个Worker上的数据量为N,我们则进行了Log(P)次传输,每次传输的数据量为NP/2,传输次数降到最低,但是如果数据量N大,那么传输量仍有些大.

https://devblogs.nvidia.com/faster-parallel-reductions-kepler/

Ring All-reduce

今天我们最想讨论的是在实际工作中采用的算法Ring All-reduce,它能更好的平衡通讯轮数与通讯量的关系,达到并行效果非常好的All-reduce.我们假设有P个进程,编号为1至P,他们之间的网络链接如Ring一样,如下图所示.对于每一个进程,我们将数据也分为P组.

https://preferredresearch.jp/2018/07/10/technologies-behind-distributed-deep-learning-allreduce/

在第一步,进程p将第p个数据块发送到进程p+1,收到来自进程p-1的数据块p-1.比如进程1发出数据块1,收到数据块4,进程2发出数据块2,收到数据块1.如下图所示:

https://preferredresearch.jp/2018/07/10/technologies-behind-distributed-deep-learning-allreduce/

然后,进程p将自己对应的数据块与收到的数据块进行一次reduce操作,然后将其发向下一个进程.进程1发出来自己的数据块4与进程4的数据块4 reduce结果到进程2,收到来自进程4的数据块3,与进程3的数据块3的reduce结果,如下图所示:

https://preferredresearch.jp/2018/07/10/technologies-behind-distributed-deep-learning-allreduce/

这样反复三次,节点4会收到所有节点对于数据块1的reduce结果,节点3会收到所有节点对于数据块4的reduce结果,节点2会收到所有节点对于数据块3的reduce结果,节点1会收到所有节点对于数据块2的reduce结果.对于还没有理解的朋友,我们假设d_{i..j,t}表示来自进程i至进程j的数据块t的reduce结果,用>p_i表示将数据传往进程i,列出以下状态转移:

  • 第1次传输: d_{1,1}>p_2, d_{2,2}>p_3, d_{3,3}>p_4, d_{4,4}>p_1
  • 第1次计算: p1=d_{14,4}, p2=d_{12,1}, p3=d_{23,2}, p4=d_{34,3}
  • 第2次传输: d_{14,4}>p_2, d_{12,1}>p_3, d_{23,2}>p_4, d_{34,3}>p_1
  • 第2次计算: p1=d_{134,3}, p2=d_{124,4}, p3=d_{1..3,1}, p4=d_{2..4,2}
  • 第3次传输: d_{134,3}>p_2, d_{124,4}>p_3, d_{1..3,1}>p_4, d_{2..4,2}>p_1
  • 第3次计算: p1=d_{1..4,2}, p2=d_{1..4,3}, p3=d_{1..4,4}, p4=d_{1..4,1}

最后,我们只需要再进行类似上文介绍的P-1次额外传输,让每一个节点都能收到每个数据块的all reduce结果即可.

我们可以分析一下Ring All-reduce的通讯量,与之前讲过的简单算法进行对比.

假设数据量为N,有P个进程.在最开始的P-1轮传输中,每一个进程发送N/P的数据量,发送到下一个进程上.在之后的P-1轮传输中,每个进程再次将每一段reduce好的结果发送到下一个进程上,数据量为N/P.这样计算,总共的数据传输为2N(P-1)/P,这种算法的好处是总数据传输量与P基本无关.

通过计算可以看出,Ring All-reduce在传输量上比上面的几种算法都更优秀,因为他们让每一个参与All-reduce的进程都尽可能的均匀分布计算量和通讯量,消除单个操作带来的瓶颈.

现在Ring All-reduce算法已经在NCCL中实现了,并且Tensorflow的CollectiveAllReduceStrategy也采用了这样的通讯方式.如果下次面试遇到了类似问题,希望大家都能顺利答出.


Recommend

About Joyk


Aggregate valuable and interesting links.
Joyk means Joy of geeK