王笑雨,董德尊

(國防科技大學(xué)計(jì)算機(jī)學(xué)院,湖南 長(zhǎng)沙 410073)

1 引言

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

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

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

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

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

2 背景介紹

2.1 分布式訓(xùn)練和框架

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

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

Figure 1 Architecture of the proposed design圖1 本文設(shè)計(jì)整體架構(gòu)

2.2 聚合通信及相關(guān)工作

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

2.3 網(wǎng)絡(luò)和通信庫

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

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

2.4 動(dòng)機(jī)

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

3 輕量級(jí)聚合通信庫

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

3.1 整體架構(gòu)

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

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

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

3.2 Gloo

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

Figure 2 Architecture of Gloo圖2 Gloo架構(gòu)

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

如圖2所示，Buffer是Gloo中的關(guān)鍵類，其既存儲(chǔ)了需要傳輸?shù)臄?shù)據(jù)又是調(diào)用通信函數(shù)的對(duì)象，是算法層和傳輸層的連接。在算法層中，通過Buffer對(duì)象來調(diào)用對(duì)應(yīng)的點(diǎn)對(duì)點(diǎn)通信函數(shù)，進(jìn)而與對(duì)端進(jìn)行數(shù)據(jù)的發(fā)送和接收，以完成算法邏輯。在傳輸層中，Buffer則通過Pair對(duì)象調(diào)用實(shí)際的通信函數(shù)。

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

3.3 UCX

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

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

3.4 實(shí)現(xiàn)細(xì)節(jié)

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

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

4 實(shí)驗(yàn)評(píng)估

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

4.1 實(shí)驗(yàn)環(huán)境

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

4.2 聚合通信操作

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

4.3 深度學(xué)習(xí)應(yīng)用

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

Figure 4 Speed of training MNIST dataset with different batch sizes圖4 MNIST數(shù)據(jù)集在不同批大小下的訓(xùn)練速度

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

5 結(jié)束語

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