基于Ring All-Reduce的高擴展性分布式機器學習架構

黃純悅 楊宇翔

摘要：如今，機器學習廣泛應用于各個行業，然而隨著當下各種應用場景的數據量的增大，分布式機器學習幾乎成為唯一的選擇。因此，各個設備之間的數據通訊的優化十分重要。在參數服務器架構中，參數同步通信量大，參數服務器節點的帶寬會成為瓶頸;而在基于Ring All-Reduce的框架下，通信時間受限于環上最慢的連接，當環中GPU節點數變多的時候，會導致延遲變大。該文提出一種基于Ring All-Reduce的分層架構，將計算節點按算力大小分成多個小組，組內使用Ring All-Reduce算法進行同步并行，小組間使用參數服務器架構實現異步并行，保證模型收斂的條件下，兼顧各個節點的負載均衡。

關鍵詞：分布式機器學習;聯邦學習;分層Ring All-Reduce

中圖分類號：TP18? ? ? 文獻標識碼：A

文章編號：1009-3044（2022）06-0054-03

開放科學（資源服務）標識碼（OSID）：

1 概述

機器學習是當下的熱門學科，神經網絡作為機器學習的核心技術[1]，廣泛應用于如自動駕駛、語音識別[2]、文本理解[3-4]、圖像分類[5-6]等領域。然而在許多互聯網應用場景下，經常會涉及量級很大的計算數據，想要利用單個設備完成機器學習模型的訓練是非常困難的，因此分布式機器學習[7]是十分必要的。所以，對設備之間數據通訊的優化具有十分重要的意義。

參數服務器架構[8-9]是最開始被提出來的同步參數的框架，但該架構會受到網絡帶寬的限制，擴展性較差。

Ring All-Reduce[10-12]架構有效減少了通訊延時和時間開銷，但是在環狀結構中隨著環上節點數目增大，各個設備算力的差異會產生大量的等待時間，從而大大增加時間開銷。各種基于Ring All-Reduce架構諸如Hierarchical All-reduce[13]等改進架構，采用了分層的思想，減少了一個組內計算節點的數目，以減少延遲，但組間仍然采用Ring All-Reduce的思想，沒有改變同步更新的特性，當有故障節點或算力很低的節點時，仍然會拖慢整體的計算速度。

因此，在本文中會提出一種新的框架：PS-enhanced Hierarchical Ring-based All-Reduce（PS-HRA），將上述兩種架構結合，用以進一步減少通信時間開銷。同時我們也會分析新算法的時間開銷并且與前述算法相比較，以證明該算法的優越性。

接下來論文的框架是：第二部分總結了我們研究的相關工作，介紹了幾種傳統架構：PS架構、Ring All-reduce架構、Hierarchical All-Reduce架構;第三部分介紹了算法PS-HRA的設計;第四部分計算了PS-HRA算法的時間代價并與傳統架構進行了性能比較;第五部分總結了全文并闡述了未來的工作。

2 相關工作

2.1 Parameter Sever架構

Parameter Server（PS）架構是一種分布式機器學習中心化的同步參數的架構，由服務器端和客戶端組成。服務器端用來接收客戶端的梯度，進行梯度平均。客戶端用來向服務器端傳送梯度以及利用從服務器端接收到的最新梯度以更新模型。其整體架構如圖1所示。

PS架構工作過程如下：每個客戶端用自己的局部數據訓練得到梯度后，將梯度傳送給參數服務器端。參數服務器端接收各個客戶端的梯度后，進行梯度平均，再傳送給各個客戶端。

設有[N]個計算節點，[C]為server計算每字節數據的耗時，[S]為數據塊大小，[B]為帶寬， [a]為兩個通信點間的延時。

則時間開銷為：

[2*a+NSB+N*S*C]

此種架構的局限性在于參數同步通信量大，由于會受到網絡帶寬的限制，擴展性差。

2.2 Ring All-Reduce

Ring All-Reduce算法中，所有的GPU節點只會與其相鄰節點進行數據通信。該算法總共有兩步： scatter-reduce和all-gather。在第一步中，每個GPU節點依次計算特定的數據塊得到梯度并和上一個GPU發送的梯度進行累加，之后發送給下一個GPU節點，假設有[N]個節點，共進行[N-1]次。在這一步結束后，每個GPU節點都會包含一個最終梯度塊;在第二步中，相鄰節點間進行梯度塊的交換，使得每個節點都有最終結果的全部梯度。其架構如圖2所示。

假設有[N]個節點，[S]為數據塊大小，[B]為帶寬，[a]為兩個通信點間的延時，[C]為每字節數據的計算耗時。

則時間開銷為：

[2*（N-1）*[a+SNB] + （N-1）*[（SN）*C]]

環上的最慢連接會大大增加該算法的通信時間，另外當GPU數量不斷增加，延遲也會不斷增大。

2.3 Hierarchical All-Reduce

Hierarchical All-Reduce是基于Ring All-Reduce進行優化的一種算法，該算法的過程如圖3所示。

Hierarchical All-Reduce算法按三步進行：第1步，Intragroup All-Reduce，將所有節點分為[k]組，每組內部進行Ring All-Reduce;第2步，Intergroup All-Reduce，在每個組中選擇一個節點，這些節點進行Ring All-Reduce。在第三個階段Intragroup Broadcast，將更新后的參數傳遞到所有節點。

假設有[N]個節點，被分為[k]組，設[S]為數據塊大小， [B]為帶寬，[a]為兩個通信點間的延時。

則需要的總時間為：

[2*Nk-1*a+kSNB+2*k-1*a+SkB+Nk-1*kSN*C+k-1*Sk*C+a+SB]

Hierarchical All-Reduce模型采用了分組策略，一定程度上改善了Ring All-Reduce節點數增加導致通信時延過大的局限性，并且通信時間依然受到環上最慢連接的限制。

3 PS-HRA：PS-enhanced Hierarchical Ring-based All-Reduce

3.1 PS-HRA模型架構

PS-HRA算法綜合了去中心化的Ring All-Reduce算法和中心化的PS框架，形成了一個分層的通訊優化框架。它由若干個去中心的節點的環和一個中心節點組成?？偟膩碚f，相比于傳統的參數服務器架構，中心節點在我們的框架下控制的是各個環而不是單個節點，因此采用我們的通訊框架可以有效地解決中心節點帶寬限制的問題，即在僅有一個中心節點的情況下，我們的框架可以擴展更多的計算節點，并保證不會因為帶寬問題使傳輸數據速率限制下降。同時環間采用異步PS架構，在環內計算節點算力相近的條件下，可以改善環間算力差異帶來的同步等待問題。

3.2 分組策略

在PS-HRA架構中，中心節點的作用主要有兩個，一是作為參數服務器，另一個是工作節點調度器。初啟階段，中心節點根據算力大小將工作節點分組。具體步驟是在中心節點設置一個隊列，初啟時，各個計算節點訓練部分測試樣本，當訓練完畢時，向中心節點發送完成標識，進入隊列。假設有[N]個節點，被分為[k]組當隊列中計算節點大于等于[N/k]時，將前[N/k]個計算節點出隊列并形成一組進行Ring All-Reduce訓練。在環內每個節點可以利用自己的數據集訓練出模型的一部分參數，然后在all-gather這一步中相鄰節點通過交換參數，使得每個節點獲得最終的全部參數。

算法1. PS-HRA分組算法

輸入：節點[id] ， 節點數目[N]， 分組數目[k]

輸出：環[ring]

1. [q=initqueue（）;] //初始化隊列

2. [worknode]←[testspeeddemo]; //測試數據發送到工作節點

3. [loop：]

4. [ifworknode[id].status= =OK q.enqueueid; ]//將完成計算的節點加入隊列

5. [ifq.size≥N/k {ring←q1…N/k;}] //前[N/k]的計算節點形成環

3.3 環間異步PS架構

Hierarchical All-Reduce算法在環間使用的Ring All-reduce是一個同步算法，沒有考慮算力大小差異的問題，無法解決同步開銷大的問題，通信時間依舊會受到最慢的連接的限制。PS-HRA在各個環之間采用異步PS架構，可以任意選取其中一個計算節點代表他所在的環和服務器進行通訊。PS架構使用的是數據并行[14]技術，即每個環獲得的數據并不一樣，每個環把通過訓練本地數據的獲得模型參數傳給中心節點，中心節點將接收來的參數進行聚合更新從而獲得最新參數并傳送給各個環。由于設置分組策略使得組內節點算力相近，組間節點算力差異大，如果采用同步的PS架構，雖然可以保證模型收斂，但是同步開銷比較大，因此采用異步PS架構。

在異步PS架構中，參數服務器收到了一個環傳來的梯度之后，不需要收集全部環的參數再進行更新，而是直接去更新參數，得到新的參數值傳給代表節點。再由代表節點在環內進行廣播，以進行下一輪的訓練。這樣不會受到算力最慢的工作節點的限制，提高了訓練速度，可以很好地解決同步開銷的問題。PS-HRA整體流程如下：

1） 執行PS-HRA分組算法，將算力接近的[N]個節點均勻分成[k]個小組。

2） 小組內并行scatter-reduce算法，執行結束之后每個計算節點具有[k/N]個數據塊的環內完整參數。

3） 小組內并行all-gather算法，執行結束后每個計算節點具有環內數據塊的全部參數。

4） 環內隨機選出一個計算節點作為代表節點，向參數服務器發送參數。

5） 參數服務器接收參數直接對參數聚合更新，無須等待其他環。

6） 代表節點接收參數。

7） 代表節點在環內進行參數廣播。

8） 各節點更新模型，進行下一輪的訓練。

對于傳統的參數服務器架構，假設有16個計算節點，中心節點最多可能面臨著16個計算節點同時和它進行通訊;假設有4個環，而PS-HRA的通訊模型最多只需要四個計算節點和中心節點通訊，這大大減少了帶寬瓶頸。

4 PS-HRA時間代價估計與性能分析

4.1 時間代價估計

假設有[N]個節點，被分為[k]組，設[S]為總的數據量，[C]為每字節數據的計算耗時，[B]為帶寬，[a]為兩個通信點間的延時。

組內scatter-reduce所需時間為 [Nk-1*a+kSNB+kSN*C]

組內all-gather 所需時間為 [Nk-1*a+kSNB]

組間使用異步PS架構所需時間為 [2*a+SB+S*C]

每組代表節點將更新后的參數傳給組內其余節點所需時間為：

[a+SB]

通過計算可得，需要的總時間為：

[2*Nk-1*a+kSNB+Nk-1*kSN*C+3*a+SB+S*C]

4.2 性能分析

為了證明PS-HRA算法性能的優越性，需要在傳輸時間延遲或節點數量變化的情況下，對比前述四種算法的通信時間。

我們模擬了兩種情形下前述四種算法的通信時間并繪制了折線圖用以比較，如圖5、圖6所示。

傳輸延遲包括算力差的節點帶來的同步開銷，所以由圖5可知，傳輸延遲越大，Ring All-Reduce的通信時間越長;然而PS架構的通信時間變化很小，其性能受傳輸延遲影響小。另外，不難發現PS-HRA的通信時間最少且隨著傳輸延遲的增大，其性能的優越性更加顯著。

由圖6可知，隨著節點數的增加，PS-HRA算法的通信時間基本不變，而其余三種算法的通信時間會隨著節點數的增加而不斷變大。當節點數較大時，PS-HRA的通信時間最少。

5 結束語

PS架構會受到帶寬的限制，而Ring All-Reduce通信時間受限于環上最慢的連接并且通信延遲會隨著環中節點數增加變大。在本文中我們將上述兩種算法相結合，進一步優化了算法的性能。在這篇文章中，首先我們介紹了相關工作和新算法PS-HRA的設計。然后，我們分析了PS-HRA算法的性能和時間開銷，并且比較證明了我們算法性能的優越性。

在未來的研究中，我們會在更多、更復雜的情況下進一步驗證我們通信策略的有效性。另外，我們會把我們的通訊策略應用到更多機器學習的場景中。

參考文獻：

[1] 杜海舟，黃晟.分布式機器學習中的通信機制研究綜述[J].上海電力大學學報，2021，37（5）：496-500，511.

[2] Deng L，Li J Y，Huang J T，et al.Recent advances in deep learning for speech research at Microsoft//2013 IEEE International Conference on Acoustics， Speech and Signal Processing. IEEE， 2013： 8604-8608.

[3] Turian J，Ratinov L，Bengio Y.Word representations：a simple and general method for semi-supervised learning[C]//ACL '10：Proceedings of the 48th Annual Meeting of the Association for Computational Linguistics.2010：384-394.

[4] Liang X D，Hu Z T，Zhang H，et al.Recurrent topic-transition GAN for visual paragraph generation[C]//2017 IEEE International Conference on Computer Vision.Venice，Italy：IEEE，2017：3382-3391.

[5] Yan Z C，Piramuthu R，Jagadeesh V，et al.Hierarchical deep convolutional neural network for image classification：US10387773[P].2019-08-20.

[6] Yan Z C，Zhang H，Wang B Y，et al. Automatic photo adjustment using deep neural networks[J].ACM Transactions on Graphics （TOG）， 2016， 35（2）：1-15.

[7] Chen T， Li M， Li Y， et al.Mxnet：A flexible and efficient machine learning library for heterogeneous distributed systems[J].arXiv preprint arXiv：1512.01274， 2015.

[8] Li M， Zhou L， Yang Z， et al. Parameter server for distributed machine learning[C]//Big Learning NIPS Workshop，2013， 6： 2.

[9] Li M，Andersen D G，Park J W，et al.Scaling distributed machine learning with the parameter server[C]//11th {USENIX} Symposium on Operating Systems Design and Implementation （{OSDI} 14）. 2014： 583-598.

[10] Tokui S，Okuta R，Akiba T，et al.Chainer： A deep learning framework for accelerating the research cycle[C]//Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining.2019：2002-2011.

[11] Tokui S，Oono K，Hido S，et al.Chainer： a next-generation open source framework for deep learning[C]//Proceedings of workshop on machine learning systems （LearningSys） in the twenty-ninth annual conference on neural information processing systems （NIPS）.2015， 5：1-6.

[12] Sergeev A，Del Balso M.Horovod： fast and easy distributed deep learning in TensorFlow[J].arXiv preprint arXiv：1802.05799，2018.

[13] Jiang Y，Gu H，Lu Y，et al.2D-HRA：Two-Dimensional Hierarchical Ring-Based All-Reduce Algorithm in Large-Scale Distributed Machine Learning[J]. IEEE Access，2020，8：183488-183494.

[14] 李曉明.數據并行計算：概念、模型與系統[J].計算機科學，2000，27（6）：1-5.

【通聯編輯：謝媛媛】