分布式機器學習網絡通信優化技術

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

1 概述

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

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

2 優化策略及理論分析

2.1 Ring all-reduce基本分析

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

圖1 Reduce_Scatter階段示意

圖2 Allgather階段示意

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

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

2.2 優化策略

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

圖3 Ring all-reduce的通信步驟

本文提出一種優化策略，設置T表示在沒有慢節點情況下的TRS的平均值，則Reduce_Scatter 階段1 個步驟的實際開銷TRS=T+△T，其中△T為因外在因素導致的時間波動；αT為節點觸發策略的時間閾值，α>1。當出現慢節點時，△T變大；當TRS=T+△T>αT時，則觸發策略，具體過程如圖4所示。

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

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

a）各個結點之間傳輸數據的時間Ttrans都一致。

b）若節點為慢節點，其TRS=T+△T且△T>0。

c）若節點為非慢節點，其TRS=T。

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

圖5 B節點發送的消息先到達D節點

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

如圖6 所示，當TRS=T+△T

圖6 C節點發送的消息先到達D節點

當△T與α取合適的值時，可以使Tdec取得一個較大值，即在有慢節點的情況下能減輕慢節點帶來的負面影響，達到優化效果。

3 仿真實驗設置及結果分析

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

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

圖7 Ring all-reduce與優化策略時間開銷比較

圖8 △T與α對優化效果的影響

當TRS=T+△T>Tt=αT+Ttrans，優化效果與α的變化無關（見圖9），結果符合式（1）預期。

圖9 α變化對優化效果影響

T+△T<αT+Ttrans時，即慢節點跳過后一節點發送的數據先到達，α 增大，說明后一節點最長等待時間將延長，觸發機制的機會將相對減少，機制帶來的增益受到限制，優化效果隨之減弱。

當△T是一個較大偏置時，有T+△T>αT+Ttrans，即后一節點發送的數據先到達而慢節點發送的數據后到達，此時獲得的時間減少值Tdec是一個定值?！鱐的增加會導致TRing的時間線性增加，Tdec與TRing的比值也隨之減少，即優化效果減弱（見圖10）。

圖10 △T變化對優化效果影響

4 總結

在分布式計算中，Ring all-reduce是一種常用的通信模式，然而由于網絡中存在慢節點，通信延遲分布不均勻，會影響整體性能。本文聚焦Ring all-reduce算法的Reduce_Scatter階段，提出一種創新性的優化策略實現對慢節點的感知與通信調整，減小慢節點對整體通信性能的負面影響。仿真結果顯示，相較于傳統的Ring all-reduce 算法，本文提出的優化策略成功降低了通信的時間開銷，訓練時長最高節省了25.3%。