分布式深度學習系統網絡通信優化技術

（國防科技大學，中國 長沙 410073）

(National University of Defense Technology, Changsha 410073, China)

伴隨人工智能研究的第3次浪潮，深度學習技術席卷了圖像分類、語音識別、自動駕駛、內容推薦等眾多應用領域。深度學習的普及促進了深度神經網絡的發展，這些網絡模型在各種各樣的任務上都取得了良好的效果。以計算機視覺為例，一些經過精心設計的深度神經網絡，例如在ImageNet數據集上訓練的GoogLeNet和ResNet-50，已經在圖像分類任務上擊敗了人類。然而，訓練高性能的網絡模型常常需要花費大量時間，短則數小時，長則數天甚至數周。

為了減少訓練時間，研究人員通常使用高性能硬件如圖形處理器（GPU）和張量處理器（TPU）等加速深度神經網絡模型的訓練。同時，在多個節點上并行訓練深度神經網絡也是行之有效的加速方法。每個節點僅僅執行整體計算任務的一部分，所以這種分布式訓練可以大幅縮短深度神經網絡的訓練時間，如圖1所示。但是，由于網絡模型訓練過程的迭代性，不同的計算節點之間往往需要頻繁地進行通信以交換大量的數據，這就導致節點間的通信成為分布式訓練中的關鍵瓶頸，圖2描述了這種情況。隨著集群規模的擴大，節點間的通信開銷會急劇地增加。這種現象極大地削弱了分布式訓練所帶來的優勢，因為很大一部分時間都花在了網絡數據傳輸上。當我們使用高性能硬件訓練網絡模型時，通信時間占整體訓練時間的比例會進一步增加，因為這些高性能硬件只減少計算開銷而沒有降低通信開銷。高昂的通信開銷限制了分布式訓練的可擴展性。為此，研究者們針對分布式訓練過程中的通信行為，展了一系列的優化工作。

1 分布式深度學習系統的基本概念

當網絡模型的訓練任務被部署到多個節點上時，我們需要考慮以下幾個問題：并行化訓練任務的哪一部分？如何組織計算資源？如何協調各個計算節點？

1.1 數據與模型并行

分布式深度學習系統面臨的第一個問題是并行化訓練任務的哪一部分。常見的兩種并行模式為數據并行和模型并行，如圖3所示。

在數據并行中，首先整個數據集被隨機且均勻地分配到各個節點中，每個節點都在本地維護一個完整的模型副本。每個節點僅讀取和處理唯一的數據子集，并在訓練期間更新本地模型。然后，將這些本地模型參數與其他節點同步以計算全局參數。因此，這些全局參數利用網絡分配到每個節點上，以便開始下一次迭代。

除了數據并行，模型并行是另一種方法。該方法將神經網絡模型分割到不同的計算節點上，每個計算節點僅負責計算模型的某一部分，只有輸入層所在的節點才負責讀取數據。當模型太大而無法放到一臺機器中時，則需要使用模型并行。本文主要關注常見的數據并行技術。

▲圖1 計算與通信相對均衡的分布式訓練

▲圖2 通信開銷占比過大的分布式訓練

▲圖3 數據并行與模型并行

1.2 中心化與去中心化架構

分布式深度學習系統所面臨的第二個問題是以何種結構組織計算節點。系統的架構會影響訓練過程中數據的傳輸方式和模型更新方式，進而影響訓練時間。常見的架構包括中心化以及去中心化架構。

參數服務器架構是分布式深度學習中最常見的中心化架構。參數服務器架構通常包括若干服務器和若干工作節點。其中服務器上存放著全局共享的模型參數。如果有多個服務器節點，那么模型的權重參數會被拆分到每臺機器上。每個工作節點都存儲一個模型副本（使用數據并行）或存儲模型的某一部分（使用模型并行）。工作節點通過推送/拉取操作與服務器進行通信，而任何工作節點之間都不會產生通信行為。每次迭代中，每個工作節點首先讀取數據，然后基于這些數據計算本地模型的梯度。隨后，節點將其本地梯度推送到參數服務器。服務器接收所有工作節點發送的梯度后，首先聚合這些梯度，然后更新全局模型，最后工作節點再從服務器拉取最新的模型權重，并使用它進行下一次的迭代。

點對點架構是一種去中心化的架構，其中每個節點在模型訓練中都扮演相同的角色。與參數服務器類似，點對點中的每個節點都擁有完整的模型副本，并首先基于小批量數據計算本地梯度。對等體系結構使用歸約和廣播等集合通信操作，而不是參數服務器體系結構中的推送/拉取操作。經過梯度計算后，每個對等方首先從其他對等方接收梯度，然后對這些梯度求平均，這就是歸約步。然后，對等方將其本地梯度廣播給所有其他對等方。節點從其他對等節點收到所有梯度后，立即更新其本地模型，然后執行下一個迭代。

1.3 同步與異步更新

將訓練任務并行化到多個節點上時，如何協調這些節點是一個大問題。同步更新和異步更新是當前的兩種主流方法。

使用同步更新時，兩次迭代之間存在全局屏障。網絡中的節點首先計算局部梯度，然后將梯度與其他節點同步。由于各種因素，每個節點都有不同的完成時間。在一次迭代中，由于存在全局屏障，計算較快的節點不得不等待計算較慢的節點。全局屏障的存在導致了“掉隊者”的問題，較快的節點需要大量的等待時間。

為了克服此限制，研究人員提出了異步更新方案。此方案中，兩次迭代之間沒有全局屏障，每個節點在從其他節點接收到參數之后（無論這些參數是舊參數還是新參數），都立即開始下一次迭代。每個節點之間沒有同步屏障，因此某些節點在更新全局模型參數時可能會使用過時的模型參數來計算梯度。數據的陳舊性將偏差和方差導入局部梯度和全局模型參數中，這會導致模型收斂緩慢且沒有收斂性保證。為了應對掉隊者和數據陳舊性的問題，延遲異步更新被提出。該方案是同步更新和異步更新的結合，并對數據的陳舊性做了限制。

2 分布式深度學習系統網絡通信優化技術

為了有效地在多節點上訓練深度神經網絡，降低訓練過程中的通信開銷，業界已經提出了多種通信優化策略。由于分布式深度學習通信優化涉及深度神經網絡、異構并行計算、分布式系統、計算機網絡等眾多技術，從芯片級到系統級有許多涉及網絡通信優化相關的技術。

本文嘗試從網絡協議棧層次的角度，從通信流量調度層和網絡通信執行層的角度，對分布式深度學習系統的網絡通信優化技術進行初步分類討論，如圖4所示。從通信流量調度層的角度來看，可以通過降低通信發生的頻次，來降低通信數據量（梯度壓縮）以及計算通信重疊等技術優化分布式訓練的通信過程。在網絡通信執行層面，不同的參數通信架構、不同的消息傳遞庫以及不同的網絡協議都會對通信產生影響。網絡與通信領域的研究人員，對網絡通信執行層面的優化技術往往更為熟悉，實際上網絡底層主要采用的還是通用的網絡技術。為了追求高性能，分布式深度學習系統的網絡通信優化必然是對通信全協議棧的定制優化和協同設計。從這個角度來看，定制優化和協同設計需要深入分析分布式深度學習訓練系統的通信需求，充分利用底層網絡技術的特點，從而對分布式深度學習訓練系統的通信流量進行高效調度與優化。

2.1 降低通信發生的頻次

訓練深度神經網絡時，整個訓練過程通常需要進行多次epoch和迭代。降低通信開銷的一種直接方法是減少通信輪次，而通信輪次與批量大小和通信周期有關。較大的批量和較長的通信周期都可以減少數據交換次數。

2.1.1 使用大批量進行訓練

▲圖4 通信優化層次

批量大小控制每次迭代時讀取的數據量。在基于數據并行的分布式訓練中，批量通常指的是全局批量，即所有節點的本地批量的總和。對于并行化訓練，節點通常在每次迭代結束時交換數據，因此，當訓練數據集的大小固定時，增加批量會減少迭代次數，從而減少通信輪次。梯度和參數的形狀及大小僅取決于神經網絡模型本身，因此單次迭代傳遞的消息大小總是保持不變，更改批量不會改變每次傳輸的消息總量[1]。大批量帶來的好處是更少的通信輪次。由于具有大批量訓練的優勢，分布式深度學習系統網絡通信優化技術最近的工作將批量增加到8 k、32 k甚至增加到64 k樣本（1 k表示1 024個樣本）。然而，當固定epoch時，與小批量相比，實際上直接使用具有大批量的隨機梯度下降（SGD）通常會帶來泛化能力的下降。由于梯度估計存在方差，小批量SGD始終收斂到平坦的最小值，而大批量SGD趨向于收斂到尖銳的最小值——尖銳的最小值導致較差的泛化能力。通常，神經網絡模型的泛化能力相較于訓練速度來說更為重要。因此，我們需要保證在運行相同的epoch下，使用大批量訓練的模型具有與小批量訓練的模型相近的泛化能力。

2.1.2 周期性通信

通信輪次不僅與批量的大小有關，還與通信周期相關。前面提到，通信操作常常發生在每次迭代的末尾。因此，我們可以控制訓練進程，每隔幾次迭代進行通信，進而降低通信操作發生的次數。實際上，最常用的方法是本地SGD[2]。為了降低多次通信帶來的巨大開銷，本地SGD讓每個工作節點先在本地進行若干次迭代，然后再與參數服務器同步梯度和全局參數，如圖5所示。實驗證明，只要周期設置合適，本地SGD可以減少訓練時間的通信開銷并且保證收斂性。

2.2 降低通信過程中傳輸的數據量

在傳統的深度神經網絡分布式訓練過程中，計算節點之間交換梯度和模型參數，以進行模型聚合和本地計算。當梯度和參數量很大時，由于交換大量32位浮點變量而導致的通信瓶頸，削弱了并行化帶來的優勢。

解決此問題的一個直接想法是壓縮傳輸的梯度，最常用的兩種方法是梯度量化和梯度稀疏化。前一種方法使用低精度數字（即8位、4位甚至1位）替換32位浮點數，以此來減少傳輸的梯度量；而后一種方法則選擇了梯度向量中的一些重要元素來更新模型參數，以避免不必要的傳輸開銷。

2.2.1 梯度量化

梯度和模型參數中的每個元素都存儲在電氣和電子工程師協會標準（IEEE 754）的單精度浮點變量中。交換大量的32位全精度變量經常占用大量網絡帶寬。為了緩解這種通信瓶頸，研究人員嘗試使用低精度浮點數來表示訓練過程中傳輸的梯度。

▲圖5 傳統并行SGD與本地SGD的對比

梯度量化的一般性描述為：首先梯度在發送之前通過量化函數量化為低精度值，然后在模型訓練期間，每個節點之間傳輸這些低精度值；接收節點通過反量化函數從量化值重構原始梯度，再聚合這些重構后的梯度向量并更新模型參數。根據所使用的量化函數的性質，量化操作可以分為確定性量化和隨機性量化。

確定性量化將梯度元素值映射到某些固定的值。1 bit SGD[3]是確定型量化的典型例子，它根據梯度元素的取值范圍，將所有的梯度值量化到0和1這兩個數字。接收端在收到量化的01序列后，會把其中的0解碼成-1，把1解碼成+1，再進行下一步的訓練。通常來說，因為確定性量化的固定量化形式會有較多的信息丟失，隨機性量化應運而生，這種方法使得量化后的梯度元素值仍然服從某一概率分布。該方法引入了額外的隨機性，因此量化后的梯度一定是原始梯度的無偏估計，并且必須具有方差約束以確保具有像傳統SGD一樣的收斂性。典型的隨機性量化方法如量化SGD（QSGD）[4]和TernGrad[5]等。

2.2.2 梯度稀疏化

由于我們需要至少1位來表示梯度向量中的每個元素，因此梯度量化方法最多只能將數據量壓縮32倍。梯度稀疏化則沒有上述壓縮率的限制，它只關注更新過程中梯度向量和模型參數中的一些重要值，即那些值遠大于零的數字。根據選擇的元素數量，稀疏化方法可以達到非常高的數據壓縮率。稀疏化方法的核心是如何從梯度向量中選擇有效值，即如何將稠密更新轉換為稀疏更新。常見的稀疏化方法是Top-K，它保留了梯度向量中前K個較大的絕對值，而將其他值設置為零。比如，基于Top-K的Gradient Dropping算法[6]在手寫字符識別和機器翻譯等任務中取得了良好的加速效果。關于梯度壓縮，如何平衡模型精度與數據壓縮比率是一個嚴峻的挑戰。誤差補償已被證明是一種行之有效且兼容大部分壓縮算法的技術[7]。此外，還有研究者使用了動量屏蔽和梯度裁剪等技術保證模型的精度[8]。另外，研究者還應該關注如何降低梯度壓縮技術的計算開銷，以進一步加速模型訓練。

2.3 計算與通信重疊

現有的深度學習框架后端引擎以先進先出（FIFO）的順序執行操作。因此，按其生成順序發送梯度，最后一層（輸出層）的梯度先被發送，然后處理中間層，最后處理輸入層。要完成當前迭代并盡早地開始下一次迭代，我們需要減少計算與第一層通信之間的延遲。

Poseidon系統提供了一種基于分布式訓練的固有特性的無等待反向傳播調度算法[9]：一旦反向傳播計算出某一層的梯度，后端引擎就開始傳輸該層的梯度數據。但是，不同層的參數量可能并不相同，因此具有不同的計算和通信時間。這就意味著Poseidon不一定會比某些特定網絡模型上的FIFO調度表現更好。

P3[10]（Priority-base Parameter Propagation）通過基于優先級的調度擴展了無等待反向傳播。在P3中，靠近輸入層的梯度向量具有較高的優先級，而靠近輸出層的梯度具有較低的優先級。在訓練階段，無論高優先級的梯度向量在何時生成，都將優先對其進行處理。這樣就保證了其他節點能夠盡早地接收到靠近輸入層的梯度，從而能夠盡早開始下一次迭代。此外，P3使用張量分割技術將各層的參數向量分解為適當的小塊，并根據其所在的層為每個切片分配優先級，以實現更細粒度的流水線化。高效的計算通信重疊率可以顯著地加速模型訓練過程。但是，當前大部分調度算法都是啟發式的，這就意味著調度算法并不是最優解。因此，尋找高效的計算、通信調度算法至關重要。目前，已經有研究者將貝葉斯優化算法[11]和強化學習[12]應用在算子調度上。

單步延遲SGD（OD-SGD）[13]算法打破了下一次迭代計算對上一次迭代中通信過程的依賴，通過實現計算過程與通信過程的高度重疊來提升分布式訓練性能。該算法結合了同步SGD和異步SGD兩種更新算法的優勢，在保證訓練精度的情況下提高分布式訓練的速度。圖6展示了在同步SGD算法訓練模式和在OD-SGD算法訓練模式下的性能對比。假設單次迭代的計算開銷和通信開銷均為3個時間單位，則在同步SGD算法訓練模式下，單次迭代訓練的時間開銷為5個時間單位；而在OD-SGD算法訓練模式下，單次迭代的時間開銷為3個時間單位，原來額外的2個時間單位的通信開銷被隱藏。

2.4 參數通信架構

在去中心化架構下，分布式深度學習的訓練過程符合全局規約操作的語義：每個節點獨立計算局部梯度，然后通過全局規約運算來計算梯度總和并將其發送給所有對等節點。顯然，上述操作可以看作一個Allreduce運算。因此，整個網絡通信過程的性能也就取決于Allreduce的性能。

在高性能計算領域，關于Allreduce算法的研究已經非常充分了。Allreduce是近兩年才在深度學習訓練系統中得到廣泛應用的，其中很有代表性的是Ring-Allreduce，由百度于2017年引入應用到主流的深度學習框架中。Ring-Allreduce由兩個階段組成：Reduce-Scatter和Allgather。當使用p個計算節點時，每個階段都包括p-1個通信步驟。每個節點都維護其本地梯度，本地梯度被平均分為p塊。在Reduce-Scatter階段，每個節點發送和接收張量的不同塊。在p-1個步驟之后，每個節點都具有了一部分塊的全局結果。在Allgather階段，每個節點發送自己維護的部分全局結果，并從其他節點接收其他部分全局結果。同樣是p-1個步驟之后，每個節點都擁有完整的全局結果，如圖7所示。因此，Ring-Allreduce總共需要2×(p-1)個通信步。

此外，關于參數服務器架構及通信特征的詳細分析，請參考文獻[14]。

2.5 網絡消息庫

在去中心化架構下，目前關于Allreduce的實現非常多，包括NVIDIA的集合通信庫、Facabook的Gloo以及百度的Allreduce等。以NCCL為例，它針對NVIDIA的GPU，實現了單機多卡以及多機多卡之間的高效GPU通信。在當前主流的跨平臺統一分布式訓練框架Horovod[15]中，NCCL、Gloo以及MPI都已經得到了支持。

對于參數服務器架構來說，由于其固有的“少對多”的通信模式，因此在常見的實現中，底層通信模塊總是依賴于點對點的消息庫，例如TensorFlow中的gRPC，以及MXNet中默認使用的ZMQ消息庫。盡管MPI中也實現了低延遲的點對點通信，但是在參數服務器架構下并不能發揮出MPI在集合通信方面的優勢。

2.6 網絡協議優化

▲圖6 傳統同步隨機梯度下降與單步延遲隨機梯度下降的對比

▲圖7 Ring-Allreduce過程

早期的分布式訓練框架的通信模塊實現通常基于傳輸控制協議（TCP）/互聯網協議（IP），需要先將參數數據復制到內核態的網絡協議棧中，再通過網絡接口發出去，這些復制操作增加了分布式訓練的通信延遲。遠程直接內存訪問（RDMA）允許用戶態進程直接讀取和寫入遠端進程的地址空間，是傳統高性能計算系統中常用的高帶寬低延遲的通信技術。在當前主流的分布式訓練框架中，RDMA原語替換了原始的基于套接字和TCP的接口。實驗表明，使用RDMA替換傳統的TCP/IP協議可以大幅降低分布式訓練的同步開銷，提升訓練速度，擴大訓練規模。

3 結束語

通信開銷是擴展大規模深度學習的障礙，訓練過程中高昂的通信代價令人難以接受。本文主要從算法和網絡兩個層面介紹了分布式深度學習中的通信優化策略，其中大部分優化策略都是正交的，它們可以進行組合，以進一步降低分布式深度學習訓練過程中的通信開銷，加速網絡模型訓練。關于神經網絡模型訓練，還有一個重要的研究方向是關于性能模型以及測量工具的探索與實現。性能模型能夠幫助我們在理論上分析各類開銷，而測量工具則允許我們找到分布式訓練中的各類瓶頸。盡管TensorFlow[16]與MXNet都提供了相關工具，但是我們仍然需要更高級的網絡分析工具如Horovod Timeline[15]以及 SketchDLC[14]等。

致謝

本文的部分研究成果和內容撰寫得到國防科技大學徐葉茂博士生的幫助，謹致謝意！