機器之心編輯部，參與：劉曉坤、王淑婷、張倩。
隨著技術、算力的發展，在 ImageNet 上訓練 ResNet-50 的速度被不斷刷新。2018 年 7 月，騰訊機智機器學習平臺團隊在 ImageNet 數據集上僅用 6.6 分鐘就訓練好 ResNet-50，創造了 AI 訓練世界紀錄。如今，這一紀錄再次被索尼刷新……隨著數據集和深度學習模型的規模持續增長，訓練模型所需的時間也不斷增加，大規模分佈式深度學習結合數據並行化是大幅減少訓練時間的明智選擇。然而，在大規模 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 之間的數據傳輸。
為了降低 all reduce 通信方案的成本，百度研究院提出了 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/
閱讀更多 機器之心 的文章
