今天無聊看到了ring-allreduce的gpu通信算法，本來希望在相關網頁上看看相關介紹的，但是在baidu research上卻搜不到相關的東西，後來看了看baidu-allreduce代碼的註釋，才明白。這是一個說明起來挺簡單的算法，自己打算隨便說說。
覺得英文好的可以直接看看GitHub上註釋，寫的很清晰：https://github.com/baidu-research/baidu-allreduce/blob/master/collectives.cu#L156
如無特殊說明，本博客的圖片都來自於知乎上的一個回答，主要原因是baidu research上找不到這篇文章，所以相關的圖例只在這個知乎回答裏找到了
一般的多卡gpu訓練有一個很大的缺陷，就是因爲每次都需要一個gpu從其他gpu上收集訓練的梯度，然後將新的模型分發到其他gpu上。如下圖：
這樣的模型最大的缺陷是gpu 0的通信時間是隨着gpu卡數的增長而線性增長的。所以就有了ring-allreduce，如下圖：
該算法的基本思想是取消Reducer，讓數據在gpu形成的環內流動，整個ring-allreduce的過程分爲兩大步，第一步是scatter-reduce，第二步是allgather。
先說第一步：首先我們有n塊gpu，那麼我們把每個gpu上的數據（均等的）劃分成n塊，並給每個gpu指定它的左右鄰居（圖中0號gpu的左鄰居是4號，右鄰居是1號，1號gpu的左鄰居是0號，右鄰居是2號……），然後開始執行n-1次操作，在第i次操作時，gpu j會將自己的第(j - i)%n塊數據發送給gpu j+1，並接受gpu j-1的(j - i - 1)%n塊數據。並將接受來的數據進行reduce操作，示意圖如下：
當n-1次操作完成後，ring-allreduce的第一大步scatter-reduce就已經完成了，此時，第i塊gpu的第(i + 1) % n塊數據已經收集到了所有n塊gpu的第(i + 1) % n塊數據，那麼，再進行一次allgather就可以完成算法了。
第二步allgather做的事情很簡單，就是通過n-1次傳遞，把第i塊gpu的第(i + 1) % n塊數據傳遞給其他gpu，同樣也是在i次傳遞時，gpu j把自己的第(j - i - 1)%n塊數據發送給右鄰居，接受左鄰居的第(j - i - 2) % n數據，但是接受來的數據不需要像第一步那樣做reduce，而是直接用接受來的數據代替自己的數據就好了。
最後每個gpu的數據就變成了這樣：
如果覺得不懂的話，我們舉一個3gpu的例子：
首先是第一步，scatter-reduce：
然後是allgather的例子：
reference:
https://github.com/baidu-research/baidu-allreduce
https://www.zhihu.com/question/57799212/answer/292494636?utm_source=ZHShareTargetIDMore&utm_medium=social&utm_oi=37729630945280