這篇文章主要介紹了一個名為Aluminum通訊庫，在這個庫中主要針對Allreduce做了一些關於計算通訊重疊以及針對延遲的優化，以加速分散式深度學習訓練過程。 ### 分散式訓練的通訊需求 #### 通訊何時發生 一般來說，神經網路的訓練過程分為三步：前向傳播、反向傳播以及引數優化。在使用資料並行進行分散式訓練的情況下，通訊主要發生在反向傳播之後與引數優化之前，在此階段各個計算節點需要進行梯度的同步。廣義上來講，梯度的同步過程符合Allreduce語義。從實現上來說，我們既可以通過中心化的引數伺服器架構來實現梯度同步，也可以通過去中心化的`MPI_Allreduce`介面實現梯度同步。本文主要關注的是去中心化的Allreduce實現。 神經網路本身是一個層次性的架構，在反向傳播的過程中，梯度是從從後向前依次計算的。換句話說，靠近輸出層的梯度先於靠近輸入層的梯度被計算出來。那麼，通訊需求就依賴於梯度同步的粒度。在極端情況下，我們可以把所有層的梯度都放到一個buffer中去做Allreduce（一般也沒人這麼做吧，頂多針對某些層做Tensor Fusion）。另一方面，我們可以針對每一層的梯度做Allreduce，只要該層的梯度被計算出來，我們就對它進行同步。本文所實現的方法屬於第二類，神經網路的每一層都有自己的buffer。對於某些需要學習引數的層如BN，本文的實現會為相應的引數申請單獨的buffer。 #### 通訊量有多少 在資料並行分散式訓練過程中，每一次迭代的通訊量一般只取決於神經網路的模型結構，與其他的因素無關（當然也有例外，比如Embedding層的通訊量會受batch size的影響。但Embedding一般只用在CTR和NLP模型，對於CV的模型，batch size一般不會影響整體通訊量）。  上圖展示了AlexNet和ResNet-50引數量的直方圖。可以看到，AlexNet是一個比較淺的網路，它包含5個卷積層以及3個較大的全連線層，這3個全連線層包含了模型大部分引數。在AlexNet中，除了最後一層，每一層都有一個偏置項，因此它包含許多small buffer。ResNet-50是一個更現代的CNN模型，它包含更多的卷積層和較少的全連線層，以及一些BN層。ResNet-50不包含偏置項，但由於BN層的存在，它依然包含許多small buffer。顯然，對於AlexNet和ResNet-50來說，我們需要對許多small buffer進行Allreduce操作。除此之外，模型中也存在很多large buffer需要通訊，這就導致單一的Allreduce演算法並不能在所有的buffer size達到最優的通訊效率。 #### 計算和通訊開銷 本小節作者主要比較了ResNet-50在訓練過程中的“計算”和“通訊”開銷。為什麼這裡給計算和通訊打引號呢，是因為作者統計的並不是真正的計算和通訊的開銷，而是計算和通訊各自的主要開銷。對於計算時間，作者使用cuDNN統計了ResNet-50所有卷積層的計算開銷；而對於通訊時間，作者使用NCCL統計了Allreduce過程的通訊開銷。由此可知，作者所統計的計算和通訊開銷並不是完整訓練過程中的開銷，只是真正開銷的一部分。 在討論計算和通訊開銷時，作者分了兩種情況進行討論：強可擴充套件和弱可擴充套件下的計算通訊開銷比例。關於強可擴充套件和弱可擴充套件的定義，可以參考[HPC Wiki](https://hpc-wiki.info/hpc/Scaling_tests#Strong_or_Weak_Scaling)： >In case of strong scaling, the number of processors is increased while the problem size remains constant. This also results in a reduced workload per processor. In case of weak scaling, both the number of processors and the problem size are increased. This also results in a constant workload per processor.  上圖展示了強可擴充套件（左）和弱可擴充套件（右）兩種情況下計算和通訊時間各自所佔的比例。在強可擴充套件的情況下，隨著GPU數量的增加，總的batch size保持不變。這樣一來，分到每塊GPU上的訓練樣本數就會減少，每塊GPU的計算開銷也會相應降低。由於總的batch size不變，所以總的迭代次數也不會變化，那麼隨著GPU的增加，節點間的通訊開銷也會增加。在弱可擴充套件的情況下，隨著GPU數量的增加，每塊GPU所處理的訓練樣本數保持不變，總的batch size會隨著GPU數量的增加而增加。這就表明訓練的迭代次數會減少，也意味著通訊發生的頻次會降低。在這種情況下，計算是線性擴充套件的，但是通訊的佔比還是會隨著GPU數量的增加而增加。從圖上可以看出，當GPU數量達到2048時，通訊開銷大概是計算開銷的12倍。 無論是強可擴充套件還是弱可擴充套件的情況下，隨著GPU數量的增加，節點間的通訊開銷所佔的比例會越來越大，這就會抵消掉分散式訓練帶來的收益。 ### 優化方法 在這一節，作者介紹了兩種通訊優化方法：計算——通訊重疊以及latency-efficient Allreduce演算法。這兩個優化方法都不是由作者首次提出，但是作者在他們的通訊庫Alumunium中實現了這些優化演算法。 計算——通訊重疊的思路非常簡單，就是儘可能地讓Allreduce操作和反向傳播、引數優化等操作並行執行。如果某一層的梯度被計算出來，那麼直接在相應的buffer上非同步地去執行Allreduce操作。其它層的反向傳播可以按照相同的方式執行，並在該層的優化階段開始時完成Allreduce。這樣一來，所有的Allreduce都可以被相關層中的反向傳播以及剩餘層中的計算所隱藏。然而，反向傳播與引數優化之類的操作一般會放在GPU上去執行，而節點間的通訊一般由MPI完成，我們並不希望CUDA操作阻塞MPI通訊，反之亦然。因此，我們需要一些額外的編碼工作來實現相應的功能。 如果我們能夠讓Allreduce演算法跑的更快，那麼通訊操作會更容易地被隱藏到計算中。目前，許多深度學習框架都採用Ring Allreduce演算法實現梯度同步，Ring Allreduce是一種頻寬最優的演算法，它在共享記憶體系統以及小型的分散式記憶體系統中表現良好。但是Ring Allreduce並不是延遲最優的，如果進行Reduce的buffer size比較小，那麼Ring Allreduce的效能就會下降。基於樹形結構的Allreduce演算法能夠針對延遲進行優化，Recursive-Doubling和Recursive-Halving/Recursive-Doubling演算法在傳輸小訊息上表現更好。假設$\alpha$是網路延遲，$\beta$是頻寬的倒數，$p$是處理器數量，$n$是buffer的大小，那麼不同Allreduce演算法的通訊時間如下表所示： |拓撲結構|演算法名稱|通訊時間| |:---:|:---:|:---:| |環形|Ring|$2(p-1)\alpha+2\frac{p-1}{p}n\beta$| |樹形|Recursive-Doubling|$\log_p(\alpha+n\beta)$| |樹形|Recursive-Halving/Recursive-Doubling|$2\log_p \alpha + 2\frac{p-1}{p}n\beta$| 可以看到，Recursive-Halving/Recursive-Doubling相比於Ring Allreduce，二者的頻寬項是相同的，而Recursive-Halving/Recursive-Doubling具有更小的延遲。本文作者開發的Aluminium通訊庫，把基於樹形拓撲的Allreduce演算法整合到NCCL中，並根據緩衝區大小和處理器數量動態選擇最優的Allreduce演算法。 ### 實現思路 當前的MPI分發版本已經針對各種型別的通訊配置做了優化，它會根據緩衝區大小以及處理器數量選擇合適的演算法。此外，一些MPI版本同樣支援CUDA，它們可以接受GPU緩衝區的指標並在該緩衝區上進行通訊操作。那麼我們為什麼不直接使用這些CUDA-aware MPI中實現的Allreduce演算法呢？這是因為，MPI不瞭解由使用者定義的CUDA流，因此MPI和CUDA程式設計模型之間會出現語義不匹配的情況，這會產生不必要的同步操作，從而導致額外的通訊和計算開銷。 使用者定義的CUDA執行流對MPI Runtime來說是透明的，因此，當用戶把一個GPU緩衝區傳給MPI程式時，MPI Runtime無法確定將要寫入該緩衝區的流上是否有待處理的計算。為了保證kernel函式寫入GPU緩衝區的正確性，使用者必須手動地同步CUDA流，這就導致我們無法重疊計算和通訊操作。同樣，MPI Runtime對於CUDA流來說也是透明的，CUDA流無法確定是否應該等待某些阻塞操作的完成。顯然，這種頻繁的同步阻塞會嚴重影響網路和GPU的利用率。此外，作者還指出，當前版本的CUDA-aware MPI不能在多執行緒環境下正確地處理GPU緩衝區上的操作。 既然當前版本的CUDA-aware MPI有以上種種問題，那麼作者就提出了幾種解決方案。第一種方案是把同步操作放到MPI中去執行，這樣能夠保證正確性。但是，MPI Runtime並不知道哪一個CUDA流會操作GPU緩衝區，因此每次同步操作都要阻塞GPU上的所有流。顯然，這種方案並沒有解決上述效能問題。第二種方案是把MPI通訊操作當做kernel函式放到CUDA流中執行。NCCL就是採用這種方案，其中的每個通訊操作都會接收一個CUDA流，把通訊操作放到該流中執行。這種方法不會阻塞CPU，但是會阻塞GPU上的CUDA流。不幸的是，MPI中的通訊操作並不能接收CUDA流作為引數。在第三種方案中，作者發現把某個CUDA流單獨繫結到MPI通訊器上是可行的，所有使用MPI通訊器和GPU緩衝區的操作都可以假定該緩衝區由該流中的某個kernel函式寫入，並且MPI Runtime僅對該流執行適當的同步。在MPI中，這種關係可以被實現為繫結在通訊器上的某個屬性。 ### Aluminum通訊庫 Aluminum是作者開發的一個開源通訊庫，它以MPI和NCCL等為後端，提供了泛化的通訊API。相比於MPI和NCCL，Aluminum更像是一個介面層，只需要簡單的改動，它就能跑在不同的硬體上。 #### 特性介紹 Aluminum基於C++11實現，API類似於MPI中的介面，通過模板技術實現了在不同後端之間的切換。下表展示了Aluminum支援的後端以及各個後端的一些特性。 |後端|演算法支援|特性| |:---:|:---:|:---:| |MPI|Ring、Recursive-Doubling、Recursive-Halving/Recursive-Doubling|Ubiquitous, optimized| |NCCL|Ring|GDR, optimized for GPUs| |MPI-CUDA|Ring、Recursive-Doubling、Recursive-Halving/Recursive-Doubling|Host-Transfer Allreduce| #### 實現細節 ##### 通訊引擎 在Aluminum中，所有在Host端進行且不阻塞主執行緒通訊的操作都會被放到後臺執行緒中執行，這個執行緒被稱為通訊引擎。這個後臺執行緒由通訊庫繫結到CPU的某個核上，並且採用某種啟發式的方法避免與其他執行緒產生衝突。所有的非同步操作都通過一個`state`物件提交到通訊引擎中，這些`state`物件被放在一個無鎖的生產者——消費者佇列中等待執行。通過呼叫`state`物件的`step`方法，相應的操作就會被非阻塞式地執行。當操作執行完成後，通訊引擎可以通過在`request`物件中自動設定一個標誌，以指示其他執行緒該操作已完成。Aluminium的MPI後端利用通訊引擎在Host端提供非同步呼叫以實現某些自定義演算法，並通過`MPI_Test`輪詢提供非阻塞式MPI操作。  Aluminum主要關注的是對GPU緩衝區的非阻塞式通訊。對於NCCL後端來說，非阻塞式的Allreduce會自動地執行在CUDA流中，整個過程如上圖所示。通過呼叫`Al::Wait`，Aluminum就可以等待通訊操作完成，這就允許Host端在GPU通訊時進行其他操作。  對於延遲是瓶頸的工作流來說，Aluminum基於MPI的樹形Allreduce實現了一種阻塞式的Allreduce演算法。這種實現相對來說比較簡單直觀：把GPU緩衝區拷貝到Host記憶體中，然後執行Allreduce，最後把結果拷貝回GPU緩衝區。為了避免Host端被阻塞，該操作相關的kernel函式以及事件都會被放入繫結在通訊器的CUDA流中，然後委託給通訊引擎執行。在資料從GPU拷貝到CPU記憶體的過程中，Aluminum會輪詢CUDA流以判斷相應的GPU緩衝區是否已經寫入完成，這樣可以有效地防止資料競爭。阻塞式Allreduce的實現如上圖所示。  上圖闡述了非阻塞式Allreduce的實現，可以看到它與阻塞式類似，只不過資料拷貝以及同步等待都被轉移到繫結在MPI通訊器的CUDA流上。 除了Allreduce，Aluminum還實現了Send和Recv操作。對於Send操作來說，Aluminum首先把資料從GPU視訊記憶體拷貝到CPU記憶體中，然後再呼叫`MPI_Isend`；對於Recv操作來說，先呼叫`MPI_Irecv`接收資料，然後再把資料拷貝到GPU視訊記憶體中。 ### 實驗測試 #### 計算——通訊重疊 按照前面章節的配置，重新使用Aluminum（NCCL作為後端）跑了一遍ResNet-50，分別比較強可擴充套件（左）和弱可擴充套件（右）下的計算和通訊佔比。  從圖上可以看到，在節點規模較小時，Aluminum基本上可以把通訊全部重疊到計算中。在強可擴充套件的情況下，如果使用32塊GPU，那麼整體訓練速度大概能加快1.4倍；如果GPU數量超過32塊，那麼我們就沒有足夠的計算量來隱藏通訊，加速效果就不太明顯。在弱可擴充套件的情況下，如果使用不超過256塊GPU，那麼通訊就可以被隱藏到計算中；但是當節點規模更大，使用1024或者2048塊GPU時，通訊開銷依然很大。 #### 延遲 為了比較NCCL和host-transfer Allreduce演算法的延遲，作者比較了不同節點數量（規模為2個節點到512個節點，每個節點上4塊GPU）以及不同緩衝區大小下的傳輸延遲。下圖是不同節點規模測試得到的效能結果。  可以看到，NCCL在節點規模很小的時候效能比較好，這是因為節點規模較小時節點間的延遲也比較小，無法體現Ring-based Allreduce和Tree-based Allredcue的差距。當節點規模增加到16時，host-transfer Allreduce的優勢就體現出來了。當使用32個節點時，它比NCCL快了2倍；使用512個節點時則快了20倍。注意到在小於8個節點（32塊GPU）的規模下，NCCL傳輸小訊息的延遲都比較低，這是因為NCCL內部用了GPUDriect RDMA以及GPU拓撲資訊來減少通訊開銷和延遲。下面這幅圖展示了給定GPU數量和緩衝區大小的情況下，使用哪種Allreduce演算法的延遲比較低。其中綠色的點表示這種配置下，host-transfer Allreduce要比NCCL快。  此外，作者還提出了一種名為minimal的動態選擇演算法，它會針對緩衝區大小和處理器數量在非阻塞式host-transfer Allreduce演算法和NCCL Allreduce中選擇合適的演算法，選擇的依據就是之前跑的benchmark。根據圖中的實驗結果可以看出，這種動態選擇演算法在弱可擴充套件的情形下作用比較大。 ![](https://img2020.cnblogs.com/blog/1023521/202010/1023521-20201009211813207-1272138