論文地址：https://arxiv.org/abs/1810.11787
深度學習已經為人工智慧領域帶來了巨大的發展進步。但是，必須說明訓練深度學習模型需要顯著大量的計算。在一臺具有一個現代 GPU 的單臺機器上完成一次基於 ImageNet 等基準資料集的訓練可能要耗費多達一週的時間，研究者已經觀察到在多臺機器上的分散式訓練能極大減少訓練時間。近期的研究已經通過使用 2048 個 GPU 的叢集將 ImageNet 訓練時間降低至了 4 分鐘。這篇論文總結了各種用於分散式訓練的演算法和技術，並給出了用於現代分散式訓練框架的當前最佳方法。更具體而言，我們探索了分散式隨機梯度下降的同步和非同步變體、各種 All Reduce 梯度聚合策略以及用於在叢集上實現更高吞吐量和更低延遲的最佳實踐，比如混合精度訓練、大批量訓練和梯度壓縮。
1 引言
A 背景和動機
資料正以前所未有的規模生成。大規模網際網路公司每天都會生成數以 TB 計的資料，這些資料都需要得到有效的分析以提取出有意義的見解 [1]。新興的深度學習是一種用於執行這種分析的強大工具，這些演算法在視覺 [2]、語言 [3] 和智慧推理 [4] 領域的複雜任務上創造著當前最佳。不幸的是，這些演算法需要大量資料才能有效地完成訓練，而這又會耗費大量時間。第一個在 ImageNet 分類任務上取得當前最佳結果的深度學習演算法在單個 GPU 上訓練足足耗費了一週時間。在如今的時代，這樣的速度已經完全不行了，因為現在的模型訓練所使用的資料甚至會讓 ImageNet 資料集的規模都相形見絀。現在存在以橫向方式延展深度學習訓練的內在需求，同時還要保證能夠維持單 GPU 模型那樣的準確度。理想情況下，這種訓練的速度應該隨機器數量的增加而線性增大，同時還要能容錯以及能在高延遲網路條件下收斂。
B 分散式訓練概述
神經網路的分散式訓練可以通過兩種方式實現：資料並行化和模型並行化。資料並行化的目標是將資料集均等地分配到系統的各個節點（node），其中每個節點都有該神經網路的一個副本及其本地的權重。每個節點都會處理該資料集的一個不同子集並更新其本地權重集。這些本地權重會在整個叢集中共享，從而通過一個累積演算法計算出一個新的全域性權重集。這些全域性權重又會被分配至所有節點，然後節點會在此基礎上處理下一批資料。
模型並行化則是通過將該模型的架構切分到不同的節點上來實現訓練的分佈化。AlexNet [2] 是使用模型並行化的最早期模型之一，其方法是將網路分攤到 2 個 GPU 上以便模型能放入記憶體中。當模型架構過大以至於無法放入單臺機器且該模型的某些部件可以並行化時，才能應用模型並行化。模型並行化可用在某些模型中，比如目標檢測網路 [5]，這種模型的繪製邊界框和類別預測部分是相互獨立的。一般而言，大多數網路只可以分配到 2 個 GPU 上，這限制了可實現的可擴充套件數量，因此本論文主要關注的是資料並行化。
本論文大致分為六個章節，第一節將介紹已有的優化訓練演算法。第二節將關注用於連線網路各節點的通訊策略。第三節會探索一些具體技術，比如大批量訓練、梯度壓縮以及用於在低功耗裝置和低速網路條件下實現有效訓練的混合精度訓練。第四節則會消化之前章節的資訊，並選擇出用於不同設定的最優的訓練演算法和通訊原語。最後兩節分別是未來研究方向和總結。
2 分散式訓練框架的元件
A 分散式訓練演算法
一個常用於分散式設定中的訓練的常用演算法是隨機梯度下降（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 資料庫上的訓練。
B 節點之間的通訊
分散式訓練還有另一個重要元件，即節點之間的資料通訊。多虧了 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 倍，使其成為了分散式訓練方面一個很有潛力的研究方向。
隨機梯度下降（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 的批量大小足夠穩健地訓練網路。