一種面向分布式深度學習的輕量級聚合通信庫*

王笑雨,董德尊

(國防科技大學計算機學院,湖南 長沙 410073)

1 引言

深度學習技術在許多領域應用廣泛，特別是在計算機視覺和自然語言處理領域。隨著處理的問題更加復雜，同時為進一步提高模型精度，深度神經網絡DNN(Deep Neural Network)層數隨之增多，數據集規(guī)模也逐漸變大，導致在單節(jié)點上進行相關模型的訓練變得格外耗時。為縮短模型訓練時間，分布式訓練成為主流深度學習框架(如MXNet[1]、TensorFlow[2]和PyTorch[3])常用的解決辦法。分布式訓練利用多個節(jié)點的計算資源來加速大規(guī)模深度學習模型的訓練，通常分為數據并行和模型并行2種方式。其中數據并行方式更為主流，其將數據集拆分為若干份并分配給各節(jié)點，各節(jié)點在拆分后的數據集上訓練本地的全局模型副本，并通過在節(jié)點間交換訓練過程中各自產生的梯度信息來更新全局模型的參數，從而完成在整個數據集上對模型的訓練。這一方法在使得更多的計算資源可以被利用的同時也引入了通信開銷。而隨著節(jié)點的增多，節(jié)點間的梯度交換變得頻繁，由此產生的通信開銷占據了整個訓練過程中的較大部分，減少了分布式并行化所帶來的優(yōu)勢，成為加速分布式訓練的瓶頸。

聚合通信操作是目前主流的用于梯度信息同步的方式之一，被主流的深度學習框架所采用。為進一步縮短分布式訓練耗時，許多研究工作針對這一操作進行了優(yōu)化。比如，通過調整聚合通信操作的執(zhí)行順序來利用計算開銷隱藏部分通信開銷，甚至通過對其進行拆分和合并來使得計算和通信盡可能重疊[4,5]。還有通過梯度量化[6,7]和稀疏化[8]等方法來減少需要交換的梯度信息的大小，從而減少通信開銷。這些工作都將聚合通信操作視作黑匣子，并從如何有效使用的角度來進行優(yōu)化，在一定程度上縮短了訓練時間。

最近一些研究工作采用了新的改進思路——根據分布式訓練的特征來調整聚合通信操作本身。eager-SGD(Stochastic Gradient Descent)[9]考慮到訓練過程的魯棒性提出了部分聚合通信操作,與原先操作不同的是，其在語義邏輯上并不要求所有節(jié)點都參與梯度同步過程，減少了等待掉隊者的開銷。此外，考慮到分布式訓練中后產生的梯度信息會先被用到，文獻[10]提出了搶占式AllReduce，使得后產生的具有更高優(yōu)先級的梯度信息在邏輯上可以中斷低優(yōu)先級的梯度同步過程，以更快地進行同步。聚合通信操作在高性能計算領域已經使用了快30年，而其近期才被引入深度學習領域，還未完全針對這一應用場景進行優(yōu)化。上述這些工作考慮到了傳統(tǒng)聚合通信操作的語義和深度學習應用的特點不匹配的現(xiàn)象，并試圖緩解這一矛盾。此外，還有一些將分布式訓練特點與網絡結合進行優(yōu)化的工作改變了聚合通信操作的實現(xiàn)[11 - 13]。這些工作都說明研究專用于分布式訓練的聚合通信操作，并進行應用和傳輸的協(xié)同設計對減少通信開銷來說很有必要。

對于現(xiàn)有去中心化的分布式訓練框架來說，進行相關的協(xié)同設計較為困難。其通常采用MPI(Message Passing Interface)標準規(guī)定的聚合通信接口來完成實際的梯度同步操作。MPI具有不同的實現(xiàn)版本，比如OpenMPI和MPICH。這些版本為了保證魯棒性、普適性和擴展性，同時為了便于開發(fā)，通常會具有復雜的架構和接口調用，多種不同功能的通信函數以及龐大的代碼量，這使得盡管其在工業(yè)生產中穩(wěn)定健壯，但對在其上進行研究和改進并不友好。

為解決上述問題，本文設計并實現(xiàn)了一個針對分布式訓練中聚合通信操作的輕量級通信庫，使得相關的實驗更易實現(xiàn)。該輕量級通信庫整體架構簡單但提供了分布式訓練中必要的聚合通信操作，且向上支持主流的深度學習框架，向下支持集群中常見的網絡環(huán)境。該庫使用較少的封裝來將聚合通信操作的執(zhí)行細節(jié)暴露出來，以便于進一步分析和改動。此外，本文希望為研究分布式訓練中的聚合通信操作提供一個統(tǒng)一的實驗環(huán)境。通過該通信庫，相關研究人員可以簡單高效地實現(xiàn)自己的想法并在更多的環(huán)境配置下進行實驗，進而獲得更廣泛的影響。而在此庫基礎上開發(fā)的算法也可以更輕松地應用到多個深度學習框架中。最后本文從聚合通信操作和深度學習應用2方面對該庫進行了評估，并將其與常用的聚合通信庫進行了對比。最終實驗結果顯示，該庫在部分情況下具有優(yōu)勢。

2 背景介紹

2.1 分布式訓練和框架

分布式訓練主要用于解決在單節(jié)點上訓練大規(guī)模DNN耗時的問題，其通過將訓練任務進行拆分并部署到多個節(jié)點上同時進行訓練來加速模型的訓練過程。在現(xiàn)有的解決方案中，數據并行是較為流行的訓練方式之一，被主流深度學習框架廣泛使用。在數據并行的方式中每一個worker(對應一個進程，可能在CPU上也可能在GPU一類的加速器上)都擁有一份全局模型的副本和一部分數據集，其在擁有的數據集上獨立訓練本地全局模型副本，并周期性地與其他的worker同步各自的梯度信息，以更新全局模型的參數，從而實現(xiàn)在多個節(jié)點上基于整個數據集來并行地訓練模型。在更新全局模型參數的過程中，所有worker交換各自訓練過程中在反向傳播階段產生的梯度，并將這些梯度進行歸約(一般是求和或求平均值)，以使用隨機梯度下降SGD的變種算法來更新全局和本地的模型參數，更新后的參數會在下一輪訓練的前向傳播階段使用。值得注意的是，梯度產生和對應模型參數使用的順序通常是相反的，也就是說,在反向傳播階段越晚產生的梯度，基于其更新的模型參數會在前向傳播中越早被用到。

前述主流的深度學習框架都為用戶使用分布式訓練提供了便捷的方式，通常只需要在原有單節(jié)點訓練的代碼上進行簡單的修改即可。但是，這些框架具有不同的編程語義和性能，比如，TensorFlow采用的是聲明式的語義，用戶只需要聲明需要執(zhí)行的操作，這些操作由TensorFlow進行定義；PyTorch則采用了命令式的語義，需要用戶規(guī)定操作具體怎么執(zhí)行；MXNet則結合了這2種語義。對于梯度同步過程，這些框架也采用了不同的通信架構，PyTorch采用的是AllReduce架構，而TensorFlow和MXNet雖然采用的都是參數服務器PS(Parameter Server)架構，但實現(xiàn)方式不同，分別是Rendezvous和KVStore,如圖1所示。這些差異使得不同的深度學習框架具有各自不同的優(yōu)勢，比如，MXNet對初學者更加友好；PyTorch在使用上更加靈活；而TensorFlow則對在CPU上訓練大規(guī)模深度神經網絡具有更好的性能。這些具有不同特性的框架在學術界和工業(yè)界都被廣泛應用，對于分布式訓練的研究也非常重要。

Figure 1 Architecture of the proposed design圖1 本文設計整體架構

2.2 聚合通信及相關工作

聚合通信操作是高性能計算領域的經典技術，但若將其應用于分布式訓練還有改進的空間。在梯度同步過程中，聚合通信操作AllReduce最常被使用。它將每個節(jié)點上的梯度數據通過指定的操作(如求和或求平均值)進行歸約，并將歸約結果同步到各個節(jié)點。從AllReduce的語義可以看出，其要求所有的節(jié)點都參與該過程。在現(xiàn)有的算法中，Ring AllReduce最常被用于分布式訓練，其可以充分利用每個節(jié)點的網絡帶寬資源，因而也被用于緩解PS架構中server節(jié)點帶寬瓶頸的問題。在此基礎上，一些工作嘗試從結合上層深度神經網絡訓練應用特點的角度合理且高效地使用該操作來減少通信開銷[4,5,10]；還有一些工作則試圖從底層聚合通信算法和網絡傳輸的角度來優(yōu)化該操作[11]，但很少有工作是結合兩者進行優(yōu)化。前面提到的部分聚合操作工作就有類似的優(yōu)化思路，其嘗試將上層的深度學習應用的魯棒性特點與底層聚合通信操作的語義相結合來進行改進，放寬聚合通信操作要求所有節(jié)點全部參加的條件，從而避免等待拖尾節(jié)點，減少了梯度同步的時間，并且不影響模型的收斂?？梢?，減少深度學習應用的需求和聚合通信操作的語義間的不匹配問題對加速分布式訓練至關重要。

2.3 網絡和通信庫

不同的集群通常有不同的網絡環(huán)境，而一個集群內通常也有多套網絡。集群的網絡環(huán)境會從多方面影響分布式訓練中的通信性能，包括帶寬、協(xié)議及拓撲結構等。相較于以太網，InfiniBand網絡具有高帶寬和低時延的特點，這得益于RDMA(Remote Direct Memory Access)可以直接繞過內核，相較TCP開銷更小，可以從網絡傳輸速率上加速梯度同步[14,15]。此外，Bcube比Fat-Tree的拓撲結構更適合AllReduce架構[16]。同時使用集群中的多套網絡設備顯然也能加速通信。鑒于集群環(huán)境的多樣性，充分利用集群的網絡資源和特性對加速通信十分重要。然而不同的網絡設備通常使用各自的接口庫來進行驅動，這些庫的使用方式和調用機制不同，給實現(xiàn)上述改進增加了困難。比如，InfiniBand使用的ibverbs庫在調用實際的通信操作前需要對許多對象進行初始化和配置，如context和QP(Queue Pair)，并且需要通信雙方交換諸如各自地址的控制信息，相比之下，以太網常用的socket庫則僅需創(chuàng)建socket對象并調用個別函數即可建立連接，而對于研究分布式訓練聚合通信操作則不需要考慮這些細節(jié)。

主流深度學習框架并不是直接使用不同網絡的底層庫進行點對點通信，而是調用諸如MPI和NCCL等聚合通信庫，來完成梯度同步操作的，由這些庫負責實現(xiàn)并選擇聚合通信算法并對底層網絡的接口進行封裝。其中，MPI實際上是一個編程接口標準，被廣泛應用于大規(guī)模并行編程中，有多種實現(xiàn)版本，比如OpenMPI和MPICH等。NCCL則是由NVIDIA公司開發(fā)的用于其GPU間交換數據的聚合通信庫。這些庫提供了便捷的使用方式，對底層網絡接口進行封裝，但對分析和優(yōu)化這些聚合操作的細節(jié)來說較為麻煩。

2.4 動機

如前文所述，通過綜合考慮深度學習應用特征和底層網絡特點來進行協(xié)同設計，可以減少其中存在的不匹配現(xiàn)象。然而,在現(xiàn)有的聚合通信庫基礎上分析和優(yōu)化分布式深度學習中的通信操作較為復雜。期望有一個架構簡單、調用邏輯清晰且代碼量較少的聚合通信庫。此外，相關的工作大部分都是針對于某種特定環(huán)境，比如某個框架或是某種網絡，想要移植到其他環(huán)境較為困難。為了滿足實驗的不同需要以及使相關研究便于拓展移植，本文希望聚合通信庫能對上層的深度學習框架和下層的集群網絡有較為廣泛的支持。

3 輕量級聚合通信庫

本文的目標是提供一個可以用于簡化聚合通信操作算法相關研究和實現(xiàn)的輕量級通信庫。為實現(xiàn)這一目標，該庫需要支持主流的框架及網絡，以滿足在其上進行相關研究的需求。本節(jié)將闡述該通信庫的整體架構和關鍵細節(jié)。

3.1 整體架構

首先，不同的分布式訓練框架通常通過其各自的框架和通信庫來進行梯度同步。比如，如圖1所示，MXNet使用基于ZMQ(ZeroMQ)庫的PS架構，而TensorFlow則采用了gRPC庫。為了實現(xiàn)對主流深度學習框架支持這一目標，本文選擇Horovod[17]為通信中間件之一。Horovod對分布式訓練中的參數同步邏輯進行了抽象，并據此針對主流框架分別定義了各自用于通信的接口以支持這些框架，并在底層統(tǒng)一實現(xiàn)。在通信的具體實現(xiàn)上，Horovod選擇了AllReduce作為其參數同步的框架，如圖1所示，底層使用MPI、NCCL和Gloo[18]作為執(zhí)行實際聚合通信操作的通信庫。Gloo是Horovod自帶的一個輕量級的通信庫，為分布式深度學習應用提供了一系列基本的聚合通信算法,使Horovod可以在缺少MPI的環(huán)境中運行。對于Horovod來說，MPI中聚合通信的內部執(zhí)行細節(jié)是透明的，僅需要按照MPI標準進行調用，具體與MPI的實現(xiàn)版本有關，而Gloo則是可以被Horovod感知的。此外，相較各種MPI實現(xiàn)版本，Gloo的架構簡單，代碼量少，更易于進一步分析和改進。而 NCCL則是專用于GPU間通信的，不支持在CPU上進行訓練。因此本文選擇了Gloo作為通信庫的基礎，從而借助Horovod來確保對上層框架的支持。

其次，集群中通常具有多種網絡設施，并且不同集群的網絡環(huán)境也不同。通過支持多種網絡來使得該庫具有更廣泛的適用性，具備充分利用網絡資源的能力。本文在對Gloo進行分析后發(fā)現(xiàn)，其對網絡環(huán)境的支持不如MPI。比如，在InfiniBand網絡模式下，Horovod無法通過使用Gloo完成AllReduce操作。因此需要對Gloo的底層網絡進行擴展。Gloo的擴展方式是使用各自網絡的接口庫來分別實現(xiàn)符合Gloo中點對點通信語義的傳輸層子類，從而為其算法層提供通信傳輸的接口。本文也可以按照這一思路來對其他網絡進行擴展，但這對不同網絡的了解和掌握程度提出了要求，工作量較大，且后續(xù)的擴展也不能得到保障。此外，使用網絡接口庫實現(xiàn)點對點通信的具體實現(xiàn)細節(jié)對研究分布式訓練中的聚合通信來說也不是必要的。本文最終使用UCX(Unified Communication X)[19]作為通信中間件來解決Gloo的網絡擴展問題。UCX整合并封裝了多種底層網絡的接口，提供了一套統(tǒng)一的高性能的點對點通信接口，在屏蔽不同網絡接口庫實現(xiàn)細節(jié)的同時，還通過提供多種語義的操作來體現(xiàn)不同網絡的特性。

本文的設計思路如圖1所示，將Gloo的底層網絡用UCX進行擴展，從而實現(xiàn)支持多種框架和網絡的目標。在這一架構中，Horovod整合并接管了不同深度學習框架的通信部分，UCX用于覆蓋盡可能多的網絡支持，Gloo作為聚合通信的算法邏輯層，為上層應用和下層網絡提供交互條件，充分整合各層信息。與圖1中使用MPI的框架相比，這種設計可以將聚合通信操作的過程加以分解，從而暴露算法的更多細節(jié)。這一改變使得對分布式訓練中諸如AllReduce一類的通信操作的研究和改進變得便捷高效。此外，由于該庫僅提供了分布式訓練必要的聚合通信操作，降低了該庫的復雜性。

3.2 Gloo

如圖1所示，Gloo是Horovod和UCX間的橋梁。本文將Gloo的架構拆分為算法層和傳輸層,其架構如圖2所示。

Figure 2 Architecture of Gloo圖2 Gloo架構

上層算法層主要包括聚合操作的算法邏輯，其基于點對點語義的通信邏輯實現(xiàn)。此外，該層提供了一些分布式訓練必要操作的算法實現(xiàn)，比如AllReduce和Broadcast,用戶可以參考這些算法實現(xiàn)來實現(xiàn)自定義算法。底層的傳輸層主要提供點對點通信的調用接口以及不同網絡環(huán)境下對應的具體實現(xiàn)。Gloo本身僅支持tcp、uv及ibverbs 3種庫，其中對于ibverbs的支持還不完善。此外，對其點對點通信進行擴展需要用網絡對應的接口重寫其基本類和函數，但這對用戶了解網絡設備硬件及熟練使用對應接口提出了更高的要求，并且對于進行聚合操作相關的研究來說引入了更多不必要的細節(jié)。

如圖2所示，Buffer是Gloo中的關鍵類，其既存儲了需要傳輸的數據又是調用通信函數的對象，是算法層和傳輸層的連接。在算法層中，通過Buffer對象來調用對應的點對點通信函數，進而與對端進行數據的發(fā)送和接收，以完成算法邏輯。在傳輸層中，Buffer則通過Pair對象調用實際的通信函數。

為了在實現(xiàn)UCX的擴展時更符合Gloo的語義，本文以Horovod調用AllReduce操作為例，分析了調用流程和類的功能。在實際調用AllReduce之前，Horovod首先通過調用CreateDevice、CreateContext和CreatePair來初始化Gloo。其中，CreateDevice配置網絡設備，包括接口和協(xié)議等；CreateContext設置通信上下文，包括節(jié)點數和節(jié)點編號等；CreatePair創(chuàng)建Pair對象，該對象代表與另一個對等節(jié)點的連接，并配置其通信語義，包括同步或異步等。需要注意的是，一個Context對象可用于創(chuàng)建多個Pair對象，同樣一個Device對象可以創(chuàng)建多個Context對象。然后，Horovod通過調用connect函數來建立全局連接，這樣所有節(jié)點都可以相互通信。初始化工作完成后，Horovod在需要進行通信時可以調用AllReduce操作。這一過程中將創(chuàng)建Buffer對象，用于存儲通信數據并調用send和recv函數。值得注意的是，send和recv是節(jié)點之間通信實際發(fā)生的位置，其語義是非阻塞的，這意味著函數在實際通信完成之前就已返回，因此需要調用waitSend和waitRecv以確保數據已傳輸完成。

3.3 UCX

UCX是一組用于高吞吐量計算的網絡API，為下一代應用程序和系統(tǒng)實現(xiàn)了高性能和可擴展的網絡堆棧。UCX框架包含3個主要組件：UCS、UCT和UCP，其中每個組件都提供一套API，并且可以作為獨立的庫來使用。UCS主要是為UCX內部提供支持;UCT是一個抽象了各種硬件體系結構之間差異的傳輸層，并提供了一個可實現(xiàn)通信協(xié)議的底層API;UCP通過使用UCT層的功能來實現(xiàn)較高級的協(xié)議傳輸功能，包括標簽匹配、流傳輸等多種通信語義。

UCP相較于UCT的優(yōu)勢在于其使用方式簡單且功能多樣。前文提到的Horovod接管了主流深度學習框架的通信部分，并調用Gloo的聚合通信操作來完成，而UCX封裝了各種網絡API，本文將前者的底層傳輸替換為后者，以在支持多框架的基礎下也能對網絡提供更好的支持。

3.4 實現(xiàn)細節(jié)

本文通過將Gloo的傳輸層替換為UCX來實現(xiàn)通信庫的設計。根據對Gloo的分析和UCP的使用方法，本文按照如下方式對其進行整合。首先，分別將ucp_context、ucp_worker和ucp_ep對象分別映射到Device、Context和Pair，這樣做是因為它們之間具有一致的對應關系。然后，在Gloo的相應初始化函數CreateDevice和CreateContext中調用ucp_init和ucp_worker_create，以初始化UCX，但在connect函數中才調用ucp_ep_create，因為其在UCP中的語義代表連接的建立。再次，根據UCP的語義及其使用方法，本文選擇ucp_worker的地址而不是節(jié)點的IP地址作為連接標識，它包括該節(jié)點所有網絡設備的相關信息。最后，在send和recv中分別調用ucp_tag_send_nb和ucp_tag_recv_nb與對等方進行通信。其中，tag的設置綜合了集合操作的標識符和發(fā)送者的rank。此外，本文創(chuàng)建了請求隊列，以維護這些操作的句柄，因為這些函數是非阻塞的，需要通過句柄來確認其傳輸是否完成，以確保waitSend和waitRecv的Buffer中數據的正確性。

Figure 3 Bandwidth of AllReduce in three collective operation library with different data size圖3 AllReduce操作在不同數據大小下的帶寬

4 實驗評估

本節(jié)將對本文的輕量級聚合通信庫中分布式訓練時最常使用的操作AllReduce進行評估 ，并與MPI的實現(xiàn)及原始的Gloo進行比較。此外，本文還分別基于該庫和MPI的實現(xiàn)在MNIST和CIFAR-10數據集上對模型進行了訓練，并測量和對比了其訓練速度。

4.1 實驗環(huán)境

實驗主要在由并行超算云提供的集群上進行，集群具有4個節(jié)點，每個節(jié)點有1個24核的E5-2678的CPU及64 GB的主存，集群網絡環(huán)境為56 Gb/s的InfiniBand和1 000 Mb/s的以太網。此外，集群上安裝有OpenMPI 4.0.1、Horovod 0.19.1、Python 3.7.0和MXNet 1.6.0。

4.2 聚合通信操作

本文在InfiniBand網絡的以太網(ib0)和IB網(mlx5_0)模式下分別測試了所設計的輕量級聚合通信庫、OpenMPI和原始Gloo的AllReduce操作的帶寬，這里帶寬指的是AllReduce處理的數據量和所花時間的比值。從圖3所示的實驗結果可以發(fā)現(xiàn)，在以太網模式下，該庫和OpenMPI及原始的Gloo的性能相近，而在高速網模式下，該庫與OpenMPI有較為明顯的性能差異，在數據小于2 MB時OpenMPI性能更好，而大于2 MB時則該庫更有優(yōu)勢，產生這樣現(xiàn)象的原因可能與底層網絡接口庫的使用方式不同以及MPI對AllReduce操作的算法選擇有關，還需要后續(xù)進一步實驗分析。需要注意的是，原始Gloo僅支持使用以太網模式。

4.3 深度學習應用

本文還在上述這些庫的基礎上在MXNet框架中進行了分布式訓練的實際應用，分別在MNIST和CIFAR-10數據集上訓練了LeNet5和ResNet18 2個模型，測量其在不同批大小下的訓練速度。圖4和圖5分別為MNIST和CIFAR-10的實驗結果。從圖4和圖5可以發(fā)現(xiàn),這些庫訓練速度較為接近，而本文提出的庫相較OpenMPI在大部分情況下具有較小優(yōu)勢。結合圖3中AllReduce性能的對比，較為可能的原因是每次同步的梯度較小，不足以使AllReduce性能呈現(xiàn)出明顯差異。

Figure 4 Speed of training MNIST dataset with different batch sizes圖4 MNIST數據集在不同批大小下的訓練速度

Figure 5 Speed of training CIFAR-10 dataset with different batch sizes圖5 CIFAR-10數據集在不同批大小下的訓練速度

5 結束語

本文提出并實現(xiàn)了一個用于分布式訓練的輕量級聚合通信庫。該庫具有較為簡潔的調用架構和較少的代碼量，并且提供了分布式訓練中基本的聚合操作，對主流的深度學習框架和集群網絡環(huán)境有較好的支持。相較于分布式訓練中傳統(tǒng)的基于MPI等的通信架構，該庫的調用邏輯簡單，能更清晰地展示聚合操作過程的細節(jié)，方便進一步分析和改進分布式訓練中的聚合通信操作實現(xiàn)。并且得益于其對框架和網絡的支持，用戶可以在更多的環(huán)境配置下進行實驗以獲得更廣的影響。此外，本文分別從聚合通信操作和深度學習應用方面對該庫的性能進行了評估，其在實際的深度學習應用方面與常用的OpenMPI性能相近，可以作為分析和研究分布式深度學習中梯度同步的聚合通信庫。