分布式機(jī)器學(xué)習(xí)網(wǎng)絡(luò)通信優(yōu)化技術(shù)

張漢鋼，鄧鑫源，宋 曄，薛旭偉，郭秉禮，黃善國（北京郵電大學(xué)，北京 100876）

1 概述

近些年隨著互聯(lián)網(wǎng)的蓬勃發(fā)展及大數(shù)據(jù)技術(shù)的不斷演進(jìn)，在圖像識別、語言翻譯、自動駕駛等相關(guān)領(lǐng)域，機(jī)器學(xué)習(xí)取得了眾多突破并被廣泛成功應(yīng)用。在進(jìn)行機(jī)器學(xué)習(xí)訓(xùn)練時，通常需要大量的數(shù)據(jù)用于計(jì)算，以保證模型的準(zhǔn)確性。在海量數(shù)據(jù)的背景下，訓(xùn)練所使用的數(shù)據(jù)集的規(guī)模越來越大，訓(xùn)練的模型越來越復(fù)雜，單個機(jī)器無法滿足存儲與算力的需求，分布式機(jī)器學(xué)習(xí)將是未來發(fā)展的必然趨勢。分布式機(jī)器學(xué)習(xí)需要快速、高效的連接網(wǎng)絡(luò)作為支撐，該網(wǎng)絡(luò)結(jié)構(gòu)有2 種主要的并行方式：數(shù)據(jù)并行與模型并行。數(shù)據(jù)并行是最為常見的并行模式，廣泛應(yīng)用于大規(guī)模分布式系統(tǒng)中。在采用數(shù)據(jù)并行的方式進(jìn)行訓(xùn)練時，各個機(jī)器之間需要對大量的數(shù)據(jù)進(jìn)行同步，例如GPT-3就擁有超過1 750 億個機(jī)器學(xué)習(xí)參數(shù)。通信時間在訓(xùn)練時間中占比較大，如何減少通信開銷成為了熱門課題。數(shù)據(jù)并行使用all-reduce操作進(jìn)行梯度同步，其中Ring all-reduce算法在高性能計(jì)算領(lǐng)域得到廣泛應(yīng)用。

Ring all-reduce基于同步通信，在計(jì)算節(jié)點(diǎn)的性能存在差異時，整體的計(jì)算性能會被較慢的計(jì)算節(jié)點(diǎn)拖慢。為解決同步通信的這一弊端，異步并行（ASP）與延遲同步并行（SSP）被相繼提出，但Ring all-reduce 架構(gòu)難以支持ASP和SSP等模型的一致性協(xié)議。本文進(jìn)一步分析Ring all-reduce 算法特性，提出一種對其同步參數(shù)的過程進(jìn)行加速的優(yōu)化方案，降低慢節(jié)點(diǎn)帶來的影響。

2 優(yōu)化策略及理論分析

2.1 Ring all-reduce基本分析

百度在2017年提出將Ring all-reduce 架構(gòu)引入到分布式機(jī)器學(xué)習(xí)中來，使這個架構(gòu)受到廣泛關(guān)注。Ring all-reduce 算法分為2 個部分：Reduce_Scatter 和Allgather。Reduce_Scatter 階段按照n個訓(xùn)練節(jié)點(diǎn)數(shù)將數(shù)據(jù)劃分為n個chunk，將從上一相鄰節(jié)點(diǎn)接收到的chunk 同本地chunk 進(jìn)行reduce 操作并發(fā)，該階段包含n-1 個步驟，如圖1 所示。Allgather 階段是將匯總后的梯度段依次傳遞，用接收到的內(nèi)容替換本地內(nèi)容，該階段包含n-1個步驟（見圖2）。

圖1 Reduce_Scatter階段示意

圖2 Allgather階段示意

在Github 上百度公開給出名為baidu-allreduce 的小型c++庫，具體演示了環(huán)形的all-reduce 與Allgather過程。Ring all-reduce 的同步過程體現(xiàn)在進(jìn)行reduce時，調(diào)用了同步CUDA 流（cudaStreamSynchronize）來等待所有GPU都完成操作。reduce操作執(zhí)行完畢也就意味著完成了一輪同步過程，隨后向后一個節(jié)點(diǎn)進(jìn)行發(fā)送操作。

當(dāng)某個節(jié)點(diǎn)的reduce 階段所用時間顯著變長，即該節(jié)點(diǎn)為慢節(jié)點(diǎn)時，若按照Ring all-reduce 過程進(jìn)行同步，則需要等慢節(jié)點(diǎn)多產(chǎn)生的時間，而這多出的時間將直接加在Ring all-reduce 原有的時間開銷上，影響到Ring all-reduce 算法的執(zhí)行效率。為了減輕慢節(jié)點(diǎn)帶來的負(fù)面影響，如何利用慢節(jié)點(diǎn)較正常節(jié)點(diǎn)多出來的時間便成為了關(guān)鍵，下文給出一種優(yōu)化策略思路。

2.2 優(yōu)化策略

以4 節(jié)點(diǎn)為例對Ring all-reduce 的Reduce_Scatter階段進(jìn)行分析。設(shè)置TRS為節(jié)點(diǎn)接收數(shù)據(jù)、進(jìn)行reduce操作和發(fā)送數(shù)據(jù)的時間，Ttrans為傳輸數(shù)據(jù)的時間。A節(jié)點(diǎn)將a0 發(fā)送至B 節(jié)點(diǎn)，B 節(jié)點(diǎn)將收到的a0 與本地的b0進(jìn)行匯總，并在下一次通信時將匯總的內(nèi)容發(fā)送至C節(jié)點(diǎn)（見圖3）。在這一時期，B節(jié)點(diǎn)因reduce過程緩慢而成為慢節(jié)點(diǎn)，TRS變長，C 節(jié)點(diǎn)需要等待的TRS+Ttrans時間也隨之變長，將導(dǎo)致訓(xùn)練效率下降。

圖3 Ring all-reduce的通信步驟

本文提出一種優(yōu)化策略，設(shè)置T表示在沒有慢節(jié)點(diǎn)情況下的TRS的平均值，則Reduce_Scatter 階段1 個步驟的實(shí)際開銷TRS=T+△T，其中△T為因外在因素導(dǎo)致的時間波動；αT為節(jié)點(diǎn)觸發(fā)策略的時間閾值，α>1。當(dāng)出現(xiàn)慢節(jié)點(diǎn)時，△T變大；當(dāng)TRS=T+△T>αT時，則觸發(fā)策略，具體過程如圖4所示。

圖4 改變策略后的通信步驟

假設(shè)B 節(jié)點(diǎn)是慢節(jié)點(diǎn)，Tt為C 節(jié)點(diǎn)等待時間。當(dāng)TRS=T+△T>Tt=αT+Ttrans時，改變C 節(jié)點(diǎn)的操作步驟。此時C 節(jié)點(diǎn)將不再等待接收B 節(jié)點(diǎn)發(fā)送的a0＋b0，而是直接將本地的c0 發(fā)送至d0；B 節(jié)點(diǎn)達(dá)到TRS時將跳過C節(jié)點(diǎn)直接與D節(jié)點(diǎn)進(jìn)行通信。為了更直觀地分析時間上的重疊，假設(shè)如下。

a）各個結(jié)點(diǎn)之間傳輸數(shù)據(jù)的時間Ttrans都一致。

b）若節(jié)點(diǎn)為慢節(jié)點(diǎn)，其TRS=T+△T且△T>0。

c）若節(jié)點(diǎn)為非慢節(jié)點(diǎn)，其TRS=T。

d）設(shè)定B節(jié)點(diǎn)為慢節(jié)點(diǎn)。如圖5 所示，當(dāng)B 節(jié)點(diǎn)是慢節(jié)點(diǎn)且C 節(jié)點(diǎn)等待至αT+Ttrans時，C 節(jié)點(diǎn)直接向D 節(jié)點(diǎn)傳遞數(shù)據(jù)，D 節(jié)點(diǎn)先收到來自C 節(jié)點(diǎn)的數(shù)據(jù)并等待B 節(jié)點(diǎn)發(fā)送來的數(shù)據(jù)。在收到2 份數(shù)據(jù)后，D 節(jié)點(diǎn)一并進(jìn)行匯總。截至D 節(jié)點(diǎn)匯總所有數(shù)據(jù)時，策略帶來的時間減少值如式（1）所示。

圖5 B節(jié)點(diǎn)發(fā)送的消息先到達(dá)D節(jié)點(diǎn)

由式（1）可知，當(dāng)△T較大時，在T+△T>αT+Ttrans情況下策略帶來的時間減少值為固定時間。

如圖6 所示，當(dāng)TRS=T+△T

圖6 C節(jié)點(diǎn)發(fā)送的消息先到達(dá)D節(jié)點(diǎn)

當(dāng)△T與α取合適的值時，可以使Tdec取得一個較大值，即在有慢節(jié)點(diǎn)的情況下能減輕慢節(jié)點(diǎn)帶來的負(fù)面影響，達(dá)到優(yōu)化效果。

3 仿真實(shí)驗(yàn)設(shè)置及結(jié)果分析

采用OMNet++框架對Ring all-reduce 算法和改進(jìn)策略進(jìn)行模擬。因?yàn)锳llgather 階段這2 種方式的開銷相同，故僅比較2 種方式完成Reduce_Scatter 所用的時間（分別是TRing和TStrategy）?；诘?.2 節(jié)做出的基本假設(shè)，對相關(guān)的仿真參數(shù)進(jìn)行設(shè)定：T=3 μs；Ttrans為固定鏈路時延5 μs 與數(shù)據(jù)包傳遞時間之和；數(shù)據(jù)包大小取8 KB；鏈路帶寬為100 Gbit/s。設(shè)置不同的△T與α的值，優(yōu)化效果即Tdec與TRing的比值。

如圖7 所示，Ring all-reduce 算法完成Reduce_Scatter 階段所用時間隨著△T的增加而線性增加，而采用優(yōu)化策略完成Reduce_Scatter 階段所用時間均得到減少，說明采取優(yōu)化策略能減少時間開銷。進(jìn)一步分析，如圖8 所示，當(dāng)α不變時，優(yōu)化效果大體上隨著△T的增大而先增后減。當(dāng)△T較小時，TRS=T+△T

圖7 Ring all-reduce與優(yōu)化策略時間開銷比較

圖8 △T與α對優(yōu)化效果的影響

當(dāng)TRS=T+△T>Tt=αT+Ttrans，優(yōu)化效果與α的變化無關(guān)（見圖9），結(jié)果符合式（1）預(yù)期。

圖9 α變化對優(yōu)化效果影響

T+△T<αT+Ttrans時，即慢節(jié)點(diǎn)跳過后一節(jié)點(diǎn)發(fā)送的數(shù)據(jù)先到達(dá)，α 增大，說明后一節(jié)點(diǎn)最長等待時間將延長，觸發(fā)機(jī)制的機(jī)會將相對減少，機(jī)制帶來的增益受到限制，優(yōu)化效果隨之減弱。

當(dāng)△T是一個較大偏置時，有T+△T>αT+Ttrans，即后一節(jié)點(diǎn)發(fā)送的數(shù)據(jù)先到達(dá)而慢節(jié)點(diǎn)發(fā)送的數(shù)據(jù)后到達(dá)，此時獲得的時間減少值Tdec是一個定值?！鱐的增加會導(dǎo)致TRing的時間線性增加，Tdec與TRing的比值也隨之減少，即優(yōu)化效果減弱（見圖10）。

圖10 △T變化對優(yōu)化效果影響

4 總結(jié)

在分布式計(jì)算中，Ring all-reduce是一種常用的通信模式，然而由于網(wǎng)絡(luò)中存在慢節(jié)點(diǎn)，通信延遲分布不均勻，會影響整體性能。本文聚焦Ring all-reduce算法的Reduce_Scatter階段，提出一種創(chuàng)新性的優(yōu)化策略實(shí)現(xiàn)對慢節(jié)點(diǎn)的感知與通信調(diào)整，減小慢節(jié)點(diǎn)對整體通信性能的負(fù)面影響。仿真結(jié)果顯示，相較于傳統(tǒng)的Ring all-reduce 算法，本文提出的優(yōu)化策略成功降低了通信的時間開銷，訓(xùn)練時長最高節(jié)省了25.3%。