2分31秒,騰訊雲創造128卡訓練ImageNet新記錄

新智元報導

來源：騰訊雲

編輯：白峰

【新智元導讀】8月21日，騰訊雲正式對外宣布成功創造了128卡訓練ImageNet業界新記錄，以2分31秒的成績一舉刷新了這個領域的世界記錄。

刷新世界紀錄：2分31秒完成 ImageNet訓練

基於騰訊公有雲25Gbps的VPC網絡環境，使用128塊V100，藉助Light大規模分布式多機多卡訓練框架，在2分31秒內訓練 ImageNet 28個epoch，TOP5精度達到93%，創造128卡訓練imagenet 業界新記錄。

若跨機網絡改為RoCE，則訓練時間進一步減少到2分2秒。之前業界同樣使用128塊V100的最好結果是友商的2分38秒。

28個epoch分別是：

模型越來越大，算力需求暴漲

在AlexNet[3]網絡模型出現後的過去幾年中，ResNet[4]、Transformer[5]、BERT[6]等優秀模型相繼出現，深度學習有了長足的發展和進步，尤其是在圖像、語音、機器翻譯、自然語言處理等領域帶來了跨越式提升。

在AlphaGo使用深度學習方法戰勝世界圍棋冠軍李世石之後，近期擁有1750億海量參數，堪稱「通用語言智能的希望」——GPT-3[7]的出世，再一次點燃了大家對人工智慧未來的期望。但與之相伴的也有很多問題：

數據量大。隨著網際網路上可使用數據資源愈加豐富，模型訓練使用的數據集往往越來越大。GPT-3模型訓練則使用了高達45TB的數據量，使得多輪訓練時數據讀取成為非常耗時的部分。

計算模型複雜。深度網絡的一個特點就是結構越深、越複雜，所表達的特徵就越豐富，在此思想下，最新的網絡結構越來越複雜，從AlexNet的8層，VGG-19的19層，ResNet-50的50層，到Inception-ResNet-V2的467層和ResNet-1000的1202層等。

參數量大。深度神經網絡由於層次很多，參數量往往很大。ResNet-50有2500萬參數量，AlexNet有6200萬的參數量，而VGG-16參數量則達到1.38億，BERT-Large參數量高達3億，而GPT-3語言模型的參數量甚至超過1750億。

超參數範圍廣泛。隨著模型複雜度的提升，模型中可供調節的超參數數量及數值範圍也在增多。例如，在CIFAR-10數據集上訓練的ResNet模型有16個可調的超參數，當多數超參數的取值為連續域的情況下，如此少量的超參數仍然可能造成組合爆炸。

訓練時間長。隨著深度學習模型越來越複雜，對算力的需求也越來越高(2012年到2018年，對算力的需求增長2000倍+ )。ResNet-50 訓練 ImageNet 128 萬張圖片 90 個 epoch 可以達到 76.5% 的 Top-1 精度，用一塊 V100 的 GPU 需要大概1天多的時間，而各大廠都分別施展了各自的大規模分布式訓練大法來縮短該標杆的訓練時間。騰訊機智團隊在2018年7月在ImageNet數據集上，用4分鐘訓練好AlexNet，6.6分鐘訓練好ResNet50，創造了當時AI訓練世界新紀錄[1]。

如何做高性能AI訓練和計算，關乎到AI生產研發效率，對於AI產品的迭代效率和成功至關重要。高效訓練的一個baseline是在更短時間內做ImageNet訓練。騰訊雲希望AI技術能成為普惠的技術幫助人們生活更便捷，通過向客戶提供高效的AI生產工具，AI高效訓練可以觸手可及。

但是要做好高效的AI生產工具，需要做大量的工作，也充滿挑戰。接下來將從具體挑戰，聯合團隊的解決方案，及最終效果幾個方面來陳述，並最終通過2分31秒訓練ImageNet這個案例為典型case來呈現成果。

多機多卡難擴展，調參耗時難收斂

數據供給制約計算

深度學習訓練數據輸入管道包括以下三個部分：從硬碟上讀取數據-解析與數據預處理-拷貝到GPU，數據輸入管道與計算部分並行執行，為保證GPU能專心來進行模型訓練，數據輸入部分使用 CPU來進行。但隨著GPU性能的提高，計算時間越來越短，CPU，內存，網絡帶寬開始成為訓練時的瓶頸，數據供給也成為訓練效率的關鍵一環。

TCP網絡下的多機多卡擴展性差

隨著GPU硬體的不斷升級，GPU的計算速度一直在提升，加上XLA和混合精度等技術也大大提升了單卡的計算速度。NVLink2使單機內部的AllReduce速度有了大幅提升，超過了120GB/s。但騰訊雲的VPC網絡只有25Gbps，與單機內的通信帶寬有著明顯的差異。再加上TCP環境下，跨機通信的數據需要從顯存拷到主存，再通過CPU去收發數據，除了實際網絡收發數據的延遲，還有很大的額外延遲，所以通信時間更長了。計算時間短加上通信時間長，使得多機多卡的擴展性受到了很大的挑戰。

大batch收斂難

為了充分利用大規模集群算力以達到提升訓練速度的目的，人們不斷的提升訓練的batch size，這是因為更大的batch size允許在擴展GPU數量的同時不降低每個GPU的計算負載。

然而，只增大batch size會對精度帶來影響和損失。這是因為在大batch size（相對於訓練樣本數）情況下，樣本隨機性降低，梯度下降方向趨於穩定，訓練就由SGD向GD趨近，這導致模型更容易收斂於初始點附近的某個局部最優解，從而抵消了計算力增加帶來的好處。如何既增大batch size，又不降低精度，是另一個重大挑戰。

超參數選擇多

超參數的選擇對模型最終的效果有極大的影響，但由於超參較多，而每一個超參分布範圍較廣，使得超參調優的耗時較長，特別是針對ImageNet這種大數據集的情況。如何快速尋找到較優解，是另一個重大挑戰。

解決方案

聯合團隊研發了 Light大規模分布式多機多卡訓練框架來進行高效訓練，並將能力平臺化。

單機訓練速度優化

1）分布式緩存與數據預取

由於訓練數據一般保存在分布式存儲，而分布式存儲機器和訓練機器並不在一個集群。為了加速訪問遠端存儲的數據，團隊利用GPU母機的SSD盤/內存，在訓練過程中為訓練程序提供數據預取和緩存。

2）自動調整最優數據預處理線程數

在預設的運行程序中，每個進程會分配大量線程用於數據預處理。僅僅在數據預處理上，單機8卡就會分配數百個線程。大量的線程的相互搶佔導致了cpu運行效率低下，同時也會與tensorflow的OP調度線程爭搶而出現OP無法正常並行的問題。

因此團隊根據經驗和採集運行時信息自動設置最優數據預處理線程數，降低cpu的切換負擔，同時也讓數據預處理可以和gpu計算並行。

3）本地預解碼圖片緩存

對於小圖片，由於每個迭代步的計算時間少，單位時間需要處理的圖片變多，cpu負載仍然過大，導致性能較差。經過分析後發現，JPEG圖片解碼為制約性能的主要部分。為了解決這個問題，團隊預先將數據集中的JPEG圖片解碼，計算時直接加載解碼數據，緩存在內存中，從而直接去除這部分計算壓力。

LightCC高效擴展多機訓練

1）自適應梯度融合技術優化通信時間

針對小數據塊通信問題，通過將小數據塊融合成大數據塊的方式，減少了通信次數，降低了通信延遲，提升了通信效率。但當應用了XLA技術後，由於TensorFlow中XlaOp的同步性，使得通信需等待編譯優化所融合的所有計算完成，導致計算通信無法很好的overlap，為解決這個問題，團隊進一步提出了自適應梯度融合技術，根據編譯優化的結果自適應的選擇梯度融合的閾值，從而達到更好的overlap。

2）層級通信+多流提升帶寬利用率

TCP網絡下全局AllReduce，除了網絡收發包的時間，還有很大的額外延時，導致帶寬利用率不高。

針對這個問題，團隊使用了2D通信[2]和多流來提升網絡收發的並發性，從而提升帶寬的利用率。以單機有8塊卡的機型為例，2D通信在TCP網絡下有不錯的效果，主要是因為可以8塊卡同時做跨機通信，從而競爭帶寬資源，帶寬空閒的時間更少，而且跨機通信時建立的環上的節點數只有總節點數的1/8，所以每塊卡做跨機通信的次數比全局AllReduce時少得多。多流又使得多個梯度的2D通信可以形成流水線，既某個梯度在做單機內通信時沒有利用到網絡帶寬，這時其它梯度的跨機通信可以填補網絡帶寬的間隙。

ResNet-50 v1.5的參數量大概是25M個FP32，如果通信時轉成FP16，則通信的梯度大小大概是50MB，但因為AllReduce的通信量大概是梯度大小的2倍，所以通信量大概是100MB。通信時間如下圖所示。

3）層級topk壓縮通信算法減少通信量，突破帶寬瓶頸

因為網絡帶寬得到充分利用之後，通信還是瓶頸，為進一步提高擴展性，團隊引入了梯度壓縮通信技術[8]。梯度壓縮通信的做法是每張卡對要通信的梯度做Topk壓縮，只剩下絕對值最大的k個value和對應的k個index，然後所有卡的value和index匯總到一起，再計算每個index的總和。通過這種方法，使通信量大大減少。

在此基礎上，團隊進一步提出了以下幾點優化：

1) 首次提出層級Topk壓縮通信算法。

還是以單機8卡的機型為例，首先第一步先用ReduceScatter在單機內通信，利用NVLink2的高速帶寬，使每塊卡快速得到單機內1/8梯度的總和。

第二步的流程與普通的Topk壓縮通信一樣（如上圖），不同的是每塊卡只對原梯度1/8大小的數據量做Topk壓縮，從而減少Topk壓縮的計算時間，而且每塊卡在AllGather時的節點數減少到1/8，因為AllGather的通信量與節點數成正比，所以8卡的總通信量也減少到普通Topk壓縮通信的1/8。而且單機的8張卡是並行跨機通信的，也提升了帶寬利用率.。

第三步再通過單機內的AllGather使每張卡得到整個梯度的通信結果。

2) Topk壓縮性能優化。

如何高效進行Topk壓縮也是提高整體通信效率至關重要的一環，尤其是在GPU上，若使用Tensorflow 自身的 Topk 操作，在8MB數據下選取最大的1%個value需要花費42ms，而圖片大小為96*96*3時的前、反向計算的總時間也才幾十ms，因此嚴重影響了整個梯度壓縮通信的性能。

為了解決Topk壓縮慢的問題，團隊調研並借鑑了多種壓縮算法，使用了性能較好的閾值二分搜索選擇算法[8]，並在平衡了效率和精度後做了一些改進，二分搜索過程中記錄最大的k'(k'

團隊也為算法實現了高性能的CUDA Kernel，首先為了最大化GPU利用率，合理配置Kernel使用的block和thread個數，並保證每個SM上擁有至少16個活動warp，以有效的隱藏訪存延遲。其次為了提高內存吞吐量，將多次傳輸合併為一次傳輸來減少CPU和GPU之間的內存傳輸；另外使用帶寬更高的共享內存，並通過動態分配的方式，提高共享內存的使用效率；保證block內的各個線程訪問連續對齊的內存塊，提高緩存命中率。

最後為了提高指令吞吐量，使用效率更高的CUDA接口，並對循環進行展開，減少條件分支。通過以上對CUDA Kernel的深度優化後，在8MB數據下選取前1% 的value只需要2ms，與Tensorflow本身的Topk實現相比，加速比為21。

3) 針對allgather通信進行優化。

當所有卡要通信的數據大小一致時，使用ncclAllGather來通信，不一致時使用ncclBroadCast來模擬allgatherv，這個優化也已經提交到Horovod社區。因為在壓縮通信這個場景下，每塊卡壓縮後的數據大小是一致的，所以使用ncclAllGather通信，比mpi的allgatherv更快。

ResNet-50 v1.5的參數量大概是25M個FP32，以壓縮率0.01為例，通信的value用FP16表示的話有0.5MB，而index是int32，所以有1MB。使用層級梯度壓縮通信時，跨機AllGather的節點數是16，所以總的跨機通信量大概是24MB。

4) LARS計算並行化

為保證大batch下收斂，團隊引入了LARS優化器，但LARS也引入了更多的計算量，每塊卡會串行地計算每一層梯度的本地學習率。團隊發現，由於該部分時間的消耗，計算時間無法與通信時間完全覆蓋，導致整體訓練時間的延長。由於LARS的計算是梯度通信之後，每張卡的針對模型計算出的本地學習率都是一致的，也就是每張卡都在做同樣的計算，且計算出來的本地學習率對每層梯度為單值。

除了LARS，LAMB也是這種梯度規約後的操作，屬於每塊做重複計算的操作。針對這類操作，根據操作的計算代價和通信代價進行評估，如果是計算代價比較大的，就會自動將其轉換成多卡並行計算再進行匯總，從而使總時間最少。

大batch收斂

1）大batch調參策略

訓練樣本上，團隊在性能和收斂性上做了權衡，使用了變解析度的多階訓練：訓練前期利用多個epoch的低解析度樣本快速收斂，並在訓練後期使用少量epoch的高解析度調高精度。

模型上，使用了業界常用的ResNet50 v1.5，同時在去除了bias及BN層參數正則化項的情況下，對BN層的beta做了一定的初始化調整，在較少迭代步數上取得了較好的收斂效果。

優化器上，選擇了業界普遍採用的MomentumSGD結合LARS層級自適應優化器，為了快速收斂，團隊對該優化器做了一些改動。對LARS優化器採用了隨迭代步變化的置信比約束策略：在Warm up階段，為了避免由於噪聲過大導致發散的問題，將lars優化器的置信比裁剪約束到較小的範圍，穩定的收斂；Warm up階段之後放鬆對lars優化器置信比的裁剪約束，加速收斂。

對於優化器中的動量值採用了全周期一次完整循環的方式，前期為了收斂的穩定性，動量值會隨著訓練的迭代不斷增大，防止訓練發散，後期為了訓練的模型適配圖片變大的解析度，動量值隨著訓練的迭代不斷減小，放大後期梯度對模型的貢獻，加速收斂。

由於考慮到訓練中使用的多解析度樣本，對學習率採用了分段線性變斜率策略，在warmup前後採用不同斜率的線性增長和線性減小，在不同解析度階段使用不同的學習率斜率，由AutoML進行端點斜率調節確定最終策略。

實際超參選擇中，在確定學習率策略等超參數方案後，團隊利用平臺大算力結合自研天風AutoML框架，以最後一個epoch的精度作為評估指標，對於各階段不同大小的樣本進行分段分塊有針對性的超參微調，並結合AutoML結果對超參方案進行迭代更新。

損失函數上，使用了label smoothing這一正則化手段對訓練標籤做了平滑處理，以防止過擬合幫助收斂。

2）梯度壓縮精度

在整個訓練過程當中，只在圖片大小為96*96*3時採用梯度壓縮通信，因為這種情況下通信的時間佔比最大，梯度壓縮通信對訓練速度的提升最好，而且對精度的影響也最小。

精度補償採用的是梯度本地累積的方法，也就是沒有被選中通信的index的值，會在本地累加，參與下一輪通信。

3）AutoML調參

針對業界開源框架研發 AutoML 算法難、集成自研訓練平臺工作量大和在大規模長任務訓練場景優化效果差的問題，研發了天風 (TianFeng)， 是聚焦於大規模訓練場景的、輕量的、易擴展的、易使用的 AutoML 框架，通過高度抽象 AutoML 的通用過程，來解決這些業界難題。

為了方便算法工程師做AutoML算法研發和新平臺接入，天風內部分為三層，算法層、邏輯層和接入層。

（1）算法層，AutoML算法都在這一層，通過抽象算法邏輯，方便算法工程師實現和接入AutoML算法。

（2）邏輯層，負責對接Client，算法邏輯下發、異常處理和持久化。

（3）接入層，負責訓練平臺的對接，當前可以很簡單的接入自研訓練平臺，結合AutoML能力和自研訓練平臺的調度和訓練能力。

在ImageNet訓練中，存在訓練規模大，超參數範圍廣，人工調參效率低且效果不好的問題。因而引入天風用於超參數空間的探索，將人力較大程度地從繁瑣的手工調參中解放出來。

依託智能鈦平臺的海量彈性算力，天風AutoML利用自研的並行貝葉斯優化算法進行超參數搜索，充分利用智能鈦平臺資源池中的空閒GPU算力資源並行進行多組彈性任務AutoML實驗，做到了高優先級任務佔用資源時實驗自動終止，資源空閒時自動加載實驗終止前的先驗數據恢復。

整體上AutoML搜索超參數為大batch size調參指明了方向，在同一模型下，AutoML僅需由工程師指定所需要搜索的超參數及其範圍，即可對超參數空間進行探索從而對調參思路進行快速驗證。

4.4 高性能機器學習平臺——智能鈦

智能鈦機器學習平臺是為 AI 工程師打造的一站式機器學習服務平臺，為用戶提供從數據預處理、模型構建、模型訓練、模型評估到模型服務的全流程開發支持。智能鈦機器學習平臺內置豐富的算法組件，支持多種算法框架，滿足多種AI應用場景的需求。自動化建模（AutoML）的支持與拖拽式任務流設計讓 AI 初學者也能輕鬆上手。

取得的成效

通過以上 Light大規模分布式多機多卡訓練框架及平臺等一系列完整的解決方案，ImageNet的訓練結果取得了新突破。同時，這些方案也廣泛應用在騰訊公有雲以及騰訊內部的業務，並形成了騰訊雲帆（深度學習框架和加速）Oteam、算力Oteam等協同組織，在更廣泛的範圍內協同共建。以ImageNet的訓練結果為例說明下。

單機優化

在單機的速度上，4種不同的圖片大小時都超過了友商。例如圖片大小為96*96*3時，優化前後的速度對比如下圖所示。

多機擴展

在25G的TCP網絡上，圖片大小為96*96*3時，由於計算時間相對較少，通信時間佔比大，所以多機擴展性是最差的。

如果使用Horovod的全局AllReduce，訓練速度只有885樣本/秒，而使用LightCC多機多卡通信庫的話，不使用壓縮通信時的訓練速度是2450樣本/秒（LARS優化前是2360樣本/秒），而使用壓縮通信的訓練速度是3100樣本/秒。而友商128卡的速度是2608樣本/秒，即不壓縮通信時由於帶寬限制導致略差於友商，而壓縮通信之後超過了友商。

在圖片大小為128*128*3時，由於計算時間變長，非壓縮通信和壓縮通信都只有最後一個梯度的通信沒有被計算隱藏，所以壓縮通信並沒有速度上的優勢，不過非壓縮通信的速度也已經超過了友商。

收斂

在收斂精度方面，通過手動設置超參與AutoML調參相結合，在28個epoch將top5精度訓練到93%。

2分31秒訓練ImageNet

在2分36秒內訓練 ImageNet 28個epoch，TOP5精度達到93%；使用壓縮通信後，在精度達到93%情況下，時間進一步減少到2分31秒，刷新了業界最好成績，創造了業界新記錄。若跨機網絡改為RoCE，則訓練時間再進一步減少到2分2秒。

構建穩定、易用、好用、高效的平臺和服務，將成為算法工程師的重要生產力工具，也會助力遊戲AI、計算機視覺AI，廣告推薦AI、翻譯AI、語音ASR AI等典型AI業務從一個成功走向另一個更大的成功。

本次破紀錄的ImageNet訓練，由騰訊機智團隊、騰訊雲智能鈦團隊、騰訊優圖實驗室、騰訊大數據團隊和香港浸會大學計算機科學系褚曉文教授團隊協同優化完成。

參考連結：

1.Highly Scalable Deep Learning Training System with Mixed-Precision: Training ImageNet in Four Minutes

https://arxiv.org/abs/1807.11205

2.ImageNet/ResNet-50 Training in 224 Seconds

https://nnabla.org/paper/imagenet_in_224sec.pdf

Krizhevsky, Alex, Ilya Sutskever, and Geoffrey E. Hinton. "Imagenet 3.classification with deep convolutional neural networks." Advances in neural information processing systems. 2012.

4.He, Kaiming, et al. "Deep residual learning for image recognition." Proceedings of the IEEE conference on computer vision and pattern recognition. 2016.

5.Vaswani, Ashish, et al. "Attention is all you need." Advances in neural information processing systems. 2017.

6.Devlin, Jacob, et al. "Bert: Pre-training of deep bidirectional transformers for language understanding." arXiv preprint arXiv:1810.04805 (2018).

7.Brown, Tom B., et al. "Language models are few-shot learners." arXiv preprint arXiv:2005.14165 (2020).

8.RedSync : Reducing Synchronization Traffic for Distributed Deep Learning

https://arxiv.org/abs/1808.04357

本項目論文《Efficient Distributed Deep Learning Training on Public Cloud Clusters: ImageNet/Resnet-50 Training in Two Minutes》，將首先在arxiv上提交供查閱。