Gloo+:利用在網計算技術加速分布式深度學習訓練*

黃澤彪,董德尊,齊星云

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

1 引言

聚合通信是分布式深度學習訓練中最主要的通信方式,是訓練時用于梯度信息同步的方式之一[1]。其中主要包括Allreduce、Barrier、Broadcast等操作,這些聚合通信操作由于涉及全局,常常對應用程序并行效率產生巨大的影響。為了進一步減少分布式訓練時間,許多研究人員針對聚合通信進行了研究,并提出了多種優化方法,例如騰訊公司2018年提出的層次Ring Allreduce[2]。雖然相關優化方法也很好地提升了聚合通信的效率,但是這些優化方法僅僅是在軟件層面上對聚合通信操作進行了改進,改進后的操作依然需要在網絡中進行多次通信才能完成整體操作,且很容易引起網絡擁塞。而且,當系統規模增大時,通信的計算步驟、計算量以及進程之間的距離都會相應增大,進一步產生較大的通信開銷,且隨著系統規模的增大,這種通信開銷的增加是非常迅速的,使得軟件層面實現的聚合通信可擴展性較差。

在網計算能夠大幅度降低聚合通信時間,極大地提高分布式深度學習訓練的速度。在傳統的基于軟件實現的聚合通信操作中,CPU在發起聚合通信操作后,會阻塞至操作完成。這導致該形式的聚合通信操作難于實現計算與通信的重疊,造成計算資源的浪費;同時,隨著通信數據量的增加,CPU在聚合通信操作中的計算負擔愈加沉重。相對于軟件實現方式,在網計算實現的聚合操作完全卸載到網絡硬件(網卡或交換機),減少了系統噪聲的影響,這進一步加快了聚合操作的執行速度。硬件卸載的方式允許程序以非阻塞的方式執行,有效地實現了計算和通信的重疊,縮短了訓練時間。目前有很多針對在網計算開展的研究[3-8],例如Barefoot公司主導提出的SwitchML交換機系統[3]、Mellanox公司提出的聚合通信網絡卸載協議SHARP(Scalable Hierarchical Aggregation and Reduction Protocol)[4]等。這些研究的實驗結果表明了在網計算有助于緩解分布式應用程序網絡通信問題,進而提升應用程序的整體性能。

聚合通信庫是分布式深度學習訓練中執行通信操作的重要部件。目前常用的聚合通信庫有Gloo[9]、MPI(Message Passing Interface)[10]、NCCL(NVIDIA Collective Communications Library)[11]等。如果能夠在這些通信庫中集成聚合通信在網計算技術,那么將很有可能極大縮短分布式深度學習訓練過程中的通信時間,進一步提升分布式深度學習訓練的整體性能。當前NVIDIA已經在Open MPI[12]和NCCL中集成了SHARP技術,但是Open MPI是個體系結構比較龐大的通信庫,且有很多分布式深度學習訓練不需要的功能,而NCCL雖然是開源的,但是關于聚合通信操作內部的實現并沒有公開,所以不方便研究人員分析和修改聚合通信操作的具體通信細節。雖然Gloo是一套開源的面向分布式深度學習的輕量級聚合通信庫,但是它只實現了軟件層面的聚合操作,并不能利用在網計算技術來加速分布式深度學習訓練。

據我們所知,我們是目前第一個設計并實現了一款輕量級、完全開源并且能夠利用在網計算技術來加速分布式深度學習訓練的通信庫。本文在實現該通信庫時解決了2個挑戰。第一個是內存注冊開銷大。SHARP每次執行聚合通信操作之前都需要進行內存注冊和綁定,即使在同一塊內存被不同聚合操作反復使用到的情境中。這是個很耗費時間的過程,需要降低這個過程所帶來的影響。第二個是功能不適配。SHARP目前實現的聚合操作比較少,有一些必要的聚合操作它沒有實現,例如Allgather。這導致某些操作沒法利用在網計算技術來加速。

本文設計并實現了Gloo+,這是一款在Gloo的基礎上設計并實現的集成SHARP技術的聚合通信庫。Gloo+使用2種方法來分別應對上述挑戰:(1)對同一個內存地址只進行一次注冊和綁定,然后采用哈希表存儲該內存的相關信息,以方便不同的聚合通信操作對其進行操作;(2)根據聚合操作的語義,利用SHARP已經實現的聚合操作來設計其未實現的一些聚合操作。

本文的主要工作包括以下3個方面:

(1) 在Gloo通信庫的基礎上,設計并實現了Gloo+。Gloo+能夠利用SHARP技術加速分布式深度學習訓練,使研究人員能夠便捷地利用在網計算技術;

(2) 評估了Gloo+對分布式深度學習訓練性能的影響;

(3) 詳細分析了基于SHARP的聚合通信操作的優勢和局限性。

2 背景介紹

2.1 聚合通信操作

聚合通信操作是高性能計算領域的經典技術,主要包括Barrier、Broadcast、Allreduce等。在梯度同步過程中,Allreduce操作最常被使用,并且衍生了很多不同的算法來實現Allreduce操作。在現有的算法中,Ring Allreduce[13]是最常被用到的一個算法,如圖1a所示。

在Ring Allreduce的基礎上,衍生了很多改良方法[14-16]。例如,騰訊公司提出的層次Ring Allreduce,其主要思想是對各個進程進行分組,然后采用組內各進程進行reduce操作,各組的主進程進行組間allreduce操作,最后組內進行主進程的broadcast操作,這種方法旨在充分利用組內高帶寬網絡的同時,降低各組之間低帶寬網絡帶來的影響;還有IBM公司提出的BlueConnect算法,其主要思想是考慮了節點間不同機器和交換機(機器內、機器間交換機、上層交換機/路由器)的帶寬不同,從而對進程進行不同的分組,以達到最優的多機通信。還有其他各種采用通信調度等軟件方法來優化Allreduce操作的工作[17,18],都在不同程度上提升了Allreduce的性能。

雖然當前有很多針對聚合通信操作進行優化的工作,但是大多數的工作都是在軟件層面上對聚合通信操作進行優化。經過優化后的聚合通信操作依然需要在網絡中進行多次通信才能完成整體操作,這樣很容易引起網絡擁塞,而且軟件實現的聚合通信可擴展性較差。在網計算將聚合操作卸載到網絡中,如圖1b所示。在網計算可以有效地提升通信效率,提高訓練性能,所以有不少研究人員開展了聚合通信操作在網計算的研究工作。Barefoot公司提出的SwitchML交換機卸載系統,其主要設計思想是使用可編程交換機替代機器學習中傳統的參數服務器,利用交換機的高吞吐率來加速參數更新。還有Mellanox公司開發的聚合通信網絡卸載技術——SHARP,是當前工業界廣泛使用的一項在網計算技術。

Figure 1 Two ways of Allreduce圖1 Allreduce的2種不同的方式

2.2 聚合通信庫

在分布式訓練中,各節點之間的聚合通信操作通常由各種通信庫來執行。最常用到的通信庫有MPI、NCCL和Gloo。

MPI是分布式和并行應用常用的消息傳遞模型的定義,其規定了若干通信原語的接口標準,主要包括點對點通信和聚合通信2類。點對點通信由MPI_Send和MPI_Recv操作組成,用于2個節點間傳遞信息。聚合通信是在點對點通信基礎上進行多節點間的通信操作,有MPI_Allreduce、MPI_Bcast等。MPI標準規定了這些操作接口的參數,所有MPI庫都需要遵循這些接口標準來實現。Open MPI和MPICH[19]是2個常用的MPI庫。

NCCL聚焦于GPU間的數據通信,主要針對GPU上的分布式訓練通信進行優化,以充分發揮PCIe、NvLink和Infiniband等硬件性能,進而實現GPU間的高性能聚合通信接口。NCCL算法豐富度不及MPI庫的,主要提供了Allreduce、Broadcast及點對點發送等操作。NCCL在GPU上的分布式訓練中應用廣泛,被大部分深度學習框架采用。

Gloo是Facebook針對分布式深度學習訓練推出的開源通信庫,為分布式深度學習訓練提供了有用的聚合通信操作。該通信庫向上提供了聚合通信操作的接口,向下提供了對不同網絡的支持,且支持主流的深度學習訓練框架,例如MXNet[20]、TensorFlow[21]和PyTorch[22]等。該庫代碼結構簡潔。相關研究人員可以通過該通信庫簡單高效地實現自己的算法并在更多的環境配置下進行實驗,進而獲得更廣的影響。

在網計算技術可以極大地提升聚合通信操作的性能。但是,目前在分布式深度學習訓練中沒有辦法直接使用在網計算技術。如果能夠在分布式深度學習常用的通信庫中集成在網計算技術,那就可以使得分布式深度學習訓練更方便地使用在網計算技術。通過調研了解到,目前NVIDIA已經在Open MPI和NCCL中集成了SHARP技術。但是,Open MPI是個龐然大物,總共有2 445個文件,包含了288 667行代碼。而且Open MPI有很多深度學習研究中不需要的功能。NCCL在GPU間的通信功能強大,但是其聚合通信操作的內部實現并沒有公開,使得研究人員很難通過它來開展對分布式深度學習聚合通信操作細節的分析。相反,Gloo是一個開源的輕量級聚合通信庫,該通信庫總共只有208個文件,僅包含25 136行代碼。Gloo提供了分布式深度學習訓練中常用的聚合通信算法,沒有其他多余的復雜功能,其體系不會很龐大,而且整體架構簡潔,便于分析和改動。能夠在Gloo中集成SHARP技術來實現該通信庫聚合通信操作的在網計算技術,對分布式深度學習的研究來說意義重大。

Figure 2 Topology structure of SHARP tree圖2 SHARP樹形拓撲結構

2.3 SHARP

SHARP是一種允許將聚合通信操作卸載到網絡中的技術。SHARP在物理拓撲的基礎上建立邏輯聚合通信樹形結構,其樹形結構圖如圖2所示。高層次通信庫中的進程子集用于形成SHARP組,該組用于定義SHARP樹中的末端節點,這些節點輸入要歸約并向上傳輸的數據。SHARP樹中的非葉子節點是聚合節點,聚合節點負責執行聚合通信操作。當數據到達SHARP的根節點時,便開始進行分發操作,將聚合通信操作完成的數據分發給SHARP組中的各個節點。其設計的網絡接口芯片與互連交換芯片硬件都具備數據聚合處理能力,共同構成邏輯樹中的聚合結點。使用SHARP的好處是可以釋放CPU資源供應用程序使用,消息通信效率高、延遲低,而且受到系統噪聲的影響極少。該技術目前已經引入到Mellanox公司開發的交換機上,在交換機芯片中集成了計算引擎單元,可以支持16位、32位及64位定點計算或浮點計算,可以支持求和、求最小值、求最大值、求與、求或及異或等計算,可以支持Reduce、Allreduce等操作。

3 Gloo+通信庫

本節將描述Gloo+的設計與實現。本文主要是在Gloo的架構上進行改動,將SHARP集成到Gloo中以實現基于SHARP的Allreduce、Reduce、Allgather操作,使其能夠利用SHARP技術進行節點間的聚合通信操作。

Gloo的代碼結構主要分為Transport、Context、Collective Operations 3個模塊,其中Transport模塊主要負責提供數據通信功能,如連接建立、數據的發送和接收等;Context模塊主要負責管理全局通信的環境,如節點的rank、size、address等信息以及建立全局連接等功能;Collective Operations模塊主要負責提供聚合通信操作,如Allreduce、Allgather、Broadcast等。3個模塊之間的關系主要是Collective Operations模塊中使用Context模塊來獲取全局的通信能力,在Context模塊中使用Transport模塊來完成具體的通信操作。本文主要在Context模塊中實現了SHARP通信域的構建,在Collective Operations模塊中實現了SHARP聚合通信操作的具體執行算法。圖3展示了Gloo+的系統結構。

Figure 3 System structure of Gloo+圖3 Gloo+整體架構

3.1 基于SHARP的Allreduce

Gloo中對于Allreduce操作的實現主要采用了ring和bcube這2個算法。本文在Gloo+中設計的基于SHARP實現的Allreduce操作是完全獨立于Gloo本身Allreduce操作實現細節的個體。其主要分為3個步驟:首先,利用Gloo的contex來對SHARP通信域及操作進行初始化,構建一個SHARP通信域;然后,進行發送緩沖區和接收緩存區的注冊綁定,使得SHARP守護進程可以對相關數據進行操作;最后,根據前2個步驟提供的一些參數調用SHARP提供的Allreduce接口,利用聚合通信在網計算技術進行Allreduce操作。

本文將Gloo+集成到Horovod中。但是,當在Horovod框架上用深度學習訓練框架進行模型訓練時,發現每一次Allreduce操作的執行都需要預先進行發送緩存區和接收緩存區的注冊與綁定,這是一項很消耗時間的操作。同時還發現,在深度學習模型訓練的過程中,訓練啟動時已經基本分配好了相應的一些內存空間來供各種數據使用,有可能多組不同的數據在不同的時間點上共享同一個內存空間,所以本文對Allreduce操作的內存注冊進行了優化。如果需要進行操作的數據對應的內存地址還沒有注冊,那么就按正常流程對該地址進行注冊和綁定,此時會獲得一個對應的內存句柄,這個內存句柄是下文提到的“reduce_spec”結構的其中一個參數。本文會將首次進行內存注冊的內存地址對應的內存句柄存放到一個哈希表中。如果需要進行操作的數據對應的內存地址已經注冊過,那么哈希表中將會有該內存地址對應的內存句柄,就只需要直接從哈希表中取出該內存句柄然后進行下一步的操作。這樣就避免了對同一個內存地址進行多次注冊綁定,大大降低了時間開銷。

在這個過程中,本文會用到SHARP提供的一個數據結構——reduce_spec。該結構定義了SHARP的聚合操作相關輸入參數,包括數據類型、操作方式、聚合模式、發送緩沖區、接收緩沖區等詳細信息。可以通過對reduce_spec中的一些參數進行設置,來滿足執行相關聚合操作所需要的條件,進一步利用它來完成聚合操作。在本文中,Allreduce、Reduce以及Allgather的實現都需要用到reduce_spec。

3.2 基于SHARP的Reduce

在對SHARP的功能進行測試和分析的時候發現一個令人不解的問題,即實現的Reduce操作是一個有缺陷的操作。它的Reduce操作在數據量的個數少于16 KB時,會出現錯誤并停止執行任務,而在數據量個數大于16 KB時,就可以正常開展作業的執行。因為SHARP是不開源的,所以對于其內部的具體實現細節也無從得知,無法分析產生該問題的原因。

當利用SHARP的Reduce實現Gloo+的Reduce操作時,針對上面發現的問題,本文進行了一些優化,使得Gloo+中的Reduce在數據量個數少于16 KB時,也能正常使用SHARP的在網計算技術。因為Allreduce和Reduce的語義很相似,只不過Allreduce是將歸約后的結果分發給所有節點,而Reduce是將歸約后的結果給指定的節點,所以本文設計的主要思想是,對于個數少于16 KB的數據量,對其采用SHARP的Allreduce功能。具體實現如下:(1)給所有進程分配一塊內存作為發送緩存區并進行注冊和綁定;(2)對需要執行操作的數據量進行判斷,如果數據量少于16 KB,給所有進程分配一塊內存作為接收緩存區并進行注冊和綁定;當數據量大于或等于16 KB、且當前進程是指定的root進程時,給當前進程分配接收緩沖區;(3)調用SHARP提供的接口進行聚合通信操作。如果數據量少于16 KB,則調用Allreduce接口,否則調用Reduce接口;(4)在數據量少于16 KB時,除了指定的根進程,其他進程均將接收緩沖區的數據丟棄。

3.3 基于SHARP的Allgather

本文在對SHARP進行分析測試時發現它只實現了Allreduce、Reduce、Broadcast和Barrier 4種操作。與此同時,在對Gloo實現的聚合操作進行分析時發現它實現了一些SHARP沒有實現的操作,比如Allgather等。因此,考慮到可以借助SHARP的Allreduce操作來實現Gloo+中基于SHARP在網計算技術的Allgather操作,進一步使得Gloo+能夠利用SHARP的在網計算技術加速Allgather操作。

眾所周知,Allgather的語義其實跟Allreduce的語義很像,同樣都是對所有進程的數據進行收集然后分發回去。唯一不同的是,Allreduce會對從各個進程收集回來的數據進行歸約操作,例如求和、求均值等。而Allgather操作則不會對各個進程的數據進行歸約操作,而只是收集起來,然后將收集到的所有數據再分發給各個進程,使得每一個進程都擁有一份所有進程的數據。

鑒于上面的分析,本文針對Allgather操作的主要設計思想是:根據進程的數量給每個進程開辟2個具有一定容量的內存作為發送緩沖區和接收緩沖區,然后根據進程號來確定所要用到的數據在發送緩沖區中的位置,接著進行Allreduce操作。如圖4所示,首先,假設有4個進程,需要執行Allgather操作的數據量是1個,那么就分配兩塊能容納4個數據的內存分別作為發送緩沖區和接收緩沖區;然后,根據每個進程的序號,來確定數據在發送緩沖區中的存放位置,比如對于rank2,它的數據從內存中的偏移量為2的位置開始存放,占用一個數據量大小的長度,其余的位置全部用0填充;接著,調用SHARP的Allreduce接口執行Allreduce操作,所有進程的數據在交換機中執行求和操作,這樣每個進程提供的數據跟其他進程提供的0進行求和,最終便會得到一個存放著所有進程的數據的大數據塊;最后,交換機將求和后的結果分發給各個進程,那么每個進程都會擁有一份所有進程數據,這樣就完成了Allgather的整個操作過程。

Figure 4 Allgather operation using SHARP’s Allreduce API圖4 使用SHARP的Allreduce接口實現Allgather操作

4 實驗評估

本節對Gloo+中實現的Allreduce、Reduce、Allgather進行測試評估,并與Gloo和MPI進行比較;還分別將Gloo+、Gloo和MPI應用于分布式深度學習訓練中進行神經網絡模型訓練,測量了其吞吐量。

4.1 實驗環境

實驗主要在5臺Intel?Xeon?Gold 6230R服務器和1臺NVIDIA Quantum 200 Gb/s InfiniBand交換機組成的集群上進行。每臺服務器都包含一個雙插槽主板,每個插槽都有一個運行頻率為2.10 GHz的26核處理器。該集群網絡包含Mellanox InfiniBand HDR適配器,在計算節點之間提供100 Gb/s的帶寬。

4.2 基準測試

本節描述了本文涉及的Allreduce、Reduce、Allgather操作在不同的消息大小上的實驗結果,并對實驗結果進行了分析。

圖5展示了在不同消息大小上對Allreduce、Reduce、Allgather操作的性能測試結果,包括消息大小對Gloo+、Gloo和MPI通信延遲的影響。其中,X坐標軸代表消息大小(以字節為單位),Y坐標軸表示延遲(以微秒為單位),4條折線分別代表Gloo、Gloo+及MPI在以太網模式(網卡接口為ib)和IB網模式(網卡接口為mlx5)的情況。

Figure 5 Latency of collective operations with Gloo+, Gloo and MPI across various message sizes圖5 使用Gloo+、Gloo和MPI 對各種消息大小的聚合操作的延遲

在Allreduce操作中,Gloo+相比于Gloo,當消息大小小于64 KB時,Gloo+加速比能達到100以上;而當消息大小大于64 KB時,其加速比也能達到10以上;Gloo+跟MPI進行比較,相比于MPI的以太網模式,Gloo+的加速比在10～57;而相比于MPI的IB網模式,Gloo+的加速比在10以內。在Reduce操作中,Gloo+相比于Gloo,當消息大小小于2 KB時,Gloo+的加速比能達到100左右;在消息大小增大到2 KB以上時,其加速比能達到十幾,甚至幾十;Gloo+跟MPI進行比較,相比于MPI的以太網模式,Gloo+的加速比在2～16;而相比于MPI的IB網模式,當消息大小小于或等于8 KB時,Gloo+的性能比MPI的還差,而當消息大小大于8 KB時,其性能相對較好,相比于MPI加速比在0～6。通過分析可知,在Allreduce和Reduce這2個聚合操作中,Gloo+相對于Gloo的性能提升都比較大,加速比從10到100以上;而跟MPI相比較,在以太網模式下,Gloo+性能也比MPI的好,不過在IB網模式中的消息大小比較小的情況下,Gloo+的性能比MPI的差,當消息大小較大時,Gloo+則更有優勢。

在Allgather操作中,SHARP同樣也提升了該操作的性能。因為本文是采用SHARP的Allreduce接口來實現基于SHARP的Allgather操作,所以每次Allgather操作所傳輸的消息大小比實際的有效消息大小要大,這會導致整體通信性能偏低。但是,通過跟Gloo比較發現,即使基于SHARP的Allreduce實現的Allgather通信性能偏低,還是比Gloo中實現的Allgather的通信效率高。當消息大小小于128 KB時,基于SHARP的Allreduce接口實現的Allgather操作相對于Gloo中實現的Allgather操作的加速比能達到10～50,可以看到這樣的加速效果還是很好的。當消息大小大于128 KB時,Gloo+所帶來的加速比也能維持在7左右。將Gloo+與MPI進行比較,在以太網模式中,Gloo+性能比MPI的好,加速比在5～24。而在IB網模式中,Gloo+跟MPI的性能比較接近。通過分析可知,在Allgather操作中,將Gloo+與Gloo進行比較,在消息大小比較小的時候,Gloo+能達到十幾甚至幾十的加速比。隨著消息大小的增大,其加速比會呈現一個下降的趨勢,但最終都會穩定在7左右。而與MPI進行比較,在以太網模式下,Gloo+僅在消息大小較大的情況下性能提升較明顯,而在IB網模式下,其性能與MPI的接近。

4.3 深度學習應用測試

本節給出Gloo+和Gloo在分布式深度學習訓練應用中的實驗結果。本文將Gloo+和Gloo集成到Horovod框架中,在該框架上采用MXNet深度學習訓練框架來進行模型訓練。本文采用ImageNet數據集分別對VGG19、AlexNet神經網絡模型進行訓練,對每個神經網絡模型都進行了16,32,64和128這4種批大小的獨立訓練。圖6和圖7分別展示了2個神經網絡模型的實驗結果。

Figure 6 Throughput of training VGG-19 with different batch size圖6 VGG-19在不同批大小下的吞吐量

Figure 7 Throughput of training Alexnet with different batch size圖7 Alexnet在不同批大小下的吞吐量

圖6展示了VGG19神經網絡模型訓練的實驗結果。該模型的參數量高達1.4億個,模型大小為534 MB,其在訓練過程中的通信量很大。從圖6可以看到,在不同批大小中,Gloo+表現都很優異。在批大小為16和32時,Gloo+相對于Gloo能夠達到1.1以上的加速比,而在批大小為64和128時,則分別能夠有0.7和0.4的加速比。相比于以太網模式下的MPI,在批大小為16和32的情況下,Gloo+能夠達到1左右的加速比,在批大小為64和128的情況下,則有0.7的加速比。相比于IB網模式下的MPI,在4種批大小中,Gloo+的加速比依次為0.54,0.49,0.47和0.35。

圖7展示了AlexNet神經網絡模型訓練的實驗結果。該模型的參數量有6 200萬個,模型大小為237 MB,該模型的大小相對VGG19來說小一些,但對于其他神經網絡模型來說也是一個比較大的模型,其通信開銷在整個訓練的過程中占比也不小。從圖7可以看到,跟VGG19實驗結果表現的一樣,在不同批大小中,Gloo+表現依然很優異。在4種批大小中相對于Gloo和以太網模式下的MPI均能夠達到1.1以上的加速比,而相對于IB網模式下的MPI其加速比依次為0.26,0.34,0.52和0.36。

通過對以上實驗結果的分析發現,當深度學習神經網絡模型參數數量較多時,其通信量則相應地會比較大,那么Gloo+就能夠使得神經網絡模型訓練性能得到很大的提升。當批大小越小時,訓練完整個數據集所需要的迭代次數越多,相應的通信頻率就越高,那么此時Gloo+也能表現出很好的訓練性能,極大地減少模型訓練的通信開銷。總而言之,Gloo+不僅在基準測試中表現優異,在分布式深度學習模型訓練的應用中也展現出了很好的效果。

5 結束語

本文設計并實現了基于SHARP的聚合通信庫Gloo+,使分布式深度學習訓練能夠利用網絡的計算能力來加速聚合通信操作。本文評估了Gloo和Gloo+中聚合通信操作的性能,并且將Gloo+集成到Horovod中,然后使用MXNet在Horovod中訓練神經網絡模型,從而評估Gloo+在實際應用場景中的真實效果。

本文對Gloo+的實驗評估結果表明,不管是在Allreduce、Reduce和Allgather等基準測試中,還是在分布式深度學習訓練的實際應用中,Gloo+的表現都極其優秀。