深度神經(jīng)網(wǎng)絡(luò)訓(xùn)練速度越來越快已經(jīng)不是新鮮事，但是，將ImageNet訓(xùn)練時間降低到200秒級別仍然讓人震撼！近日，索尼發(fā)布新的方法，在ImageNet數(shù)據(jù)集上，使用多達(dá)2176個GPU，在224秒內(nèi)成功訓(xùn)練了ResNet-50，刷新了紀(jì)錄。

隨著用于深度學(xué)習(xí)的數(shù)據(jù)集和深度神經(jīng)網(wǎng)絡(luò)模型的規(guī)模增大，訓(xùn)練模型所需的時間也在增加具有數(shù)據(jù)并行性的大規(guī)模分布式深度學(xué)習(xí)可以有效縮短訓(xùn)練時間。

然而，由于大型 mini-batch 訓(xùn)練的不穩(wěn)定性和梯度同步的開銷，將分布式深度學(xué)習(xí)擴(kuò)展到大規(guī)模的GPU集群級別很有挑戰(zhàn)性。

日本索尼公司的Hiroaki Mikami等人近日提出一種新的大規(guī)模分布式訓(xùn)練方法，通過控制batch size解決了大型mini-batch訓(xùn)練的不穩(wěn)定性，用2D-Torus all-reduce解決了梯度同步的開銷。

具體而言，2D-Torus all-reduce將GPU排列在一個邏輯2D網(wǎng)格中，并以不同的方向執(zhí)行一系列操作。

這兩種技術(shù)都是基于索尼的神經(jīng)網(wǎng)絡(luò)庫NNL（Neural Network Libraries）實(shí)現(xiàn)的。最終，索尼的研究人員在224秒內(nèi)（使用多達(dá)2176個GPU）成功訓(xùn)練了ImageNet/ResNet-50，并在ABCI 集群上沒有明顯的精度損失。

在ImageNet數(shù)據(jù)集上訓(xùn)練ResNet-50是用于測量深度學(xué)習(xí)分布式學(xué)習(xí)速度的一般行業(yè)基準(zhǔn)，該研究刷新了這個基準(zhǔn)的速度。

論文地址：

https://arxiv.org/pdf/1811.05233.pdf

224秒！刷新深度學(xué)習(xí)紀(jì)錄

在大型GPU集群中，大規(guī)模分布式深度學(xué)習(xí)存在兩個技術(shù)問題。第一個問題large mini-batch訓(xùn)練造成的收斂精度下降。第二個問題是GPU間梯度同步的通信開銷。解決這兩個問題需要一種新的分布式處理方法。

近年來，許多研究人員提出了多種方案來解決這兩個問題（見原文參考文獻(xiàn)）。這些工作利用ImageNet/ResNet-50訓(xùn)練來衡量訓(xùn)練效果。ImageNet/ResNet-50分別是最流行的數(shù)據(jù)集和最流行的DNN模型，用于對大規(guī)模分布式深度學(xué)習(xí)進(jìn)行基準(zhǔn)測試。

表1比較了近期一些工作的訓(xùn)練時間和top-1驗(yàn)證精度。其中，F(xiàn)acebook使用256個Tesla P100 GPU，在1小時內(nèi)訓(xùn)練完ResNet-50，是加速了這一任務(wù)的著名研究。

表1：ImageNet/ResNet-50訓(xùn)練時間及top-1 -crop驗(yàn)證精度

之前的其他一些業(yè)界最好水平來自：

日本Perferred Network公司Chainer團(tuán)隊(duì)，15分鐘訓(xùn)練好ResNet-50 [5]

騰訊機(jī)智團(tuán)隊(duì)，6.6分鐘訓(xùn)練好ResNet-50 [6]

索尼團(tuán)隊(duì)的研究著重于解決大型mini-batch訓(xùn)練的不穩(wěn)定性和梯度同步開銷，他們使用2176 個Tesla V100 GPU，將訓(xùn)練時間縮短至224秒，驗(yàn)證精度為75.03％。

研究人員還嘗試在不造成明顯精度損失的情況下提高GPU scaling效率，使用1088個Tesla V100 GPU將GPU scaling效率提高到91.62％（表2）。

表2：ImageNet/ResNet-50訓(xùn)練的GPU scaling 效率

2大方法解決不穩(wěn)定問題

大規(guī)模分布式訓(xùn)練有兩個主要問題：大型mini-batch訓(xùn)練的不穩(wěn)定性和同步通信的開銷。

眾所周知， large mini-batch訓(xùn)練是不穩(wěn)定的，會產(chǎn)生泛化差距。

數(shù)據(jù)并行分布式訓(xùn)練需要在每個訓(xùn)練迭代之間增加一個步驟，以便在參與的GPU之間同步和平均梯度。這個步驟是使用一個all-reduce的集合操作來實(shí)現(xiàn)的。在一個大型GPU集群中，all-reduce集合操作的開銷使得線性縮放變得非常具有挑戰(zhàn)性。

針對這兩個問題，我們使用Batch Size控制技術(shù)來解決不穩(wěn)定問題，并開發(fā)了2D-Torus all-reducing方案，有效地跨GPU交換梯度。

Batch Size Control

以往的工作已經(jīng)證明，在訓(xùn)練期間逐漸增加總的mini-batch size可以減少大型 mini-batch訓(xùn)練的不穩(wěn)定性。直觀地說，隨著訓(xùn)練的損失情況變得“平坦”而增加批大小有助于避免局部最小值。

在這項(xiàng)工作中，我們采用 Batch Size Control來減少精度下降， batch size超過了32K。在訓(xùn)練期間采用了預(yù)定的batch-size來更改調(diào)度。

2 D-Torus All- reduce

有效的通信拓?fù)鋵τ跍p少集體操作的通信開銷至關(guān)重要。

為了解決這個問題，我們開發(fā)了2D-Torus all-reduce。2D-Torus拓?fù)浣Y(jié)構(gòu)如圖1所示。集群中的GPU排列在2D網(wǎng)格中。在2D-torus拓?fù)渲?，all-reduce由三個步驟組成：reduce-scatter，all-reduce和all-gather。

圖1：2D-Torus拓?fù)溆伤胶痛怪狈较虻亩鄠€環(huán)組成。

2D-Torus all-reduce的示例如圖2所示。

圖2：在2x2網(wǎng)格中，一個4-GPU集群的2D-Torus all-reduce步驟

評估：實(shí)驗(yàn)設(shè)置和訓(xùn)練設(shè)置

實(shí)驗(yàn)設(shè)置

軟件：使用神經(jīng)網(wǎng)絡(luò)庫（NNL）及其CUDA擴(kuò)展，作為DNN訓(xùn)練框架。通信庫使用NCCL和OpenMPI。2D-Torus all-reduce在NCCL上實(shí)現(xiàn)。以上軟件打包在Singularity容器中，用于運(yùn)行分布式DNN訓(xùn)練。

硬件：使用AI橋接云基礎(chǔ)設(shè)施（ABCI）作為GPU集群。ABCI是日本先進(jìn)工業(yè)科技研究所（AIST）運(yùn)營的GPU集群。它包括1088個節(jié)點(diǎn)，每個節(jié)點(diǎn)有4個NVIDIA Tesla V100 GPU，2個Xeon Gold 6148處理器，376 GB內(nèi)存。同一節(jié)點(diǎn)的GPU由NVLink互連，而節(jié)點(diǎn)由2個InfiniBand EDR互連。

數(shù)據(jù)集和模型：使用ImageNet數(shù)據(jù)集。使用ResNet-50作為DNN模型。模型中的所有層都由[9]中描述的值初始化。

訓(xùn)練設(shè)置：

使用LARS [9]，系數(shù)為0.01，eps為1e-6更新權(quán)重。學(xué)習(xí)率（LR）通過以下公式計(jì)算：

用以下公式計(jì)算出總的mini-batch size和學(xué)習(xí)率。

我們還采用了[15]中介紹的混合精度訓(xùn)練。前向/后向計(jì)算和同步梯度的通信在半精度浮點(diǎn)（FP16）中進(jìn)行。

我們調(diào)整每個worker和總batch size，如表3所示，直到將總batch size增到最大。通過增加GPU的數(shù)量（Exp.1到Exp.4）來嘗試提高最大總batch size。

表3：per-worker/total mini-batch size

但是，當(dāng)使用超過2176個GPU時，訓(xùn)練效率變低了。因此，由于這個問題， Exp. 5 和Exp. 6僅使用2176個GPU。

表4：實(shí)驗(yàn)中使用的2D-Torus拓?fù)涞木W(wǎng)格尺寸。

結(jié)果：精度無損失，訓(xùn)練時間只需224秒

我們在224秒內(nèi)完成了ResNet-50的訓(xùn)練，沒有明顯的精度損失，如表5所示。

表5：Top-1 1-crop 驗(yàn)證精度和訓(xùn)練時間

訓(xùn)練誤差曲線與參考曲線非常相似（圖3）。雖然最大的batch size可以增加到119K也不會造成明顯的精度損失，但進(jìn)一步增大會使精度降低約0.5％（表5中的實(shí)驗(yàn)6）。

圖3：訓(xùn)練誤差曲線

我們描述了與單個節(jié)點(diǎn)（4個GPU）相比的訓(xùn)練速度和GPU縮放效率。

表6顯示了當(dāng)每個worker的批大小設(shè)置為32時的GPU數(shù)量和訓(xùn)練吞吐量。雖然當(dāng)使用超過2176個GPU時，GPU scaling效率降低到70％，但當(dāng)使用1088 GPU時，scaling效率超過了90％。

在之前的研究[6]中，當(dāng)使用1024個Tesla P40，每個worker的批大小設(shè)置為32時，GPU scaling效率為87.9％。因此，與之前的研究相比，我們的通信方案通過更快、更多的GPU實(shí)現(xiàn)了更高的GPU scaling效率。

表6：2D-Torus all-reduce的訓(xùn)練吞吐量和scaling效率

結(jié)論

大規(guī)模分布式深度學(xué)習(xí)是減少DNN訓(xùn)練時間的有效方法。我們采用了多種技術(shù)來減少精度下降，同時在使用一個龐大的GPU集群進(jìn)行訓(xùn)練的同時保持了較高的GPU scaling效率。

這些技術(shù)是用神經(jīng)網(wǎng)絡(luò)庫（NNL）實(shí)現(xiàn)的，我們使用了2176個 Tesla V100 GPU，訓(xùn)練時間224秒，驗(yàn)證精度75.03%。我們還通過1088個Tesla V100 GPU達(dá)到了90%以上的GPU擴(kuò)展效率。