刷新世界紀錄：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'<k且最接近k)個 value 并確保在最終結果中，從而使結果的k個 value 最接近真正的 topk value。<="" font="">

　　團隊也為算法實現了高性能的 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

總結

以上是生活随笔為你收集整理的2分31秒，腾讯云创造128卡训练ImageNet新记录的全部內容，希望文章能夠幫你解決所遇到的問題。

如果覺得生活随笔網站內容還不錯，歡迎將生活随笔推薦給好友。