2048塊GPU再次加速SGD：富士通74.7秒完成ResNet-50 訓練

選自arXiv

作者：Masafumi Yamazaki 等

機器之心編譯

參與：杜偉、Chita

隨著訓練數據越來越大，我們需要一種高效的訓練方法，而今年提出了很多大批量、大學習率的大規模分散式訓練方法。例如 2018 年 8 月，fast.ai 學員在18 分鐘內成功完成訓練 ImageNet的任務，準確率達到 93%（128 塊 V100 GPU）；騰訊在 ImageNet 數據集上以 64K 的批量大小用4 分鐘訓練好 AlexNet、6.6 分鐘訓練好 ResNet-50（1024-2048 塊 Tesla P40）；索尼隨後以 34K 到 64K 的批量大小用224 秒在 ImageNet 訓練 ResNet-50（2176 塊 V100 GPU）。

這一次，來自富士通的研究人員用上了 2048 塊 GPU，以 81,920 的批量大小用 74.7 秒完成了 ResNet-50 訓練。

在大規模數據集上訓練的深度神經網路（DNN）在目標檢測、語言翻譯等多個領域取得了令人印象深刻的結果。但是，隨著 DNN 模型和數據集規模不斷增加，DNN 訓練的計算成本也隨之增長。眾所周知，數據並行的分散式深度學習在集群上執行訓練，是一種有效的加速方法。在這種方法中，集群上啟動的所有進程具有相同的 DNN 模型和權重。每一個進程用不同的小批量（mini-batch）來訓練模型，並結合來自所有進程的權重梯度以更新所有權重。這種通信開銷成為大型集群的一大問題。為了減少大型集群的開銷，來自富士通的研究人員增加 DNN 的批量大小，並以並行方式執行 DNN 訓練。但是，large mini-batch 的訓練通常會導致 DNN 模型的驗證準確率下降。所以，研究人員使用了幾項技術來增加批量大小（每次迭代中計算的輸入圖像數量），同時不影響驗證準確率。

來自富士通的研究人員使用 AI Bridging Cloud Infrastructure（ABCI）集群（具有 2048 個 GPU）和自優化的 MXNet 深度學習框架來執行實驗，結果 74.7 秒內在 ImageNet 上使 ResNet-50 的驗證準確率達到 75.08%，使用的批量大小為 81920。

一般而言，大型集群上分散式深度學習的批量大小應該很大。Goyal 等人 [2] 提議使用預熱（warm-up）技術來保持使用批量大小為 8192 時的驗證準確率。谷歌 [3] 和索尼 [7] 採用可變且增加了的批量大小實現了高度並行化處理。

所以，每一層的權重梯度範數和權重範數之間的差異導致了訓練的不穩定，[10] 使用層級自適應速率縮放（Layer-wise Adaptive Rate Scaling，LARS）規範了每一層的差異，使 DNN 可以在不影響驗證準確率的情況下以 32768 的批量大小進行訓練。

方法

準確率提升

研究者採用了深度學習優化器中常用的隨機梯度下降法（SGD）。當使用 large mini-batch 訓練時，SGD 更新數量會隨著小批量的增加而減少。所以，提高 large mini-batch 上最終驗證準確率是一個很大的挑戰，研究者採取了以下方法：

1）控制學習率：由於更新數量較小，研究人員需要採用較高的學習率來加速訓練。但是，較高的學習率會導致模型在早期階段訓練不穩定。因此，研究人員採用能夠逐步提升學習率的預熱技術來穩定 SGD。此外，所有層都用相同的學習率對於某些層來說太高了，所以研究人員還使用層級自適應速率縮放（LARS）[10] 來穩定訓練，該方法根據權重範數和梯度範數調整每一層的學習率。

對於權重收斂，研究人員嘗試了步長、多項式以及線性等多種學習率衰減模式。此外，研究人員在多次試驗的基礎上採用了優化的權重衰減模式。

2）其他技術：先前的研究 [7] 使用標籤平滑（label smoothing）[11] 在小批量大小為 32,768 的情況下實現了準確率提升。因此本論文作者也採用了這種方法，並證實該方法可在批量大小為 81,920 的情況下提高準確率。

批歸一化層均值和方差的移動平均值（moving average）在每一進程中單獨進行計算，其中權重也是同步的。這些值在 large mini-batch 訓練中變得不準確，所以研究人員微調一些超參數以優化移動平均值。

表 1：ImageNet 上 ResNet-50 的訓練時間和 top-1 驗證準確率。

框架優化

研究人員使用以 C 和 CUDA C 語言寫就的開源深度學習框架MXNet，該框架提供多種語言綁定。MXNet 具有靈活性和可擴展性，能夠高效地在集群上訓練模型。但是，在中小規模集群環境中僅佔總時間一小部分的部分進程可能成為大型集群環境的瓶頸。研究人員使用多種分析器來分析 CPU 和 GPU 性能，並找出瓶頸問題。研究人員用以下方法優化了瓶頸，以提高訓練吞吐量：

1）並行 DNN 模型初始化：在數據並行分散式深度學習中，所有層必須初始化，以使所有進程中的權重具有相同的值。一般而言，根進程（root process）初始化模型的所有權重。之後，根進程將這些權重廣播到所有進程。廣播時間隨進程數量的增加而延長，當分散式深度學習具有數千個進程時，廣播操作成本不可忽視。因此，研究人員採用其它初始化方法：每一進程擁有相同的種子，然後以並行化方式初始化權重。該方法可以在沒有廣播操作的情況下同步初始權重。

2）GPU 上的批量範數計算：每層的範數計算都是必要的，以使用 LARS 更新權重。與 GPU 上的內核數量相比，ResNet-50 網路的大部分層沒有足夠數量的權重。如果我們計算 GPU 上每層的權重範數，就會發現線程的數量不足以佔據所有 CUDA 內核。因此，研究人員實現了一個特殊的 GPU 內核，以在 MXNet 中批量進行範數計算。該 GPU 內核可以啟動足夠數量的線程，並且可以並行計算所有層的範數。

通信優化

分散式並行深度學習需要 allreduce 通信，以在所有進程之間交換每一層的梯度。在大型集群環境中，allreduce 通信開銷是不可忽略的，因為由於每個 GPU 的批量大小變小，通信時間變得更長，而計算時間變得更短。為了解決這些問題，研究人員採用了以下兩種優化方法。

1）調節通信的數據大小：深度學習模型由很多層組成，而梯度的數據大小因層而異。每層的 allreduce 操作會由於頻繁調用通信而導致大量開銷，而如果梯度的數據大小變小了，情況會變得更糟，因為網路帶寬無法被高效使用。因此，擴大 allreduce 的數據大小非常重要。研究人員收集了各層的梯度，然後將 allreduce 的數據大小調整為幾兆位元組。

2）做最佳的通信調度：研究人員開始對部分層執行 allreduce 操作，而沒有等所有層完成。這導致 allreduce 操作與反向處理重疊。開始 allreduce 操作的時機應是梯度的數據大小大於閾值時。每個進程需要與其它進程保持同步，因為每個進程只能發送相同層的梯度。使用 allgather 操作通常可以找到完成的層，但是這會導致額外的開銷。為了去掉這項開銷，研究人員預先以統計的方式將所有層分為好幾組。一旦每個進程完成一組中所有層的反向處理，就計劃開始 allreduce 操作。

環境和實驗結果

研究人員使用 ABCI 集群來評估其基於 MXNet 的優化框架的性能。ABCI 集群的每個節點由兩個 Xeon Gold 6148 CPU、四個英偉達 Tesla V100 SXM2 GPU 組成。另外，一個節點上的 GPU 通過 NVLink 連接，節點還具有兩個 InfiniBand 網卡。

圖 1：ABCI 集群中的計算節點圖解。它由兩個 CPU、四個 GPU 以及兩個連接相應 CPU 的 HCA 組成。

圖 3：在 49,152 或更大規模的批量大小上訓練的 top-1 驗證準確率。

論文：Yet Another Accelerated SGD: ResNet-50 Training on ImageNet in 74.7 seconds

論文鏈接：https://arxiv.org/abs/1903.12650

摘要：人們對於能夠快速執行機器學習的演算法有有著強烈的需求，而深度學習的速度僅在過去兩年就加快了 30 倍。使用 large mini-batch 的分散式深度學習是解決這個需求的關鍵技術，也是一項巨大的挑戰，因為在不影響準確率的情況下，很難在大型集群上達到很高的可擴展性。在本文中，研究人員引入優化方法來應對這個挑戰。他們應用這些方法在 ABCI 集群上使用 2,048 個 GPU 在 74.7 秒時間內完成訓練。訓練吞吐量為 173 萬張圖像/秒，top-1 驗證準確率為 75.08%。

本文為機器之心編譯，轉載請聯繫本公眾號獲得授權。

------------------------------------------------

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！