獨立研究者 Karanbir Chahal 和 Manraj Singh Grover 與 IBM 的研究者 Kuntal Dey 近日發布了一篇論文，對深度神經網絡的分布式訓練方法進行了全面系統的總結，其中涉及到訓練算法、優化技巧和節點之間的通信方法等。機器之心摘取了論文主干內容進行介紹，更多有關數學推理過程和算法步驟的解讀請參閱原論文。

[[248730]]

論文地址：https://arxiv.org/abs/1810.11787

深度學習已經為人工智能領域帶來了巨大的發展進步。但是，必須說明訓練深度學習模型需要顯著大量的計算。在一臺具有一個現代 GPU 的單臺機器上完成一次基于 ImageNet 等基準數據集的訓練可能要耗費多達一周的時間，研究者已經觀察到在多臺機器上的分布式訓練能極大減少訓練時間。近期的研究已經通過使用 2048 個 GPU 的集群將 ImageNet 訓練時間降低至了 4 分鐘。這篇論文總結了各種用于分布式訓練的算法和技術，并給出了用于現代分布式訓練框架的當前最佳方法。更具體而言，我們探索了分布式隨機梯度下降的同步和異步變體、各種 All Reduce 梯度聚合策略以及用于在集群上實現更高吞吐量和更低延遲的最佳實踐，比如混合精度訓練、大批量訓練和梯度壓縮。

1. 引言

(1) 背景和動機

數據正以前所未有的規模生成。大規模互聯網公司每天都會生成數以 TB 計的數據，這些數據都需要得到有效的分析以提取出有意義的見解 [1]。新興的深度學習是一種用于執行這種分析的強大工具，這些算法在視覺 [2]、語言 [3] 和智能推理 [4] 領域的復雜任務上創造著當前最佳。不幸的是，這些算法需要大量數據才能有效地完成訓練，而這又會耗費大量時間。第一個在 ImageNet 分類任務上取得當前最佳結果的深度學習算法在單個 GPU 上訓練足足耗費了一周時間。在如今的時代，這樣的速度已經完全不行了，因為現在的模型訓練所使用的數據甚至會讓 ImageNet 數據集的規模都相形見絀。現在存在以橫向方式延展深度學習訓練的內在需求，同時還要保證能夠維持單 GPU 模型那樣的準確度。理想情況下，這種訓練的速度應該隨機器數量的增加而線性增大，同時還要能容錯以及能在高延遲網絡條件下收斂。

(2) 分布式訓練概述

神經網絡的分布式訓練可以通過兩種方式實現：數據并行化和模型并行化。數據并行化的目標是將數據集均等地分配到系統的各個節點(node)，其中每個節點都有該神經網絡的一個副本及其本地的權重。每個節點都會處理該數據集的一個不同子集并更新其本地權重集。這些本地權重會在整個集群中共享，從而通過一個累積算法計算出一個新的全局權重集。這些全局權重又會被分配至所有節點，然后節點會在此基礎上處理下一批數據。

模型并行化則是通過將該模型的架構切分到不同的節點上來實現訓練的分布化。AlexNet [2] 是使用模型并行化的最早期模型之一，其方法是將網絡分攤到 2 個 GPU 上以便模型能放入內存中。當模型架構過大以至于無法放入單臺機器且該模型的某些部件可以并行化時，才能應用模型并行化。模型并行化可用在某些模型中，比如目標檢測網絡 [5]，這種模型的繪制邊界框和類別預測部分是相互獨立的。一般而言，大多數網絡只可以分配到 2 個 GPU 上，這限制了可實現的可擴展數量，因此本論文主要關注的是數據并行化。

本論文大致分為六個章節，第一節將介紹已有的優化訓練算法。第二節將關注用于連接網絡各節點的通信策略。第三節會探索一些具體技術，比如大批量訓練、梯度壓縮以及用于在低功耗設備和低速網絡條件下實現有效訓練的混合精度訓練。第四節則會消化之前章節的信息，并選擇出用于不同設定的最優的訓練算法和通信原語。最后兩節分別是未來研究方向和總結。

2. 分布式訓練框架的組件

(1) 分布式訓練算法

一個常用于分布式設定中的訓練的常用算法是隨機梯度下降(SGD)，該算法將是我們進一步討論中的核心點。需要指出一個重點，針對 SGD 提及的原則可以輕松地移植給其它常用的優化算法，比如 Adam [6]、RMSProp [7] 等等 [8]。分布式 SGD 可以大致分成兩類變體：異步 SGD 和同步 SGD。

同步 SGD [9] 的目標是像分布式設置一樣對算法進行復制，由此將網絡中的節點緊密地耦合到一起。而異步 SGD [10] 則會通過降低節點之間的依賴程度而去除不同節點之間的耦合。盡管這種解耦能實現更大的并行化，但卻不幸具有副作用，即穩定性和準確性要稍微差一些。人們已經提出了一些對異步 SGD 的修改方式，以期能縮小其與同步 SGD 的準確度差距。最近的研究趨勢是擴展同步 SGD，更具體而言，即用大批量來訓練網絡已經得到了出色的結果。

較大的 mini-batch 大小具有一些優勢，主要的一個優勢是：在大 mini-batch 上訓練時，模型能以更大的步幅到達局部最小值，由此能夠加快優化過程的速度。但是在實踐中，使用大批量會導致發散問題或「泛化差距」，即網絡的測試準確度有時會低于在更小批量上訓練的模型。最近的一些研究通過與批量大小成比例地調整學習率而實現了在大批量上的訓練。實驗發現，增加批量大小就相當于降低學習率 [11]，而使用大批量進行訓練還有一個額外的好處，即在訓練中所要更新的總參數更少。通過將批量大小增大為 8096，以及使用線性學習率調整，已經能在一小時內完成在 ImageNet [12] 上的訓練了 [9]。一種名為 LARS [13] 的技術支持使用高達 32k 的批量大小，[14] 中最近還與混合精度訓練結合到了一起，使用 64k 的批量大小，在 4 分鐘內成功完成了在 ImageNet 數據庫上的訓練。

(2) 節點之間的通信

分布式訓練還有另一個重要組件，即節點之間的數據通信。多虧了 GFS [15] 和 Hadoop [16] 等一些分布式文件系統，這已經是一個成熟的研究主題了。以點對點的方式在節點之間實現高效且可感知帶寬的通信需要集合通信原語(collective communication primitives)[17]，這首先是在高性能計算(HPC)系統中引入的，之后由 [18] 帶進了深度學習領域。TensorFlow [19] 和 PyTorch 等現代深度學習框架使用了這些原語來進行 All Reduce 過程，因為其允許在最優的時間內完成互連節點之間的梯度傳輸。All Reduce [17] 在實際應用中具有多種變體，比如 Ring All Reduce、遞歸減半或倍增(Recursive Halfing/Doubling)、Binary Blocks 算法。

在分布式訓練中，為了實現有效的橫向擴展，計算與通信的配置必須保證最優。如果通信步驟是高效的且能與各個機器的計算保持同步，即整個集群中的計算應該在大致同一時間結束，那么訓練才是最優的。如果網絡速度慢，那么節點之間的通信就會成為瓶頸。梯度壓縮和混合精度訓練都是很有潛力的技術，能夠增大網絡的整體吞吐量。近期的一項研究 [20] 已經發現使用周期性的學習率能將實現網絡收斂所需的 epoch 數量降低 10 倍，使其成為了分布式訓練方面一個很有潛力的研究方向。

(3) 隨機梯度下降(SGD)的變體

隨機梯度下降 [21] 是一種用于訓練神經網絡的優化算法。這是梯度下降的一種變體，是一種用于調整權重的算法，能在每次反向傳播步驟之后使結果更接近最小值。SGD 不同于單純的梯度下降，因為其處理的是 mini-batch，而非單個訓練樣本。其形式如下：

其中 w_(t+1) 是為當前批計算出的權重，n 是 mini-batch 中的訓練樣本的數量，∇l(x, w_t) 是為前一個訓練樣本計算出的梯度。

對于分布式的情況，SGD 大致可分為兩類：異步 SGD 和同步 SGD。后續的章節會詳細介紹這兩種 SGD 和它們的變體。

3. 同步 SGD

同步 SGD 是一種分布式梯度下降算法，這是當前用于分布式訓練的最常用優化方法之一。網絡中的節點首先在它們的本地數據批上計算出梯度，然后每個節點都將它們的梯度發送給主服務器(master server)。主服務器通過求這些梯度的平均來累積這些梯度，從而為權重更新步驟構建出新的全局梯度集。這些全局梯度會通過使用與單機器 SGD 同樣的配方來更新每個節點的本地權重。這整個過程都類似于在單臺機器上通過單個數據 mini-batch 計算前向通過和反向傳播步驟，因此同步 SGD 能保證收斂。但是同步 SGD 也存在一些局限性。

4. 異步 SGD

異步 SGD 是一種分布式梯度下降算法，允許在不同節點上使用不同的數據子集來并行地訓練多個模型副本。每個模型副本都會向參數服務器請求全局權重，處理一個 mini-batch 來計算梯度并將它們發回參數服務器，然后參數服務器會據此更新全局權重。因為每個節點都獨立計算梯度且無需彼此之間的交互，所以它們可以按自己的步調工作，也對機器故障更為穩健，即如果一個節點故障，其它節點還能繼續處理，因此能消除由同步 SGD 引入的同步屏障(synchronization barrier)問題。

5. Ring 算法

[17] 中的 Ring All Reduce 結合了兩種算法：scatter-reduce 和 all gather。

圖 1：Ring 算法

scatter reduce 工作 p-1 個步驟，其中 p 是機器的數量。梯度向量被分為 p 塊。

算法 1：Ring 算法的 scatter reduce

在 scatter reduce 過程結束后，每臺機器都會有一部分最終結果。現在，每臺機器只需將自己那部分結果廣播給所有其它機器即可。這是使用 all gather 過程完成的，非常類似于 scatter gather，只是在接收數據上有些約簡，這部分只是被看作是最終結果存儲起來。

算法 2：Ring 算法的 all gather

6. 遞歸減半和倍增算法

[17] 中的遞歸距離倍增和向量減半算法使用了 4 種不同的原語，如下所示：

• 遞歸向量減半：向量在每個時間步驟減半。
• 遞歸向量倍增：分散在各個進程的向量的各個小塊被遞歸式地收集起來，構建成一個大的向量。
• 遞歸距離減半：機器之間的距離在每次通信迭代后減半。
• 遞歸距離倍增：機器之間的距離在每次通信迭代后倍增。

圖 2：減半和倍增算法

類似于 Ring 算法，這種 All Reduce 算法也由兩個過程構成：scatter-reduce 和 all gather。該算法與 Ring 算法的不同之處是這些過程執行運算的方式。

算法 3：Scatter Reduce 向量減半算法

算法 4：All Gather 向量倍增算法

7. Binary Blocks 算法

Binary Blocks 算法是對遞歸距離倍增和向量減半算法的延展，它的目標是當機器數量不是 2 的乘方數時降低負載的不平衡程度。在用于非 2 的乘方數情況的原始算法中，有一些機器會被閑置在一旁，直到算法執行完成，之后它們會接收結果得到的向量。這種方法會導致某些情況下大量機器空閑，比如，對于 600 臺機器構成的集群，會有 86 臺機器空閑，進程只會在另外 512 臺機器上執行。使用這種方法會出現很顯著的負載不平衡。

Binary Blocks 算法是通過將機器的數量分成 2 的乘方數的多個模塊來緩解這個問題。舉個例子，對于 600 臺機器構成的集群，可以分成 4 組，其中每組各有 2^9、2^6、2^4、2^3 臺機器。

算法 5：Binary Blocks 主服務器算法

算法 6：Scatter Reduce 主客戶端算法

8. 容錯式 All Reduce

如果一個分布式集群是由低可靠度的設備構成的，那么一旦遇到機器故障就需要重啟 All Reduce 算法。我們提出了一種容錯式 Binary Blocks 算法，將 Raft 共識算法 [31] 的元素整合進了 All Reduce 中。該算法能應對機器故障，只要備份的副本仍可運行，就能繼續執行。

圖 3：Raft 算法

9. 調整批量大小

在實際訓練深度神經網絡時，學習率會隨著訓練經過多個 epoch 后而緩慢逐漸減小。這背后的直觀思想是讓權重在訓練初期邁開更大的步子，隨著模型接近收斂，步子也越來越小。這在實踐中的效果相當好，并且能比使用固定學習率訓練的模型達到更好的收斂程度。這也是一種很高效的方法，因為初期采用較大的學習率能夠在使用更小的學習率微調之前取得很好的進展。但是，使用較大的批量大小進行訓練是一個很有前景的研究方向，因為這能顯著加速訓練過程，能將訓練時間從數天降至幾分鐘，正如 [9,14,13] 中證明的那樣;[11] 的研究更是增強了這一趨勢，通過實驗證明增大批量大小就相當于降低學習率。

使用更大的批量大小進行訓練的優勢是模型執行的整體權重更新的數量更少，這能讓訓練速度更快，如圖 4 所示。但是，相比于使用更小批量大小訓練的模型，直接使用大批量進行訓練也會出現問題，比如過早發散以及最終驗證準確度更低。

圖 4：衰減學習率與增大批量大小 [11]

10. 張量融合

對于 ResNet 等某些常見的模型，研究者已經觀察到為梯度計算的張量的大小是相當小的。更具體而言，用于卷積層的梯度張量大小比全卷積層的要小得多。這個現象是很突出的，因為通過線路發送少量數據可能會導致出現大量延遲，同時也沒有充分利用網絡帶寬。一種解決這一問題的簡單方法是張量融合 [14]，簡單來說就是將多個小張量融合到一起，形成一個至少超過某個最小大小的張量，之后再在網絡中發送這個融合得到的張量。執行這種融合的好處是降低每臺機器起始時間的負載以及降低網絡流量的整體頻率。這能讓網絡變得不再亂七八糟以及在最優時間內完成任務。但是，為小張量使用張量融合可能會導致 Ring All Reduce 變得低效而緩慢，[14] 提出了一種分層式 All Reduce，使用了多層主從設置，觀察表明這種方法能帶來更低的延遲。這種分層式 All Reduce 的工作方式是將所有機器分成不同的批，之后每一批都選擇一個主節點來聚合梯度。這些主節點在它們自身上執行 Ring All Reduce，之后該主節點將更新后的梯度分配給它們各自的從節點。這種策略是通過降低批的數量來降低延遲開銷，進而逐步降低 Ring All Reduce 的開銷。使用張量融合能減少小網絡的處理，并提升網絡的整體速度，這是很值得推薦的。這種方法在 Horovord [38] 和騰訊的框架 [14] 等產業系統中得到了廣泛的使用，使其成為了現代分布式訓練框架中的重要一員。

11. 低精度訓練

在 ImageNet 數據庫 [12] 上訓練一個 ResNet 模型 [29] 的最快時間目前是 4 分鐘 [14]。在那項研究中，研究者使用低延遲的零副本 RDMA 網絡連接了 2048 個 GPU，并且組合性地使用了 LARS 算法、混合 All Reduce 算法、張量融合和混合精度訓練。混合 All Reduce 算法結合了 Ring All Reduce 和分層式的版本，并根據當前步驟的張量大小來進行切換。他們還使用了一種全新的方法才實現了這樣的整體訓練速度增益：混合精度訓練 [39]。由此所帶來的吞吐量和帶寬效率增長將訓練速度提升了 8 倍。顧名思義，混合精度訓練是使用兩種不同的數據類型來訓練神經網絡——更小的數據類型用于大多數運算，更大的數據類型用于對精度要求嚴格的運算。

神經網絡最初是使用單精度或雙精度數作為默認數據類型，因為這些數據類型在獲取網絡想要建模的任務的表征上表現很好。單精度數是 32 位浮點數，雙精度數是 64 位浮點數。近期有研究表明通過在更低精度數據類型進行訓練，可將神經網絡的速度和大小降低 50%-80% [40, 41]。一種常用方法是使用 16 位浮點數來訓練網絡(FP16 訓練)，但與使用單精度訓練的同樣網絡相比，這些網絡的測試準確度會更差一些 [39]。這種情況的原因主要是權重更新步驟的精度更低。更具體而言，將低精度梯度與學習率相乘有時會導致數值溢出 16 位的范圍，從而導致計算不正確，進而導致最終驗證準確度損失。

混合精度訓練 [39] 的目標是通過使用單精度(32 位)的權重主副本并以半精度(16 位)運行其它一切來解決這個問題。

混合精度訓練能實現更高的吞吐量，從而能降低計算與通信瓶頸。但是，混合精度訓練也存在一些需要注意的是像，即損失丟失和算術精度更低。

12. 梯度和參數壓縮

擴展分布式訓練過程的一個主要瓶頸是節點之間的模型權重和梯度通信具有很高的帶寬成本。當在使用了聯盟學習(federated learning)[42] 的設備(尤其是移動設備)上訓練時，這個瓶頸會尤其顯著，因為其存在網絡帶寬低和連接速度慢的問題。針對這一問題，研究者已經提出了多種用于高效利用網絡帶寬的方法。使用 SGD 的異步和同步變體允許節點各自獨立地通信，同時還能實現并行化并將網絡帶寬利用提升到一定程度 [10,43,44];另外在梯度壓縮上已經取得一些顯著進展也很有希望 [45]。這些方法主要基于兩大思想：量化和稀疏化。

算法 7：用于節點 k 上單純的動量 SGD 的深度梯度壓縮

13. 未來研究

過去幾年來，分布式訓練領域已經取得了很大的進展，并已為進一步創新做好了準備。相關領域目前正在增大 mini-batch 大小的限制，直到網絡不會發散或降低同步 SGD 的最終驗證準確度的程度，還在研究解決異步 SGD 的 stale 梯度和最終驗證準確度降低的問題。在更低功耗的設備上進行訓練已經以聯盟學習(federated learning)[42] 的形式獲得一些發展勢頭，也出現了一些用于安全和去中心化訓練的現代深度學習框架構建模塊。在消費級設備上訓練有一些好處，其中之一是能夠打造能基于消費者交互學習特定習慣的智能應用，而且通過在設備上存儲數據和執行訓練，還能保證用戶的數據隱私。這類應用的一個案例是 Android Predictive Keyboard(安卓預測鍵盤)，它能在設備上學習用于預測下一個詞的小型個性化模型。在低功耗設備上訓練的一個關鍵難題是可用的網絡帶寬和計算資源很少。高效的框架也許可以實現在手機和物聯網設備上的大規模訓練，從而實現便攜式的深度學習應用。[54] 已經在低功耗設備上的分布式訓練方面做出了一些出色的工作，其使用了強化學習算法來調度在異構式設備集群上的訓練任務。在商品級設備上進行分布式訓練需要在訓練和通信方面進行多種優化。梯度壓縮和混合精度訓練是少數幾種結果優良且很有希望的方向。整體而言，這個領域的研究方向很活躍且存在很多創新，并已經做好準備，即將成為更廣泛智能應用的核心組件。

總結

我們介紹和總結了分布式訓練框架的各個組件。為了打造一個高效且可擴展的分布式訓練框架，我們推薦使用以下技術：

推薦使用同步 SGD 算法來進行訓練，因為它有嚴格的收斂保證。

應該使用 Binary Blocks 算法來執行梯度累積的 All Reduce 流程，因為其運行時間更優。

為了有效地使用硬件和網絡帶寬，應該在框架中使用梯度壓縮、量化和混合精度訓練等多種技術。我們推薦組合式地使用深度梯度壓縮和混合精度訓練。

應該使用非常大的批量大小來進行訓練，以最大化并行能力和最小化運行時間。我們推薦使用 LARS 算法，因為已有研究證明它能夠以高達 64000 的批量大小足夠穩健地訓練網絡。

【本文是51CTO專欄機構“機器之心”的原創譯文，微信公眾號“機器之心( id: almosthuman2014)”】

戳這里，看該作者更多好文