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

現(xiàn)狀

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

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

較大的mini-batch由于存在泛化誤差，會(huì)導(dǎo)致較低的測(cè)試精度：

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

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

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

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

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

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

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

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

但顯然，他們?cè)趯懻撐臉?biāo)題的時(shí)候也漏掉了重要內(nèi)容，就是這個(gè)速度背后是2000多塊GPU，相信這個(gè)真相會(huì)讓一些研究人員興奮，也會(huì)讓大批學(xué)者和實(shí)驗(yàn)室感到內(nèi)心拔涼。

研究成果概述

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

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

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

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

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

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

1. 提出了一種混合精度訓(xùn)練方法，可以顯著提高單個(gè)GPU的訓(xùn)練吞吐量而不會(huì)降低精度。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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