提升12倍！香港浸會大學(xué)與MassGrid發(fā)布低帶寬高效AI訓(xùn)練新算法

安妮 2019-05-09 13:11:58 來源：量子位

允中 發(fā)自 凹非寺

量子位 報道 | 公眾號 QbitAI

隨著訓(xùn)練數(shù)據(jù)量的增加和深度神經(jīng)網(wǎng)絡(luò)（DNN）的日益復(fù)雜，分布式計(jì)算環(huán)境（如GPU集群）被廣泛采用，以加速DNN的訓(xùn)練。分布式計(jì)算網(wǎng)絡(luò)在機(jī)器學(xué)習(xí)方面的瓶頸在于節(jié)點(diǎn)之間的數(shù)據(jù)傳輸效率，那如何在這一網(wǎng)絡(luò)下高效地進(jìn)行AI訓(xùn)練？

2018年，香港浸會大學(xué)異構(gòu)計(jì)算實(shí)驗(yàn)室與MassGrid合作，通過研究提出一種可用于低帶寬網(wǎng)絡(luò)的全局Top-k稀疏化的分布式同步SGD算法，并通過實(shí)驗(yàn)論證出在低帶寬網(wǎng)絡(luò)下也能高效進(jìn)行AI訓(xùn)練。目前實(shí)驗(yàn)結(jié)果及相關(guān)論文已被ICDCS workshop收錄。

數(shù)據(jù)并行的分布式同步隨機(jī)梯度下降（S-SGD）方法是訓(xùn)練大規(guī)模神經(jīng)網(wǎng)絡(luò)常用的優(yōu)化器之一。與單節(jié)點(diǎn)的SGD相比，S-SGD將工作負(fù)載分配給多個計(jì)算節(jié)點(diǎn)以加速訓(xùn)練，但它也引入了在每次迭代中交換模型參數(shù)或梯度的通信開銷。

舉例說明

假設(shè)有P個節(jié)點(diǎn)用S-SGD訓(xùn)練DNN模型。在每次迭代中，所有計(jì)算節(jié)點(diǎn)都會采用不同的小批量(mini-batch)數(shù)據(jù)來并行計(jì)算模型的梯度。然后，對每個節(jié)點(diǎn)的梯度進(jìn)行平均后來更新模型，這便引入較大的通信開銷。

由于計(jì)算節(jié)點(diǎn)的加速器（如GPU和TPU）的計(jì)算能力比網(wǎng)絡(luò)速度的增長快得多，網(wǎng)絡(luò)通信性能通常成為訓(xùn)練的性能瓶頸，特別是當(dāng)通信與計(jì)算比率很高時。

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

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

• 一方面，許多大批量SGD技術(shù)已經(jīng)提出了一些優(yōu)化策略來提高mini-batch而不會丟失模型準(zhǔn)確性。
• 另一方面，研究人員也已經(jīng)提出了梯度稀疏化，量化和壓縮等方法，在不影響收斂速度的情況下顯著減小交換梯度的數(shù)量以降低通信量。

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

Top-k稀疏化是一種較有效的梯度壓縮方法，相關(guān)研究已做進(jìn)行了實(shí)驗(yàn)和理論論證。

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

在稠密梯度上使用基于環(huán)形的AllReduce方法（DenseAllReduce）的通信復(fù)雜度為 O(P + m)，其中 P為計(jì)算節(jié)點(diǎn)個數(shù)，m為參數(shù)/梯度的數(shù)量。而在Top-k稀疏化中，假設(shè)每個節(jié)點(diǎn)的梯度稠密度為ρ，即 k = ρ×m ，因?yàn)槊總€節(jié)點(diǎn)非零值的對應(yīng)的索引在不同節(jié)點(diǎn)是不一致的。

因此，每次通信需要傳輸 2k個值（梯度值和索引）。采用AllGather對這2k個值進(jìn)行聚合（簡稱TopKAllReduce）則需要O(kP)的通信復(fù)雜度。當(dāng)擴(kuò)展到大規(guī)模集群時（即P很大），即使k較小也仍會產(chǎn)生顯著的通信開銷。

Top-k稀疏化的主要思想是基于這樣一個事實(shí)，即具有較大絕對值的梯度可以為模型收斂做出更多貢獻(xiàn)。因?yàn)樵赥op-k算法中，即使P個計(jì)算節(jié)點(diǎn)在聚合后最大可生成 k×P個非零值梯度，但最高絕對值最大的 k個梯度值對于模型更新則更重要。

基于這一觀察，香港浸會大學(xué)異構(gòu)計(jì)算實(shí)驗(yàn)室與MassGrid的研究人員，提出了一種有效的Top-k稀疏化方法來解決TopKAllReduce的低效問題。

具體而言，根據(jù)不同節(jié)點(diǎn)的梯度絕對值來選擇全局的Top-k（簡稱gTop-k）梯度。在本文中，研究人員主要討論使用AllReduce的分布式S-SGD來應(yīng)用gTop-k稀疏化，但它也適用于基于參數(shù)服務(wù)器的分布式S-SGD。

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

在實(shí)驗(yàn)研究及所發(fā)表的論文中主要貢獻(xiàn)如下：

1. 研究人員觀察到Top-k稀疏化的聚合結(jié)果可以在更新模型前做進(jìn)一步稀疏化。
2. 在分布式SGD上提出了一種有效的全局Top-k稀疏化算法，稱為gTop-k S-SGD，用于加速深度神經(jīng)網(wǎng)絡(luò)的分布式訓(xùn)練。
3. 在主流的深度學(xué)習(xí)框架PyTorch和MPI上實(shí)現(xiàn)了gTop-k S-SGD，代碼和參數(shù)配置將在GitHub上開源。
4. 在多種神經(jīng)網(wǎng)絡(luò)和數(shù)據(jù)集進(jìn)行了實(shí)驗(yàn)驗(yàn)證，實(shí)驗(yàn)結(jié)果表明gTop- k S-SGD在低帶寬網(wǎng)絡(luò)GPU集群上（MassGrid礦機(jī)集群）顯著地提高系統(tǒng)的擴(kuò)展效率。
5. MassGrid分布式計(jì)算網(wǎng)絡(luò)使用虛擬化技術(shù)將全球范圍可用的計(jì)算資源虛擬成通用計(jì)算設(shè)備，可快速部署連接組網(wǎng)進(jìn)行大規(guī)模分布式計(jì)算，具有成本低，使用靈活等優(yōu)勢。本次實(shí)驗(yàn)中MassGrid提供了配置如下的礦機(jī)集群：

訓(xùn)練方法

gTop-k的關(guān)鍵思想

在Top-k S-SGD中，每個節(jié)點(diǎn)在本地選出k個梯度值，然后所有節(jié)點(diǎn)進(jìn)行聚合得到。研究人員發(fā)現(xiàn)并非所有（其非零元素?cái)?shù)量為且）都有助于模型收斂。

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

△?圖1 從Top-k算法中的最多k×P個非零值進(jìn)一步選擇k個梯度示例

為驗(yàn)證相比對收斂速度沒有影響，通過訓(xùn)練ResNet 來進(jìn)行對比，實(shí)驗(yàn)結(jié)果如圖2所示。

△?圖2 從Top-k算法中選擇k個梯度進(jìn)行模型更新的收斂結(jié)果

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

從表1可以看到AllGather集合對從不規(guī)則索引進(jìn)行AllReduce操作效率很低，特別是P對通信性能的影響。而新提出的有效算法的主要目的是減輕變量P對通信開銷的影響。因?yàn)樽罱K只需要選擇k個值對模型進(jìn)行更新，所以在通信過程中，每次只需要傳輸k個非0值。

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

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

△?圖3 對8個節(jié)點(diǎn)，共需要3輪通信，每輪通信只傳輸k個非0值

△?圖4 gTopKAllReduce算法偽代碼

實(shí)驗(yàn)結(jié)果

香港浸會大學(xué)異構(gòu)計(jì)算實(shí)驗(yàn)室與MassGrid的研究人員在32臺礦機(jī)環(huán)境上進(jìn)行實(shí)驗(yàn)，每臺礦機(jī)使用一個Nvidia P102-100顯卡。

首先驗(yàn)證gTop-k S-SGD的收斂性。之后，對三種S-SGD算法（即基于稠密梯度的S-SGD，Top-k S-SGD和gTop-k S-SGD）的訓(xùn)練效率進(jìn)行了比較。

實(shí)驗(yàn)具體硬件配置如表II所示

用于實(shí)驗(yàn)的神經(jīng)網(wǎng)絡(luò)配置如表III所示

gTop-k的收斂性能

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

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

gTop-k的擴(kuò)展性能

與S-SGD和Top-k S-SGD相比，在32個計(jì)算節(jié)點(diǎn)的集群環(huán)境上，gTop-k S-SGD比S-SGD快6.7倍，比Top-k S-SGD平均快1.4倍。不同的模型和不同節(jié)點(diǎn)數(shù)加速比如圖6和表IV所示。

△?圖6不同節(jié)點(diǎn)數(shù)的擴(kuò)展效率對比

△?表7不同模型的擴(kuò)展效率對比

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

一方面，在VGG-16和AlexNet型號的時間細(xì)分中通信開銷遠(yuǎn)大于計(jì)算。因?yàn)閂GG-16和AlexNet有三個完全連接的層，具有大量參數(shù)，而計(jì)算速度相對較快。這些也反映出即使使用gTop-k稀疏化，圖6中S-SGD的縮放效率也很低。

另一方面，通信和稀疏化的時間遠(yuǎn)小于使用ResNet20和ResNet-50計(jì)算的時間，這表明通信計(jì)算比率低，因此即使在低帶寬網(wǎng)絡(luò)上，擴(kuò)展效率也可高達(dá)80％。

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

實(shí)驗(yàn)總結(jié)

分布式同步隨機(jī)梯度下降（S-SGD）已經(jīng)廣泛用于訓(xùn)練大規(guī)模深度神經(jīng)網(wǎng)絡(luò)（DNN），但是它通常需要計(jì)算工作者（例如，GPU）之間非常高的通信帶寬以迭代交換梯度。

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

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

在擴(kuò)展效率方面，研究人員在具有32個GPU機(jī)器的集群上（MassGrid礦機(jī)集群）評估gTop-k，這些機(jī)器與1 Gbps以太網(wǎng)互連。

實(shí)驗(yàn)結(jié)果表明，該方法比S-SGD實(shí)現(xiàn)了2.7-12倍的縮放效率，比現(xiàn)有的Top-k S-SGD提高了1.1-1.7倍。

傳送門

論文原文鏈接：https://arxiv.org/abs/1901.04359

— 完 —

誠摯招聘

量子位正在招募編輯/記者，工作地點(diǎn)在北京中關(guān)村。期待有才氣、有熱情的同學(xué)加入我們！相關(guān)細(xì)節(jié)，請?jiān)诹孔游还娞?QbitAI)對話界面，回復(fù)“招聘”兩個字。

量子位 QbitAI · 頭條號簽約作者

?’?’ ? 追蹤AI技術(shù)和產(chǎn)品新動態(tài)