AllReduce

背景

深度學習領域的關鍵挑戰之一，就是耗時。耗時的原因來自於兩個方面，一方面是模型越來越大，參數越來越多，千萬甚至數十億，另一個方面是訓練數據量的增多。所以更低的耗時、更快速的實驗，對於模型的的開發效率至關重要，使研究者有機會探索更多的超差組合。

分佈式深度學習算法是降低耗時問題的重要方法。2017年6月，Facebook發佈論文，文中說他們用32臺服務器256塊GPU，將Resnet-50模型在ImageNet數據集上的訓練時間從兩週縮短到1小時。他們採用的就是分佈式學習算法。

在分佈式訓練中，數據併發是經常採用的架構，這種架構下每個節點都有完整的模型結構，共同消費訓練數據，達到加速訓練的目的：

學習的步驟爲：

1. 每個節點分別消費minibatch， 計算各自的梯度
2. reduce所有梯度，一般reduce操作採用求平均值
3. 梯度的reduce結果用來更新模型參數
4. 回到步驟1，直到數據消耗完成

這裏問題就出現了，就是reduce怎麼實現呢 ？

首先我們把這個問題抽象成一個一般性的問題.

問題定義

存在N個節點，$P_i$表示第$i$個節點，$i\in[0, N-1]$。每個節點都有一個數組$A_i$，數組長度也爲N, 即 $len（A_i） = N$。求目標數組R，滿足以下條件：

$R[k]=f(A_0[k],A_1[k],...,A_{N-1}[k])$
$len(R) = N$

其中reduce函數 $f$ 可能是平均值函數、求和函數、最大值函數、最小值函數，等等。

比如，具體對應到上面GPU數據並行訓練的場景，節點$P_i$表示每個GPU，數組$A_i$表示每個GPU各自計算的梯度，數組$R$表示梯度的聚合結果:

定義完問題之後，我們來討論下這個問題的算法設計。一個常見的算法思路，是所有GPU將各自的數據發往一箇中心節點，中心節點讓後將reduce的結果返回給GPU，如下所示:

這個算法的缺點主要有兩點：

• 中心節點需要等待所有其他節點的數據
• 中心節點的通信量比較大，並且通行量正比與節點的數量

Ring-AllReduce

下面用四個節點的例子，說明ring-allreduce的過程。Ring-AllReduce算法分爲兩個階段。

第一階段, scatter-reduce階段：

step 1:

step 2:

step 3:

step 4:

step 5:

step 6:

step 7:

第二階段，all-gather階段：

step 1:

step 2:

step 3:

step 4:

step 5:

Tree-AllReduce

Tree-AllReduce算法過程也分爲兩個階段:reduce與broadcast.