一種具有混合精度的高度可擴展的深度學習訓練系統(tǒng)

編者按：昨天，騰訊和香港浸會大學的研究人員在arxiv上發(fā)表了一篇文章，介紹了一種具有混合精度的高度可擴展的深度學習訓練系統(tǒng)?；蛟S是覺得這個名字不夠吸引人，他們又在后面加了一個頗具標題黨意味的后綴——“4分鐘訓練ImageNet”。那么這樣的措辭是否夸大其詞了呢？讓我們來讀讀他們的論文。

現(xiàn)狀

現(xiàn)如今，隨著數(shù)據(jù)集和大型深度神經(jīng)網(wǎng)絡規(guī)模不斷擴大，研究人員訓練模型使用的時間也在不斷延長，短則幾天，長則幾周，但過長的訓練時間會給研發(fā)進度帶來阻礙。由于計算資源有限，針對這個問題，現(xiàn)在一種常見解決方案是使用分布式同步隨機梯度下降（SGD），它可以跨硬件作業(yè)，前提是必須給每個GPU分配合理的樣本數(shù)量。

雖然這種做法可以利用系統(tǒng)的總吞吐量和較少的模型更新來加速訓練，但它也存在兩個不可忽視的問題：

較大的mini-batch由于存在泛化誤差，會導致較低的測試精度：

如果增加mini-batch里的樣本個數(shù)，我們確實可以通過計算平均值來減少梯度變化，從而提供更準確的梯度估計，此時模型采用的步長更大，優(yōu)化算法速度也更快。但正如ImageNet training in minutes這篇論文所論證的，一旦mini-batch的大小增加到64K，ResNet-50在測試集上的準確率會從75.4%下降到73.2%，這達不到基線模型的精度要求。

當使用大型GPU集群時，訓練速度并不會隨著GPU數(shù)量的增加而呈線性上升趨勢，尤其是對于計算通信比較高的模型：

訓練模型時，分布式訓練系統(tǒng)會為每個GPU分配訓練任務，然后在每個訓練步之間插入一個梯度聚合步驟，GPU數(shù)量越多，這個梯度聚合步就越容易成為系統(tǒng)瓶頸。假設GPU數(shù)量N固定，如果要提升系統(tǒng)總吞吐量T，我們就要同時提高單個GPU的吞吐量S和它的縮放效率e，但這兩者提升需要額外的算力資源，這就和N固定有矛盾。

實驗結(jié)果

在介紹研究成果前，我們先來看看最引人注意的“4分鐘訓練ImageNet”。

根據(jù)論文實驗部分的內(nèi)容，研究人員選取的模型是AlexNet和ResNet-50，它們各自代表一種典型的CNN。AlexNet的參數(shù)數(shù)量是ResNet-50的2.5倍，而ResNet-50的計算卻是前者的5.6倍。因此它們的瓶頸分別是通信和計算，這正代表上節(jié)提出的兩個問題。

上表是兩個模型的訓練結(jié)果和對比，可以發(fā)現(xiàn)，在ImageNet數(shù)據(jù)集上，研究人員用4分鐘訓練好AlexNet，又用6.6分鐘訓練好ResNet-50，batch size非常大，但精度卻和其他模型沒什么區(qū)別。從數(shù)據(jù)角度看，這確實是個歷史性的突破。

而根據(jù)騰訊機智團隊自己的介紹，在這之前，業(yè)界最好的水平來自：日本Perferred Network公司Chainer團隊，他們用15分鐘就訓練好了ResNet-50；UC Berkely等高校的團隊，他們用11分鐘訓練好AlexNet。相較之下，騰訊和香港浸會大學的這個成果創(chuàng)造了AI訓練的世界新紀錄。

但顯然，他們在寫論文標題的時候也漏掉了重要內(nèi)容，就是這個速度背后是2000多塊GPU，相信這個真相會讓一些研究人員興奮，也會讓大批學者和實驗室感到內(nèi)心拔涼。

研究成果概述

關(guān)于論文技術(shù)的詳細細節(jié)，騰訊技術(shù)工程官方已經(jīng)有長文分析，所以這里只根據(jù)論文內(nèi)容的一點簡介（才不說是寫完才發(fā)現(xiàn)人家已經(jīng)發(fā)了呢）。

這篇論文為為密集GPU集群構(gòu)建了一個高度可擴展的深度學習訓練系統(tǒng)，它很好地解決了上述兩個問題，下圖是它結(jié)構(gòu)概覽：

從圖中我們可以看到，這個系統(tǒng)可以大致分為三個模塊：輸入管道模塊、訓練模塊和通信模塊。

輸入管道負責在當前步驟完成之前就為下一步提供數(shù)據(jù)，它使用pipelining來最小化CPU和GPU的閑置時間。

訓練模塊包括模型構(gòu)建和變量管理。在這個模塊中，研究人員結(jié)合了各類優(yōu)化方法，如使用混合精度訓練前向/反向傳播和用LARS更新模型。

通信模塊使用tensor fusion和論文提出的混合Allreduce，根據(jù)張量大小和GPU集群大小優(yōu)化縮放效率。

1. 提出了一種混合精度訓練方法，可以顯著提高單個GPU的訓練吞吐量而不會降低精度。

之前，Micikevicius等人已經(jīng)在研究中提出過在訓練階段使用半精度（FP16）有助于降低內(nèi)存壓力并增加計算吞吐量的想法，前者可以通過把相同數(shù)量的值儲存進更少的bit來實現(xiàn)，后者則可以降低數(shù)學精度，讓處理器提供更高的吞吐量。

而Yang You等人提出了一種為分布式培訓提供更大mini-batch的算法——LARS（自適應速率縮放），它會為每一層引入局部學習率，也就是用系數(shù)η加權(quán)的L2正則化權(quán)重和梯度權(quán)重的比率，能大幅度提高大batch size場景下的訓練精度。

這兩個成果非常互補，但它們不能直接結(jié)合使用，因為會導致梯度消失。為了融合兩種思想，研究人員做出的改進是用LARS進行混合精確訓練，如下圖所示，當進行前向傳播和反向傳播時，系統(tǒng)先把參數(shù)和數(shù)據(jù)轉(zhuǎn)成半精度浮點數(shù)，然后再做訓練，而訓練權(quán)重和梯度時，參數(shù)和數(shù)據(jù)則是單精度浮點數(shù)。

2. 提出了一種針對超大型mini-batch（最大64k）的優(yōu)化方法，可以在ImageNet數(shù)據(jù)集上訓練CNN模型而不會降低精度。

模型架構(gòu)改進是提高模型性能的一種常見手段，在論文中，研究人員從以下兩個方面改進了模型架構(gòu)：1)消除偏差和batch normalization的權(quán)重衰減；2）為AlexNet增加了一個batch normalization層。

除了這一點，深度學習中耗時占比較重的還有超參數(shù)調(diào)整。為了優(yōu)化這一過程，研究人員的思路是：

參數(shù)步長由粗到細：調(diào)優(yōu)參數(shù)值先以較大步長進行劃分，可以減少參數(shù)組合數(shù)量，當確定大的最優(yōu)范圍之后再逐漸細化調(diào)整。

低精度調(diào)參：分析相關(guān)數(shù)據(jù)，放大低精度表示邊緣數(shù)值，保證參數(shù)的有效性。

初始化數(shù)據(jù)的調(diào)參：根據(jù)輸入輸出通道數(shù)的范圍來初始化初始值，一般以輸入通道數(shù)較為常見；對于全連接網(wǎng)絡層則采用高斯分布即可；對于shortcut的batch norm，參數(shù)gamma初始化為零（也可以先訓練一個淺層網(wǎng)絡，再通過參數(shù)遞進初始化深層網(wǎng)絡參數(shù)）。

3. 提出了一種高度優(yōu)化的allreduce算法，使用這種算法后，相比NCCL計算框架，AlexNet和ResNet-50在包含1024個Tesla P40 GPU的集群上的訓練速度分別提高了3倍和11倍。

在張量足夠多的情況下，Ring Allreduce可以最大化利用網(wǎng)絡，但工作效率和速度都不如張量少的情況。針對這種現(xiàn)象，研究人員利用分層同步和梯度分段融合優(yōu)化Ring Allreduce

分層同步與Ring Allreduce有機結(jié)合：對集群內(nèi)GPU節(jié)點進行分組，減少GPU數(shù)量對整體訓練用時的影響。

梯度融合，多次梯度傳輸合并為一次：根據(jù)具體模型設置合適的Tensor size閾值，將多次梯度傳輸合并為一次，同時超過閾值大小的Tensor不再參與融合；這樣可以防止Tensor過度碎片化，從而提升了帶寬利用率，降低了傳輸耗時。

GDR技術(shù)加速Ring Allreduce：在前述方案的基礎上，將GDR技術(shù)應用于跨節(jié)點Ring，這減少了主存和顯存之間的Copy操作，同時為GPU執(zhí)行規(guī)約計算提供了便利。

這三大成果的直接效果是在不降低分類準確率的同時，把AlexNet和ResNet-50訓練時所用的mini-batch size擴大至64K。同時，通過優(yōu)化All-reduce算法，并讓系統(tǒng)支持半精度訓練，研究人員最后構(gòu)建了一個高吞吐量分布式深度學習訓練系統(tǒng)，可以在GPU數(shù)量N不變的情況下，提高單個GPU性能S和縮放效率e。