神經網絡在過去幾年中規模不斷擴大，訓練需要大量的數據和計算資源。為了提供所需的計算能力，我們可以使用高性能計算（HPC）中常見的技術將模型擴展到幾十個 GPU，但該技術在深度學習中未被充分利用。這項技術，Ring Allreduce，還能減少不同 GPU 之間的通信時間，從而允許將更多時間用在有用計算上。在百度的矽谷人工智慧實驗室（SVAIL），我們已經成功地使用這些技術訓練了當前最先進的語音識別模型。我們很高興以庫和 TensorFlow 軟體補丁的形式推出 Ring Allreduce 的實現。我們也希望通過發布這些庫可以使深度學習社區更有效地擴展他們的模型。
引言
在過去的幾年中，神經網絡已被證明是解決各種問題的非常有效的工具，並在規模和計算需求上快速增長。在用兩個 GPU 運行一周並調節了 6000 萬參數之後，用於圖像識別的 SuperVision 卷積網絡在物體識別方面取得了巨大成功 [1]。在 2016 年，對一個有超過 10 億個參數的網絡在 32 個 GPU 上訓練了 3 周之後，研究人員在語言建模方面取得了突破性進展 [2]。在 SVAIL，2014 年我們的 Deep Speech 語音識別系統的第一次迭代約有 1100 萬個參數 [5]，而一年後的下一次迭代已經增長到 1 億個參數 [3]。
隨著參數數量以及神經網絡計算需求的不斷增長，在多節點、多 GPU 上進行高效並行的神經網絡訓練已經變得越發重要，因為等待幾個月時間訓練大型網絡會減慢試驗進程，限制進一步開發。在這篇博文中，我們提出了一種來自高性能計算（HPC）領域的技術，並演示如何將其應用於深度學習以在神經網絡訓練中取得顯著的表現。
通信問題
當在多個 GPU 上並行訓練一個神經網絡，你必須選擇如何將不同的運算分布到不同的 GPU 上。本文中，我們將介紹一種被稱為數據並行隨機梯度下降（data parallel stochastic gradient descent）的技術。在標準隨機梯度下降（SGD）中，梯度下降通過使用數據的子集（minibatch）來完成，它們通過進行多次迭代來遍歷整個數據集。然而，在數據並行訓練中，每個 GPU 都有一個完整的神經網絡模型的副本，並且每一次迭代只會被分配 minibatch 樣本中的一個子集。對於每次迭代，每個 GPU 在自己處理的數據上將神經網絡向前傳播，隨後再進行誤差反向傳播（error backpropagation）來計算相對於神經網絡參數的損失的梯度。
最後，GPU 通過相互通信來平均不同 GPU 計算的梯度，將平均梯度應用於權重來獲取新權重。GPU 在鎖步（lock-step）中都進行迭代，並且一旦一個 GPU 完成了自己的迭代，它必須要等待其它所有 GPU 都完成，這樣以保證權重可以被適當地更新。這等價於在單塊 GPU 上處理 SGD，但是我們通過把數據分配給多個 GPU 來並行運算，從而獲得了計算速度的提升。
當你僅僅只有兩塊 GPU 和數以兆字節（MB）的參數時，這些 GPU 如何通信可能看上去沒什麼影響。但是，當你的模型有數十億個參數時，梯度就會佔用千兆字（GB）節的空間（因為每個參數都有一個梯度值），並且你還在同時協調幾十個 GPU，那麼此時 GPU 之間的通信機制就顯得非常重要了。
例如，我們考慮一下可能的最直接的通信機制。每個 GPU 都在 minibatch 上的一個子集裡計算一個梯度。然後，每個 GPU 都將該子集的梯度發送給同一個 GPU，讓這個 GPU 來計算所有梯度的平均值，最後它會將平均值發送回給其它 GPU。
與單個 reducer GPU 之間的數據傳輸
如果存在越多需要被發送的數據，那麼發送的時間就越長；每個通信信道都有一個最大吞吐量（帶寬）。例如，一個好的網絡連接可以提供 15MB/s 的帶寬，一個千兆乙太網連接能提供 125MB/s 的帶寬。搭載在高性能計算集群（HPC cluster）上的專業網絡硬體（比如 InfiniBand）可以在結點之間提供高達數 GB/s 的帶寬。
在數據於單個 GPU 上傳輸的直接機制（straight-forward mechanism）中，這個 GPU 必須接收來自所有其它 GPU 的所有參數，並且它還要將所有參數發回給所有 GPU。於是，系統中存在的 GPU 越多，通信成本就越大。
現在，讓我們來評估一下這種通信機制在真實模型上的能力，例如，有一個基於百度語音識別系統 Deep Speech 2 開發的語音識別網絡 [3]，它有 3 億個可訓練參數，每個參數佔 4 字節（Byte），也就是大概 1.2 GB 的數據量。讓我們假設你系統上的網絡硬體能夠支持 1GB/s 的帶寬，那也就是說，如此將你的系統在如上所述的兩塊 GPU 上並行訓練，將會讓每次迭代都變慢 1.2 秒。如果在 10 個 GPU 上並行訓練，將會讓每次迭代都變慢 10.8 秒。隨著 GPU 數量的增加，處理每次迭代的時間都會線性增長。即便每個迭代會花個幾秒鐘，這種在通信成本中的快速線性增長也會使接下來的並行處理變得不現實，它大大降低了你訓練的效率。
一種替代辦法就是放棄訓練算法的同步性質，去除所有 GPU 在鎖步中遍歷梯度下降迭代的約束。然而，儘管這可以使得你模型的並行處理更加簡便，去除了這些限制的算法（各種異步隨機梯度下降）還是會很難調試，因為有些模型會收斂到欠佳的結果上。不過由於這篇博文意不在此，我們就不在這裡考慮它了。
不過，我們可以通過使用來自高性能計算領域的分布式簡約算法（distributed reduction algorithms）並利用帶寬優化環衰減（bandwidth-optimal ring allreduce）來解決通信問題。
Ring Allreduce
上述簡單通信策略的主要問題是通信成本隨系統中的 GPU 數量線性增長。相反，ring allreduce 是這樣一種算法——其通信成本是恆定的，與系統中的 GPU 的數量無關，並且僅由系統中的 GPU 之間的最慢連接來確定。事實上，如果在通信成本上你只考慮帶寬這一因素（並忽略延遲），那麼 ring allreduce 就是一個最佳的通信算法 [4]（當你的模型較大時，這是一個很好的通信成本估算，你需要在較少的次數內發送大量數據）。
Ring Allreduce 中的 GPU 被布置在一個邏輯環路（logical ring）之中。每個 GPU 左右兩個各有一個 GPU，並且只從左邊的 GPU 接收數據，再把數據發送至右邊的 GPU。
被布置在邏輯環中的 GPU
算法的進行分兩步：第一步，scatter-reduce；第二步，allgather。在第一步中，GPU 將交換數據，使得每個 GPU 最終都有一個最終結果的數據塊。在第二步中，GPU 將交換那些塊，使得所有 GPU 最終得到完整的最後結果。
Scatter-Reduce
為了簡單起見，讓我們假設目標是以元素方式求和浮點數的單個大數組的所有元素。在系統中有 N 個 GPU, 其中每個 GPU 有一個相同大小的數組。在 allreduce 的最後，每個 GPU 都應該有一個同樣大小的包含了原數組中數值的和的數組。
一開始，GPU 把數組分割成 N 個較小的塊（其中 N 是 GPU 在環中的數量）。
接著，GPU 會執行 N-1 次迭代 scatter-reduce。在每一次迭代中，GPU 將發送其中一個塊到右邊的 GPU，並從左邊的 GPU 接收一個塊，把數據累積進該塊。在每一次迭代中，被發送的塊和被接收的塊是不同的。第 n 個 GPU 以發送塊 n 和接收塊 n – 1 開始，並從那兒接著向後運行。每次迭代發送的塊即是上次迭代所接收的塊。
例如，在第一次迭代中，上圖表中的 5 個 GPU 將會發送和接收以下的塊：
GPU 發送 接收
0 Chunk 0 Chunk 4
1 Chunk 1 Chunk 0
2 Chunk 2 Chunk 1
3 Chunk 3 Chunk 2
4 Chunk 4 Chunk 3
在 scatter-reduce 的第一次迭代中的數據傳輸
在第一次的發送和接收完成之後，每個 GPU 會有一個由兩個不同 GPU 中的相同塊的總和組成的塊。例如，第二個 GPU 上的第一塊將是來自第二個 GPU 和第一個 GPU 的那個塊中的值的和。
scatter-reduce 的第一次迭代完成之後的中間和
在下一次迭代中，進程繼續，直到最後，每個 GPU 會有一個塊包含所有 GPU 中的那塊的所有值的和。下面的圖像演示了所有的數據傳輸和中間結果，從第一次迭代開始，一直持續到 scatter-reduce 結束。
scatter-reduce 數據傳輸（迭代 1）
scatter-reduce 數據傳輸（迭代 2）
scatter-reduce 數據傳輸（迭代 3）
scatter-reduce 數據傳輸（迭代 4）
所有 scatter-reduce 傳輸結束之後的最後狀態
Allgather
scatter-reduce 這一步完成之後，每個 GPU 有一個值的數組，其中這些值（每個 GPU 一個塊）中的一些是包含來自所有 GPU 貢獻的最後值。為了完成 allreduce，GPU 必須交換這些塊，從而所有的 GPU 獲得所有的必需值。
該環路的 allgather 的執行等同於 scatter-reduce（通過 N-1 次發送和接收的迭代），除了不是累加 GPU 接收的值，而是簡單地重寫塊。第 n 個 GPU 通過發送第 n+1 個塊和接收第 n 個塊開始，並在未來的迭代中一直發送它剛接收的塊。
例如，在我們的 5 個 GPU 設置中的第一次迭代，GPU 會發送和接收以下的塊:
GPU 發送 接收
0 Chunk 1 Chunk 0
1 Chunk 2 Chunk 1
2 Chunk 3 Chunk 2
3 Chunk 4 Chunk 3
4 Chunk 0 Chunk 4
allgather 的第一次迭代中的數據傳輸
在首次迭代完成之後，每個 GPU 將有最終數組（final array) 的兩個塊。
在接下來的迭代中，該進程會繼續運行，一直到最後每個 GPU 都會有整個數組的全部累計值。下面的圖像演示了從第一次迭代到 allgather 完成的所有數據傳輸和中間結果。
Allgather 數據傳輸（第 1 次迭代）
Allgather 數據傳輸（第 2 次迭代）
Allgather 數據傳輸（第 3 次迭代）
Allgather 數據傳輸（第 4 次迭代）
所有 allgather 完成後的最終狀態
Allreduce 通信成本
回到引言中描述的簡單的通信算法，通信成本（communication cost）會隨 GPU 的數量而線性增長。allreduce 效果良好的主要原因是情況已經發生了變化。
在我們描述的系統中，每 N 個 GPU 都會因為 scatter-reduce 而接收 N-1 次值，還為 allgather 接收 N-1 次值。每一次，GPU 都會發送 K/N 個值，其中 K 是指數組中值的總數量，這是在不同 GPU 上相加得到的。因此，每個 GPU 的傳入和傳出數據總量為：
其獨立於 GPU 的數量，這是很關鍵的。
因為在離散的迭代中，所有的傳輸都是同時發生的，所以 allreduce 的速度受限於該環路中相鄰 GPU 之間的最慢（最低帶寬）的連接。為每個 GPU 選擇合適的鄰居，該算法能在所有 allreduce 上做到帶寬最優並且可能是最快的算法（假設相比於帶寬，其延遲成本可以忽略不計）[4]。總的來說，如果在一個節點上的所有 GPU 都臨近該環路中的其它 GPU，那麼該算法可以得到最好的效果；這能最小化網絡連接的數量，從而顯著增加這種 GPU-GPU 連接的有效帶寬。
將該 Allreduce 應用於深度學習
ring allreduce 是高性能計算領域內一個眾所周知的算法，但在深度學習領域內的應用相對較少。在我們的實驗室中，我們已經成功地將這種工具用作我們所有的數據並行訓練的基礎，讓我們可以將訓練有效地擴展到十餘個 GPU。
為了最小化通信負載，我們可以利用神經網絡的結構。在每一次迭代中，每一個 GPU 都會運行前向傳播來計算錯誤，然後再運行反向傳播來為該神經網絡的每一個參數計算梯度。反向傳播是從輸出層開始計算梯度，然後向輸入層移動，這意味著輸出層參數的梯度在更早的層的梯度之前是顯著可用的。因為該 allreduce 能一次操作該網絡的參數的一個子集，所以我們可以在其輸出層參數上開始 allreduce，同時還能計算其它梯度。這樣做使得該計算可以和反向傳播步驟中的其它計算一起進行，所以可以減少每個 GPU 用於等待通信完成的總時間。
比如，假設有一個類似於 [2] 中的語言模型，但帶有大約 3 億個可學習的參數（因此總梯度大小為 1.2 GB）。使用 allreduce，每個 GPU 必須發送和接收大約 2.4 GB 的數據。使用一種 CUDA-aware MPI 實現（比如 OpenMPI），我們可以使用 GPUDirect RDMA 在 GPU 之間傳輸數據，帶寬大約為 10 GB/s；但是，在我們的集群中的節點之間的連接更慢——Infiniband 提供了大約 6 GB/s 的帶寬。因為 Infiniband 連接是這裡的限制因素，那麼單次迭代就需要大約
因為更深入到網絡中的層一開始就有可用的梯度，所以我們可以在整個反向傳播通過完成之前就開始進行數據傳輸，所以其真正的開銷可能會少於 400 ms；通信和計算之間的這種重疊可能會隨被優化的神經網絡的本質而發生改變。
我們實現了之前提到的語言模型，並且在我們從單個 GPU（沒有通信開銷）擴展到 40 個 GPU 時測試了每次迭代所用的時間。這 40 個 GPU 被布置成了 5 個節點，每個節點 8 個 GPU，它們之間通過 Infiniband 連接。我們以 32 的批大小對該語言模型運行了 300 迭代，然後計算了其每秒所處理的樣本的數量。
對於一個有 3 億個參數的語言模型，每秒所處理的樣本的數量會隨同時進行同步訓練的 GPU 的數量而線性增長。
如你所見，整個系統的吞吐量會隨 GPU 的數量線性擴展；但超過一個特定的點之後，增加更多 GPU 並不會導致每次迭代的顯著減速。在 40 個 GPU 上運行該模型的時間是每次迭代 650-700 ms，而在單個 GPU 上該數字為大約 370 ms。因為據我們估計通信大約需要 400 ms，所以我們通過將反向傳播和數據傳輸重疊起來進行而為每次迭代節省了 70-120 ms 的時間。
總結
來自於高性能計算領域的 Ring Allreduce 技術讓我們可以高效地在跨多設備和多節點的神經網絡上對梯度進行平均。通過在訓練過程中使用這種帶寬優化算法，你可以極大地減少通信負載並擴展到遠遠更多的設備，同時仍能保留同步隨機梯度下降的確定性和可預測的收斂性。該算法並不特定於任何網絡架構和深度學習框架，能夠為數據並行訓練的效率提供顯著的和直接的好處，同時其部署實現也是相當直接和容易的。
為了讓你能更輕鬆地使用這些技術，今天我們也發布了一個演示該 allreduce 算法的 C 語言庫：，你可以將其嵌入到任何使用 MPI 的應用中。此外，我們也已經將該 allreduce 整合到 TensorFlow 中（可在 tensorflow.contrib.mpi 模塊獲取文檔）。
我們希望其它深度學習框架也能在合適的地方使用類似的技術；通過使用這些工具，你可以輕鬆有效地將你的神經網絡模型擴展到許多機器，而且不論你選擇的是什麼框架。
參考文獻
1.Krizhevsky, Alex, Ilya Sutskever, and Geoffrey E. Hinton.「ImageNet classification with deep convolutional neural networks.」Advances in neural information processing systems. 2012.
2.Jozefowicz, Rafal, et al.「Exploring the limits of language modeling.」arXiv preprint arXiv:1602.02410 (2016).
3.Amodei, Dario, et al.「Deep speech 2: End-to-end speech recognition in english and mandarin.」arXiv preprint arXiv:1512.02595 (2015).
4.Patarasuk, Pitch, and Xin Yuan.「Bandwidth optimal all-reduce algorithms for clusters of workstations.」Journal of Parallel and Distributed Computing 69.2 (2009): 117-124.
5.Hannun, Awni, et al.「Deep speech: Scaling up end-to-end speech recognition.」arXiv preprint arXiv:1412.5567 (2014).
February 21st, 2017