基于Ring All-Reduce的高擴(kuò)展性分布式機(jī)器學(xué)習(xí)架構(gòu)

黃純悅 楊宇翔

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

關(guān)鍵詞：分布式機(jī)器學(xué)習(xí);聯(lián)邦學(xué)習(xí);分層Ring All-Reduce

中圖分類號：TP18? ? ? 文獻(xiàn)標(biāo)識碼：A

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

開放科學(xué)（資源服務(wù)）標(biāo)識碼（OSID）：

1 概述

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

參數(shù)服務(wù)器架構(gòu)[8-9]是最開始被提出來的同步參數(shù)的框架，但該架構(gòu)會(huì)受到網(wǎng)絡(luò)帶寬的限制，擴(kuò)展性較差。

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

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

接下來論文的框架是：第二部分總結(jié)了我們研究的相關(guān)工作，介紹了幾種傳統(tǒng)架構(gòu)：PS架構(gòu)、Ring All-reduce架構(gòu)、Hierarchical All-Reduce架構(gòu);第三部分介紹了算法PS-HRA的設(shè)計(jì);第四部分計(jì)算了PS-HRA算法的時(shí)間代價(jià)并與傳統(tǒng)架構(gòu)進(jìn)行了性能比較;第五部分總結(jié)了全文并闡述了未來的工作。

2 相關(guān)工作

2.1 Parameter Sever架構(gòu)

Parameter Server（PS）架構(gòu)是一種分布式機(jī)器學(xué)習(xí)中心化的同步參數(shù)的架構(gòu)，由服務(wù)器端和客戶端組成。服務(wù)器端用來接收客戶端的梯度，進(jìn)行梯度平均?？蛻舳擞脕硐蚍?wù)器端傳送梯度以及利用從服務(wù)器端接收到的最新梯度以更新模型。其整體架構(gòu)如圖1所示。

PS架構(gòu)工作過程如下：每個(gè)客戶端用自己的局部數(shù)據(jù)訓(xùn)練得到梯度后，將梯度傳送給參數(shù)服務(wù)器端。參數(shù)服務(wù)器端接收各個(gè)客戶端的梯度后，進(jìn)行梯度平均，再傳送給各個(gè)客戶端。

設(shè)有[N]個(gè)計(jì)算節(jié)點(diǎn)，[C]為server計(jì)算每字節(jié)數(shù)據(jù)的耗時(shí)，[S]為數(shù)據(jù)塊大小，[B]為帶寬， [a]為兩個(gè)通信點(diǎn)間的延時(shí)。

則時(shí)間開銷為：

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

此種架構(gòu)的局限性在于參數(shù)同步通信量大，由于會(huì)受到網(wǎng)絡(luò)帶寬的限制，擴(kuò)展性差。

2.2 Ring All-Reduce

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

假設(shè)有[N]個(gè)節(jié)點(diǎn)，[S]為數(shù)據(jù)塊大小，[B]為帶寬，[a]為兩個(gè)通信點(diǎn)間的延時(shí)，[C]為每字節(jié)數(shù)據(jù)的計(jì)算耗時(shí)。

則時(shí)間開銷為：

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

環(huán)上的最慢連接會(huì)大大增加該算法的通信時(shí)間，另外當(dāng)GPU數(shù)量不斷增加，延遲也會(huì)不斷增大。

2.3 Hierarchical All-Reduce

Hierarchical All-Reduce是基于Ring All-Reduce進(jìn)行優(yōu)化的一種算法，該算法的過程如圖3所示。

Hierarchical All-Reduce算法按三步進(jìn)行：第1步，Intragroup All-Reduce，將所有節(jié)點(diǎn)分為[k]組，每組內(nèi)部進(jìn)行Ring All-Reduce;第2步，Intergroup All-Reduce，在每個(gè)組中選擇一個(gè)節(jié)點(diǎn)，這些節(jié)點(diǎn)進(jìn)行Ring All-Reduce。在第三個(gè)階段Intragroup Broadcast，將更新后的參數(shù)傳遞到所有節(jié)點(diǎn)。

假設(shè)有[N]個(gè)節(jié)點(diǎn)，被分為[k]組，設(shè)[S]為數(shù)據(jù)塊大小， [B]為帶寬，[a]為兩個(gè)通信點(diǎn)間的延時(shí)。

則需要的總時(shí)間為：

[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節(jié)點(diǎn)數(shù)增加導(dǎo)致通信時(shí)延過大的局限性，并且通信時(shí)間依然受到環(huán)上最慢連接的限制。

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

3.1 PS-HRA模型架構(gòu)

PS-HRA算法綜合了去中心化的Ring All-Reduce算法和中心化的PS框架，形成了一個(gè)分層的通訊優(yōu)化框架。它由若干個(gè)去中心的節(jié)點(diǎn)的環(huán)和一個(gè)中心節(jié)點(diǎn)組成?？偟膩碚f，相比于傳統(tǒng)的參數(shù)服務(wù)器架構(gòu)，中心節(jié)點(diǎn)在我們的框架下控制的是各個(gè)環(huán)而不是單個(gè)節(jié)點(diǎn)，因此采用我們的通訊框架可以有效地解決中心節(jié)點(diǎn)帶寬限制的問題，即在僅有一個(gè)中心節(jié)點(diǎn)的情況下，我們的框架可以擴(kuò)展更多的計(jì)算節(jié)點(diǎn)，并保證不會(huì)因?yàn)閹拞栴}使傳輸數(shù)據(jù)速率限制下降。同時(shí)環(huán)間采用異步PS架構(gòu)，在環(huán)內(nèi)計(jì)算節(jié)點(diǎn)算力相近的條件下，可以改善環(huán)間算力差異帶來的同步等待問題。

3.2 分組策略

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

算法1. PS-HRA分組算法

輸入：節(jié)點(diǎn)[id] ， 節(jié)點(diǎn)數(shù)目[N]， 分組數(shù)目[k]

輸出：環(huán)[ring]

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

2. [worknode]←[testspeeddemo]; //測試數(shù)據(jù)發(fā)送到工作節(jié)點(diǎn)

3. [loop：]

4. [ifworknode[id].status= =OK q.enqueueid; ]//將完成計(jì)算的節(jié)點(diǎn)加入隊(duì)列

5. [ifq.size≥N/k {ring←q1…N/k;}] //前[N/k]的計(jì)算節(jié)點(diǎn)形成環(huán)

3.3 環(huán)間異步PS架構(gòu)

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

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

1） 執(zhí)行PS-HRA分組算法，將算力接近的[N]個(gè)節(jié)點(diǎn)均勻分成[k]個(gè)小組。

2） 小組內(nèi)并行scatter-reduce算法，執(zhí)行結(jié)束之后每個(gè)計(jì)算節(jié)點(diǎn)具有[k/N]個(gè)數(shù)據(jù)塊的環(huán)內(nèi)完整參數(shù)。

3） 小組內(nèi)并行all-gather算法，執(zhí)行結(jié)束后每個(gè)計(jì)算節(jié)點(diǎn)具有環(huán)內(nèi)數(shù)據(jù)塊的全部參數(shù)。

4） 環(huán)內(nèi)隨機(jī)選出一個(gè)計(jì)算節(jié)點(diǎn)作為代表節(jié)點(diǎn)，向參數(shù)服務(wù)器發(fā)送參數(shù)。

5） 參數(shù)服務(wù)器接收參數(shù)直接對參數(shù)聚合更新，無須等待其他環(huán)。

6） 代表節(jié)點(diǎn)接收參數(shù)。

7） 代表節(jié)點(diǎn)在環(huán)內(nèi)進(jìn)行參數(shù)廣播。

8） 各節(jié)點(diǎn)更新模型，進(jìn)行下一輪的訓(xùn)練。

對于傳統(tǒng)的參數(shù)服務(wù)器架構(gòu)，假設(shè)有16個(gè)計(jì)算節(jié)點(diǎn)，中心節(jié)點(diǎn)最多可能面臨著16個(gè)計(jì)算節(jié)點(diǎn)同時(shí)和它進(jìn)行通訊;假設(shè)有4個(gè)環(huán)，而PS-HRA的通訊模型最多只需要四個(gè)計(jì)算節(jié)點(diǎn)和中心節(jié)點(diǎn)通訊，這大大減少了帶寬瓶頸。

4 PS-HRA時(shí)間代價(jià)估計(jì)與性能分析

4.1 時(shí)間代價(jià)估計(jì)

假設(shè)有[N]個(gè)節(jié)點(diǎn)，被分為[k]組，設(shè)[S]為總的數(shù)據(jù)量，[C]為每字節(jié)數(shù)據(jù)的計(jì)算耗時(shí)，[B]為帶寬，[a]為兩個(gè)通信點(diǎn)間的延時(shí)。

組內(nèi)scatter-reduce所需時(shí)間為 [Nk-1*a+kSNB+kSN*C]

組內(nèi)all-gather 所需時(shí)間為 [Nk-1*a+kSNB]

組間使用異步PS架構(gòu)所需時(shí)間為 [2*a+SB+S*C]

每組代表節(jié)點(diǎn)將更新后的參數(shù)傳給組內(nèi)其余節(jié)點(diǎn)所需時(shí)間為：

[a+SB]

通過計(jì)算可得，需要的總時(shí)間為：

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

4.2 性能分析

為了證明PS-HRA算法性能的優(yōu)越性，需要在傳輸時(shí)間延遲或節(jié)點(diǎn)數(shù)量變化的情況下，對比前述四種算法的通信時(shí)間。

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

傳輸延遲包括算力差的節(jié)點(diǎn)帶來的同步開銷，所以由圖5可知，傳輸延遲越大，Ring All-Reduce的通信時(shí)間越長;然而PS架構(gòu)的通信時(shí)間變化很小，其性能受傳輸延遲影響小。另外，不難發(fā)現(xiàn)PS-HRA的通信時(shí)間最少且隨著傳輸延遲的增大，其性能的優(yōu)越性更加顯著。

由圖6可知，隨著節(jié)點(diǎn)數(shù)的增加，PS-HRA算法的通信時(shí)間基本不變，而其余三種算法的通信時(shí)間會(huì)隨著節(jié)點(diǎn)數(shù)的增加而不斷變大。當(dāng)節(jié)點(diǎn)數(shù)較大時(shí)，PS-HRA的通信時(shí)間最少。

5 結(jié)束語

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

在未來的研究中，我們會(huì)在更多、更復(fù)雜的情況下進(jìn)一步驗(yàn)證我們通信策略的有效性。另外，我們會(huì)把我們的通訊策略應(yīng)用到更多機(jī)器學(xué)習(xí)的場景中。

參考文獻(xiàn)：

[1] 杜海舟，黃晟.分布式機(jī)器學(xué)習(xí)中的通信機(jī)制研究綜述[J].上海電力大學(xué)學(xué)報(bào)，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] 李曉明.數(shù)據(jù)并行計(jì)算：概念、模型與系統(tǒng)[J].計(jì)算機(jī)科學(xué)，2000，27（6）：1-5.

【通聯(lián)編輯：謝媛媛】