提升12倍！香港浸會大學與MassGrid發布低帶寬高效AI訓練新演算法

允中 發自 凹非寺

量子位 報道 | 公眾號 QbitAI

隨著訓練數據量的增加和深度神經網路（DNN）的日益複雜，分散式計算環境（如GPU集群）被廣泛採用，以加速DNN的訓練。分散式計算網路在機器學習方面的瓶頸在於節點之間的數據傳輸效率，那如何在這一網路下高效地進行AI訓練？

2018年，香港浸會大學異構計算實驗室與MassGrid合作，通過研究提出一種可用於低帶寬網路的全局Top-k稀疏化的分散式同步SGD演算法，並通過實驗論證出在低帶寬網路下也能高效進行AI訓練。目前實驗結果及相關論文已被ICDCS workshop收錄。

數據並行的分散式同步隨機梯度下降（S-SGD）方法是訓練大規模神經網路常用的優化器之一。與單節點的SGD相比，S-SGD將工作負載分配給多個計算節點以加速訓練，但它也引入了在每次迭代中交換模型參數或梯度的通信開銷。

舉例說明

假設有P個節點用S-SGD訓練DNN模型。在每次迭代中，所有計算節點都會採用不同的小批量(mini-batch)數據來並行計算模型的梯度。然後，對每個節點的梯度進行平均後來更新模型，這便引入較大的通信開銷。

由於計算節點的加速器（如GPU和TPU）的計算能力比網路速度的增長快得多，網路通信性能通常成為訓練的性能瓶頸，特別是當通信與計算比率很高時。

許多大型IT公司使用昂貴的高速網路（如40 / 100Gbps IB或乙太網）來減少通信壓力，但仍有許多研究人員和小公司只能使用由1Gig-Ethernet等低帶寬網路連接的消費級GPU。

為了克服通信的性能瓶頸，可以通過使用更大的mini-batch來增加工作負載從而降低通信與計算比，或者減少每次通信過程中所需的通信量：

一方面，許多大批量SGD技術已經提出了一些優化策略來提高mini-batch而不會丟失模型準確性。

另一方面，研究人員也已經提出了梯度稀疏化，量化和壓縮等方法，在不影響收斂速度的情況下顯著減小交換梯度的數量以降低通信量。

在模型/梯度壓縮技術中，Top-k稀疏化是關鍵方法之一，它可以將每個節點梯度稀疏到約為原來的千分之一（即 99.9%的梯度置為零而無需傳輸這些零值）。

Top-k稀疏化是一種較有效的梯度壓縮方法，相關研究已做進行了實驗和理論論證。

Top-k稀疏化的S-SGD在每次迭代中只需傳輸少量的梯度進行平均也不影響模型收斂或精度。然而，稀疏化後的非零值梯度所在的索引位置在不同的計算節點是不一致的，這使得高效的稀疏化梯度聚合成為一項挑戰。

在稠密梯度上使用基於環形的AllReduce方法（DenseAllReduce）的通信複雜度為 O(P m)，其中 P為計算節點個數，m為參數/梯度的數量。而在Top-k稀疏化中，假設每個節點的梯度稠密度為ρ，即 k = ρ×m ，因為每個節點非零值的對應的索引在不同節點是不一致的。

因此，每次通信需要傳輸 2k個值（梯度值和索引）。採用AllGather對這2k個值進行聚合（簡稱TopKAllReduce）則需要O(kP)的通信複雜度。當擴展到大規模集群時（即P很大），即使k較小也仍會產生顯著的通信開銷。

Top-k稀疏化的主要思想是基於這樣一個事實，即具有較大絕對值的梯度可以為模型收斂做出更多貢獻。因為在Top-k演算法中，即使P個計算節點在聚合後最大可生成 k×P個非零值梯度，但最高絕對值最大的 k個梯度值對於模型更新則更重要。

基於這一觀察，香港浸會大學異構計算實驗室與MassGrid的研究人員，提出了一種有效的Top-k稀疏化方法來解決TopKAllReduce的低效問題。

具體而言，根據不同節點的梯度絕對值來選擇全局的Top-k（簡稱gTop-k）梯度。在本文中，研究人員主要討論使用AllReduce的分散式S-SGD來應用gTop-k稀疏化，但它也適用於基於參數伺服器的分散式S-SGD。

gTop-k可以很好地利用樹結構從所有節點中選擇全局top-k值，並稱之為gTopKAllReduce，而通信複雜性從原來的O(kP)減少到O(klogP)。表1中總結了不同梯度聚合方法的通信複雜度。

在實驗研究及所發表的論文中主要貢獻如下：

研究人員觀察到Top-k稀疏化的聚合結果可以在更新模型前做進一步稀疏化。

在分散式SGD上提出了一種有效的全局Top-k稀疏化演算法，稱為gTop-k S-SGD，用於加速深度神經網路的分散式訓練。

在主流的深度學習框架PyTorch和MPI上實現了gTop-k S-SGD，代碼和參數配置將在GitHub上開源。

在多種神經網路和數據集進行了實驗驗證，實驗結果表明gTop- k S-SGD在低帶寬網路GPU集群上（MassGrid礦機集群）顯著地提高系統的擴展效率。

MassGrid分散式計算網路使用虛擬化技術將全球範圍可用的計算資源虛擬成通用計算設備，可快速部署連接組網進行大規模分散式計算，具有成本低，使用靈活等優勢。本次實驗中MassGrid提供了配置如下的礦機集群：

訓練方法

gTop-k的關鍵思想

在Top-k S-SGD中，每個節點在本地選出k個梯度值，然後所有節點進行聚合得到。研究人員發現並非所有（其非零元素數量為且）都有助於模型收斂。

具體來說，可以進一步稀疏化為，這樣每次模型更新只需要更少數量的非零梯度。換句話說，可以進一步從中選擇top-k個最大絕對值的梯度（表示為）來更新模型，同時保證模型的收斂速度。一個4節點的示例如圖1所示。

圖1 從Top-k演算法中的最多k×P個非零值進一步選擇k個梯度示例

為驗證相比對收斂速度沒有影響，通過訓練ResNet 來進行對比，實驗結果如圖2所示。

圖2 從Top-k演算法中選擇k個梯度進行模型更新的收斂結果

gTopKAllReduce：gTop-k稀疏化的高效AllReduce演算法

從表1可以看到AllGather集合對從不規則索引進行AllReduce操作效率很低，特別是P對通信性能的影響。而新提出的有效演算法的主要目的是減輕變數P對通信開銷的影響。因為最終只需要選擇k個值對模型進行更新，所以在通信過程中，每次只需要傳輸k個非0值。

利用樹狀結構進行兩兩通信，每次只傳輸k個非0值，而接收者則會有2k個值，為下次傳輸也只傳輸k個值，接收者從2k個值中再選出k個值傳遞給下一個接收者。由於兩兩通信是可以同時進行，因此對於P個節點，只需要logP輪通信，一個8節點的示例如圖3所示。

由圖3可以看出，第1個節點在每一輪通信中都會接收另一個節點的k個非0元素，在最後一輪通信後，第1個節點則選出了k個非0元素然後廣播給其他節點，所以整體的通信開銷為：2αlogP 4kβlogP。當機器規模P變大時，gTop-k還能維持較高的擴展性。演算法偽代碼為圖4所示。

圖3 對8個節點，共需要3輪通信，每輪通信只傳輸k個非0值

圖4 gTopKAllReduce演算法偽代碼

實驗結果

香港浸會大學異構計算實驗室與MassGrid的研究人員在32台礦機環境上進行實驗，每台礦機使用一個Nvidia P102-100顯卡。

首先驗證gTop-k S-SGD的收斂性。之後，對三種S-SGD演算法（即基於稠密梯度的S-SGD，Top-k S-SGD和gTop-k S-SGD）的訓練效率進行了比較。

實驗具體硬體配置如表II所示

用於實驗的神經網路配置如表III所示

gTop-k的收斂性能

總體來看，在不同數據集上三種不同類型的DNN的收斂結果表明研究人員提出的gTop-k S-SGD在訓練期間不會損壞模型性能。

圖5 gTop-k S-SGD收斂性能

gTop-k的擴展性能

與S-SGD和Top-k S-SGD相比，在32個計算節點的集群環境上，gTop-k S-SGD比S-SGD快6.7倍，比Top-k S-SGD平均快1.4倍。不同的模型和不同節點數加速比如圖6和表IV所示。

圖6不同節點數的擴展效率對比

表7不同模型的擴展效率對比

局部稀疏化時間（tcompr.）和通信時間（tcommu.）。結果如圖11所示。

一方面，在VGG-16和AlexNet型號的時間細分中通信開銷遠大於計算。因為VGG-16和AlexNet有三個完全連接的層，具有大量參數，而計算速度相對較快。這些也反映出即使使用gTop-k稀疏化，圖6中S-SGD的縮放效率也很低。

另一方面，通信和稀疏化的時間遠小於使用ResNet20和ResNet-50計算的時間，這表明通信計算比率低，因此即使在低帶寬網路上，擴展效率也可高達80％。

此外，應注意梯度稀疏化所用的時間是與VGG-16和AlexNet型號的計算時間相當。主要原因是GPU上的Top-k選擇效率低下，並且在SIMD架構上高度並行化可能並非易事。研究人員將此作為未來的優化方向。

實驗總結

分散式同步隨機梯度下降（S-SGD）已經廣泛用於訓練大規模深度神經網路（DNN），但是它通常需要計算工作者（例如，GPU）之間非常高的通信帶寬以迭代交換梯度。

最近，已經提出了Top-k稀疏化技術來減少工人之間要交換的數據量。Top-k稀疏化可以將大部分梯度歸零，而不會影響模型收斂。

通過對不同DNN的廣泛實驗，這一研究驗證了gTop-k S-SGD與S-SGD具有幾乎一致的收斂性能，並且在泛化性能上只有輕微的降級。

在擴展效率方面，研究人員在具有32個GPU機器的集群上（MassGrid礦機集群）評估gTop-k，這些機器與1 Gbps乙太網互連。

實驗結果表明，該方法比S-SGD實現了2.7-12倍的縮放效率，比現有的Top-k S-SGD提高了1.1-1.7倍。

傳送門

喜歡這篇文章嗎？立刻分享出去讓更多人知道吧！