隨著資料集和深度學習模型的規模持續增長，訓練模型所需的時間也不斷增加，大規模分散式深度學習結合資料並行化是大幅減少訓練時間的明智選擇。然而，在大規模 GPU 叢集上的分散式深度學習存在兩大技術難題。第一大難題是大批量訓練下的收斂準確率下降；第二大難題是在 GPU 之間進行梯度同步時的資訊交流成本。我們需要一種解決這兩大難題的分散式處理新方法。
在過去的幾年裡，研究者們為這兩大問題提出了很多解決方法。他們使用 ImageNet/ResNet-50（在 ImageNet 資料集上訓練 ResNet-50 分類器）作為訓練效能的基準。ImageNet 和 ResNet-50 是最流行的用作大規模分散式深度學習基準的資料集和深度神經網路之一。表 1 對比了近期研究的訓練時間和 top-1 驗證準確率。
表 1：ImageNet/ResNet-50 訓練時間和 top-1 1-crop 驗證準確率。
從表中可以看出，隨著技術、算力的發展，在 ImageNet 上訓練 ResNet-50 的速度被不斷重新整理。日本 Perferred Network 公司的 Chainer 團隊曾在 15 分鐘訓練好 ResNet-50；2018 年 7 月，騰訊機智機器學習平臺團隊在 ImageNet 資料集上，僅用 6.6 分鐘就可以訓練好 ResNet-50，創造了 AI 訓練世界新紀錄。這一次，訓練記錄被再一次重新整理。
本文作者著眼於大批量訓練的不穩定性和梯度同步成本問題，成功地使用 2176 塊 Tesla V100 GPU 將訓練時間減少到 224 秒，並達到 75.03% 的驗證準確率。研究者還嘗試在不顯著降低準確率的前提下提升 GPU 擴充套件效率。如表 2 所示，他們最終使用 1088 塊 Tesla V100 GPU 實現了 91.62% 的 GPU 擴充套件效率。
表 2：ImageNet/ResNet-50 訓練的 GPU 擴充套件效率。
論文：ImageNet/ResNet-50 Training in 224 Seconds
論文地址：https://arxiv.org/ftp/arxiv/papers/1811/1811.05233.pdf
摘要：由於大規模批量訓練的不穩定性和梯度同步的開銷，將分散式深度學習擴充套件至 GPU 叢集級頗具挑戰。我們通過批量控制來解決大規模批量訓練不穩定的問題，用 2D-Torus all-reduce 來解決梯度同步的開銷。具體來說，2D-Torus all-reduce 將 GPU 排列在邏輯 2D 網格中，並在不同方向上執行一系列叢集計算。這兩種技術是用神經網路庫（NNL）實現的。我們在 224 秒內成功訓練了 ImageNet/ResNet-50，並且在 ABCI 叢集上沒有顯著的準確性損失。
方法
我們採用了 [4], [10] 和 [11] 中引入的批量控制技術來解決大規模批量不穩定的問題。我們開發了一個 2D-Torus all-reduce 方案來有效交換 GPU 上的梯度。
批量控制
根據之前的研究，在訓練期間通過逐漸增加批量的總規模可以減少大批量訓練的不穩定性。隨著訓練損失曲面變得「平坦」，增加批量有助於避開區域性最小值 [4] [10] [11]。在本文中，我們採用了批量控制的方法來減少批量超過 32K 時的準確率下降。訓練期間採用預定的批量變化方案。
2D-Torus All- reduce
高效的交流拓撲對於減少叢集計算的交流成本而言非常重要。人們已經提出了包括 Ring all-reduce[12] 和層級 Ring all-reduce[6] 等多種交流拓撲來提高 all-reduce 計算的效率。Ring all-reduce 演算法不能完全利用超過 1000 塊 GPU 的極大規模叢集的頻寬。這是因為 [12] 中展示的網路延遲的存在，使得演算法的交流成本和 GPU 的數量成正比。我們開發了 2D-Torus all-reduce 來解決這個問題。
2D-Torus 拓撲如圖 1 所示。叢集中的 GPU 按照 2D 網格排列。在 2D-Torus 拓撲中，all-reduce 包含三個步驟：reduce-scatter、all-reduce、all-gather。圖 2 展示了 2D-Torus all-reduce 的案例。首先，水平地執行 reduce-scatter。然後，豎直地執行 all-reduce。最後，水平地執行 all-together。
2D-Torus all-reduce 的交流成本比 Ring all-reduce 的更低。設 N 為叢集中的 GPU 數量，X 為水平方向的 GPU 數量，Y 為豎直方向的 GPU 數量。2D-Torus all-reduce 只需執行 2(X-1) 次 GPU-to-GPU 運算。相比而言，Ring all-reduce 需要執行 2(N-1) 次 GPU-to-GPU 運算。儘管層級 Ring all-reduce 執行的 GPU-to-GPU 運算次數和 2D-Torus all-reduce 相同，2D-Torus all-reduce 方案的第二步的資料量只有層級 all-reduce 的 1/X。
圖 1：2D-Torus 拓撲由水平和豎直方向上的多個閉圈構成。
圖 2： 4-GPU 叢集的 2D-Torus all-reduce 運算步驟。
評估結果
我們在 224 秒內完成了 ResNet-50 的訓練，並且準確率沒有顯著降低，如表 5 所示。訓練誤差曲線和參考曲線很相似（圖 3）。儘管最大批量可以增加到 119K 而不會有明顯的準確率降低，進一步增加最大批量到 136K 會減少 0.5% 的準確率（表 5 中的 Exp. 6）。表 6 展示了當每塊 GPU 的批量被設為 32 的時候，使用的 GPU 數量和訓練吞吐量。
儘管當我們使用超過 2176 塊 GPU 時，GPU 擴充套件效率降低了 50% 到 70%，但在使用 1088 塊 GPU 時 GPU 擴充套件效率也能超過 90%。過去的研究 [6] 報告稱當使用 1024 塊 Tesla P40 並且將批量設為 32 時，GPU 擴充套件效率達到了 87.9%。對比過去的研究，我們的 GPU 交流方案可以使用更快和更多的 GPU（Tesla V100）達到更高的 GPU 擴充套件效率。
表 5：Top-1 1-crop 驗證準確率和訓練時間
圖 3：訓練誤差曲線
表 6：2D-Torus all-reduce 的訓練吞吐量和擴充套件效率
討論
分散式深度學習的瓶頸
從過去在 ImageNet 上的大規模分散式研究來看，基本上都將工作重點放在減少通訊成本上。
深度學習的分散式訓練分為同步和非同步兩種，它們的主要區別在於引數在各個 GPU（工作器）上的計算是否獨立。
具體來說，非同步式訓練在初始化時在每個 GPU 上有一個相同的模型，然後每個 GPU 接收不同的樣本進行訓練。各個 GPU 在一個批量訓練完成時會將引數更新到一個公有的伺服器，但這個伺服器僅保留一個模型引數版本。當其它工作器訓練完一個批量時，會直接在公有伺服器上用新的模型引數覆蓋。
這種訓練方式的通訊成本較低，並且獨立工作的方式可以在時間協調上更加靈活，但也存在問題。由於「引數覆蓋」，這些 GPU 之間就像是在互相競爭，看看誰先完成訓練，就能成為下一次更新的初始引數。而更新了之後，有些 GPU 還在之前的引數上進行更新，這就導致了每個 GPU 無法獲取完整的梯度資訊，從而其更新方向不一定是朝著收斂的方向進行。也就是說，GPU 通訊的不足導致了模型更新的梯度資訊不足。
同步式訓練可以克服非同步式訓練的這個缺點。同步式訓練就是在每一次更新模型引數前，先對所有的 GPU 計算出的梯度求平均。如此就能保證每個 GPU 獲取相同的梯度資訊，但這需要等待所有的 GPU 完成訓練才行。所以很自然，這種訓練方式的一大缺點就是通訊成本大和延遲問題。經典的 all reduce 方案的通訊成本與 GPU 的數量成正比。
all reduce 通訊方案：所有 GPU 與單個 reducer GPU 之間的資料傳輸。通訊次數正比於 GPU 數量 N。
為此，百度研究院提出了 Ring all reduce 通訊方案，將 GPU 連線成一個圓環，進行引數資訊傳遞，這種通訊方式可以有效減少通訊成本。在理想情況下，這種通訊方案的成本是和 GPU 數量無關的。至於本文研究者提到的 Ring all reduce 存在的網路延遲問題，這裡就不討論了。簡單來說，Ring all reduce 將通訊分成兩個階段，在第一個階段中可以先將傳遞過來的部分引數值進行合併，然後在每個 GPU 中都有部分的最終合併引數值，並在第二個階段中再次進行傳遞。
百度將高效能運算引入深度學習：可高效實現模型的大規模擴充套件
Ring all reduce 通訊方案：被佈置在邏輯環中的 GPU，與近鄰的 GPU 接收和傳送引數值。
和 all reduce 方案對比一下，我們可以發現，Ring all reduce 就是將單階段通訊變成了多階段，每個階段可以先合併部分引數值，這樣引數數量就會在每個階段不斷縮減。也就是說，對於同類的可合併的數值通訊而言，分層通訊才是根本。如此，我們也不難理解騰訊之後提出的層級 all reduce 的思想了。
層級 all reduce 的三階段通訊方案圖示。
根據論文方法部分的解釋，2D-Torus All- reduce 的通訊拓撲比 Ring all reduce 多一個維度，並且和層級 all reduce 的通訊次數相同，但在第二個步驟的通訊資料量更小。2D-Torus All- reduce 將通訊過程分成多個維度和多個階段進行，經過了兩個階段的合併之後在每個 GPU 都有部分的最終引數值，和層級 all reduce 有著異曲同工之妙。
之後的研究可能會在什麼方向？可能的變數還是很多的，無論是 GPU 效能增長，通訊頻寬成本的降低，還是 GPU 叢集拓撲的新方案......
竟然花費那麼大成本來訓練 ImageNet？
首先，完成這種訓練的成本到底有多大？我們可以參考 Reddit 上網友的粗略計算：Tesla V100 大約是單價 6000 美元，訓練最快的結果中使用了 2176 塊 V100，總價約 1300 萬美元。你可能覺得這對於索尼大法而言不算什麼，但考慮到維護成本和佔地，這就很不經濟了。當然，索尼也可能像迪士尼一樣利用這些 GPU 來渲染動畫（遊戲）。
另外也有網友計算了一下租用谷歌雲 Tesla V100 的訓練成本，總價約：
2176GPU * $0.74/GPU·hr * 1hr/3600s * 224s ≈ $100
100 美元...Unbelievable~小夥伴們，有空跑一個？
參考內容：https://www.reddit.com/r/MachineLearning/comments/9wxcw7/r_imagenetresnet50_training_in_224_seconds/