綜述

因為我個人最近在從事可能是AI領域對性能挑戰最大的方向，自動駕駛領域，所以對整個深度學習訓練的優化尤為關注，最近一直在學習相關內容，謹以此篇文章做一個總結。

我一直很看好深度學習訓練優化這個方向,因為從大的環境上來看，似乎大模型會成為未來的一個趨勢,目前以Google、OpenAI、阿里等廠商為代表的一系列頭部的AI研究機構，已經把模型尺寸做到十萬億參數級別，明年應該可以達到百萬億參數，總體呈指數倍增長態勢。而核心計算芯片的性能，卻無法保證指數倍增長。

(上圖因為制作時間較久,具體數字可能存在偏差)

所以，未來在人工智能領域所面對的一個主要矛盾就是指數倍增長的算力需求與線性增長的硬件計算能力的矛盾。而深度學習訓練優化技術，是解決這樣矛盾的一個關鍵。深度學習訓練優化技術，又分為主觀和客觀性兩個方面，有模型開發者的主動選擇，用什么樣的并行化模式，怎樣處理通信算子，也有來自客觀的優化，比如框架的編譯優化。本文將嘗試對常見的優化技術做一個概述。

具體結構如下：

分布式策略優化

前面介紹了在大模型時代，算力往往會成為短板，那么如何解決算力問題，比較好的方案就是并行化計算，或者說是分布式計算。分布式計算有許多種策略組成，常見的分布式策略有數據并行、模型并行、流水并行、算子并行。

1、數據并行[1]

數據并行是比較常見的并行化計算邏輯，是將訓練數據做打散和分片，然后每張GPU卡取一部分數據進行計算，每張GPU卡保存著全部的模型信息，通過All Reduce的方式進行梯度的更新，這里涉及到很多通信優化的方案，我們下一部分通信優化模塊會介紹。數據并行經常跟模型并行、流水并行等模式一起使用。

開源的一些框架，比如Uber開源的Horovod對數據并行及其它并行方式提供了很好的支持。

2、模型并行（算子拆分）[2]

當一些模型的部分模塊計算量明顯高于其它模塊的時候，會使用到算子并行的策略。舉個例子，比如Resnet模型做圖像分類，當分類器的類別極大，比如要做幾十萬類別的區分，那么全連接層會變得特別的大。如下圖所示。

分類數為10 W時，ResNet50部分的模型權重大小89.6 MB，全連接層FC（Fully Connected Layer）部分的模型權重大小為781.6 MB，是前者的8.7倍。當采用數據并行進行分布式訓練時，在后向階段（Back Propagation過程），FC部分計算得到梯度后，立刻通過AllReduce進行梯度同步，同時ResNet50部分的梯度會繼續進行后向計算。然而由于FC部分同步的梯度過大，當ResNet50部分的梯度計算完成時，FC部分的梯度通常還在通信過程中。因此，FC部分的梯度同步無法良好地與ResNet50的后向計算Overlap，導致整輪迭代過程中的通信占比十分高，性能低下。

針對ResNet50的大規模分類任務，將模型分為兩個Stage。將Stage 0的ResNet50部分通過數據并行策略復制N份至不同的卡中，進行數據并行。將Stage 1的FC和Softmax部分通過算子拆分策略分片至不同卡中，進行模型并行。假設共六張卡（GPU0、GPU1、GPU2、GPU3、GPU4及GPU5），將GPU分為兩組，分別進行數據并行和算子拆分，如下圖所示。

3、流水并行[3]

上面介紹了模型并行，在執行模型并行的過程中極容易出現GPU等待的問題。比如一部分模型還在做梯度更新，而另一部分的模型這一輪已經更新好了。在Bert類模型經常容易出現這樣的狀況。

比如把Bert模型進行拆分，不同GPU處理模型的不同模塊：

模型并行計算過程中會出現閑時等待問題，如下圖所示。

可以看到同一個時間段只有一張卡在做前向或者是后向的訓練。流水并行要解決的問題是，當一張卡訓練完之后馬上通知下一張卡進行訓練，讓整個計算過程像流水線一樣連貫，流水并行在大模型場景會大大提升計算效率，減少GPU的等待時間。

無論是模型并行還是流水并行，這些并行模式通常會混合使用，也就是說混合并行計算。業內有許多開源的框架可以實現上述能力，比如微軟開源的DeepSpeed可以提供模型并行和流水并行的能力。DeepSpeed項目地：https://github.com/microsoft/DeepSpeed

通信優化

通信優化有很多方面，比如通信系統優化、通信模式優化和通信內容優化。接下來逐一介紹下這些通信優化領域的基礎概念。

1、通信系統優化[5]

通信系統可以分為兩個大的類別，一種是Allreduce結構，目前更常用的是Ring-Allreduce。另一種是PS，也就是Parameter Server參數服務器結構。

1.1、PS框架[6]

首先介紹下PS架構，在這種架構下，通常將計算節點分為PS節點和Worker節點兩個類型。其中PS節點主要用來存放參數，Worker節點負責計算梯度，Worker計算好的梯度數據傳輸給PS節點，由PS節點聚合后再反饋給各個Worker進行下一步訓練。

但是隨著模型逐漸增大，導致每次PS節點和Worker節點需要傳輸的數據越來越多，PS的通信結構使得PS節點的帶寬成為通信瓶頸，并不一定適合于超大規模的分布式深度學習訓練，于是便有了下文提到的，Ring Allreduce通信結構。

1.2 、Ring-Allreduce框架[7]

關于Ring-Allreduce，我之前寫過相關文章，這里直接引用，Ring-Allreduce主要解決兩方面問題。

問題一，每一輪的訓練迭代都需要所有卡都將數據同步完做一次Reduce才算結束。如果卡數比較少的情況下，其實影響不大，但是如果并行的卡很多的時候，就涉及到計算快的卡需要去等待計算慢的卡的情況，造成計算資源的浪費。

問題二，每次迭代所有的計算GPU卡多需要針對全部的模型參數跟Reduce卡進行通信，如果參數的數據量大的時候，那么這種通信開銷也是非常龐大，而且這種開銷會隨著卡數的增加而線性增長。

為了解決這樣的問題，就引入了一種通信算法Ring Allreduce，通過將GPU卡的通信模式拼接成一個環形，從而減少隨著卡數增加而帶來的資源消耗，如下圖所示：

將GPU卡以環形通信之后，每張卡都有一個左手卡和右手卡，那么具體的模型參數是如何傳遞的呢，可以看下圖：

因為每張卡上面的網絡結構是固定的，所以里面的參數結構相同。每次通信的過程中，只將參數send到右手邊的卡，然后從左手邊的卡receive數據。經過不斷地迭代，就會實現整個參數的同步，也就是reduce。形成以下這張圖的樣式：

通過Ring Allreduce的方式，基本上可以實現當GPU并行卡數的增加，實現計算性能的線性增長。

通過上圖不難發現，Ring-Allreduce通信框架在數據并行方面有著是否優異的效果。

2、通信模式

除了通信系統架構，通信模式也對性能有較大影響。通信模式通常分為異步通信和同步通信兩種。在PS架構中通常會涉及到通信模式的問題。同步通信是比較低效的一種方案，指的是利用PS節點做參數的存儲和Reduce，每次PS會等所有的worker都計算完梯度，才做一次梯度的平均，然后分發給各個worker進行更新。這種模式很容易造成堵塞，當PS和worker的通信阻斷了，整個訓練就無法進行下去，當然如果增加PS節點來緩解阻斷問題，又會增加整個系統的負責度。

異步更新會比同步更新機制更完善，指的是PS節點不必每次都等所有worker計算完再做梯度更新，當PS節點收到worker的梯度更新請求后，立刻進行梯度的下發。當然，異步更新也會引出新的問題需要優化，就是每一輪更新的梯度可能是幾輪以前計算出來的，另外如果參數讀取未加鎖，有可能worker會從ps讀取到剛更新一般的參數。

3、通信內容優化[8][9]

以上介紹的都是通信的架構和模式，那么通信管道的內容也是有優化空間的。比如稠密通信可以做信道合并，稀疏通信可以數據存儲和格式轉換的優化。

通常做深度學習訓練的時候，使用的數據都是TFrecord類型，protobuf格式的文件，或者是hdf5、pickle，這種文件都是行存，每次讀取數據都需要獲取全量字段。但是很多場景，可能每次讀取只需要一個字段的數據，這樣列存的數據IO模式就比較高效，自動駕駛就是這樣的場景。舉個例子區分行存和列存，

以下是行存數據結構：

下圖是列存的數據結構：

列存允許用戶在訓練過程中加載列子集，從而節約通信信道的帶寬。列存通常會使用Parquet格式，uber開源了一個工作叫做Petastorm ，可以在TensorFlow、Pytorch等框架支持Parquet數據格式。

地址：https://github.com/uber/petastorm

顯存優化

首先分析下GPU中占用顯存的元素有哪些，神經網絡占據顯存的部分主要是模型自身的參數以及模型的輸出[10]。從模型角度分析，只有有參數的層才會占用顯存，比如卷積、全連接、Embedding，那些沒有參數的層其實不占用顯存，比如激活、池化、dropout。從模型輸出角度分析，基本上batch-size越大，顯存占用越大。總結下，降低顯存的占用的方案是降低batch-size，或者減少全連接層這種參數較多的神經網絡結構。

另外在系統工程層面也有顯存優化的方案，可以分為分時復用和合并共享兩種模式。

1、分時復用[11]

從工程層面，如果能讓GPU的顯存一直100%的被占用，其實是最大限度地利用了GPU的顯存資源。分時復用的原理是將GPU的顯存實現某種層面的資源隔離，利用Schedule的能力，將任務合理的分配到不同時間段內。

2、合并共享[11]

合并共享是指多個任務合并成一個上下文，因此可以共享GPU資源，同時發送kernel到GPU上，也共同使用顯存。NV提供了這樣的一套能力，叫做MPS（CUDA_Multi_Process_Service），感興趣通信可以看下文章的引用信息和NV的官方文檔：Multi-Process Service

整體上MPS的內存合并共享機制架構圖如下：

編譯優化

框架的編譯優化是比較底層的一種優化手段，傳統的編譯器優化依賴于工程師手動的去優化算子，那么計算框架的編譯優化希望通過一種偏自動化的方式去實現算子在各種硬件設備層面的編譯優化。為什么要做框架的編譯優化[12]，核心問題是人的精力是有限的，不可能基于無數可被優化的算子做人肉的優化。

編譯器將深度學習框架描述的模型在各種硬件平臺上生成有效的代碼實現，其完成的模型定義到特定代碼實現的轉換將針對模型規范和硬件體系結構高度優化。目前主流的編譯器對于環境的支持情況如下：

在論文”The Deep Learning Compiler: A Comprehensive Survey“中，還給出了深度學習編譯器的整體架構，如下圖所示：

在編譯器中可以分為三個主要部分[14]，前端計算圖和后端以及IR（Intermediate Representation）。通常IR是程序的抽象，用于程序優化。具體地說，DL模型在DL編譯器中被轉換成多級IR，其中高層IR駐留在前端，而底層IR駐留在后端。編譯器前端基于高層IR，負責與硬件無關的轉換和優化。基于底層IR，編譯器后端負責特定于硬件的優化、代碼生成和編譯。

整個優化分為前端圖級別的優化以及后端優化，接下來分別介紹下這兩個優化方向。

1、前端圖優化

前端圖優化總體上是通過輸入的神經網絡圖結構，找到其中可以優化的模塊并重新構造圖。圖優化可以分為節點優化、塊優化、數據優化。

其中節點優化是消除不必要的節點，比如A是零元張量，B是常數張量，A+B的計算就可以替換成B，這樣就減少了圖的復雜度。塊優化可以實現一些運算符層面的簡化，比如A*B+A*C是三次計算，(B+C)*A就變成了兩次計算，這樣就通過一些算子的融合機制實現了計算的精簡。數據全局優化，是通過全圖鏈路精簡計算復雜度，比如一個數值A已經計算過，那么接下來就不用重復計算，直接使用。

2、后端優化

后端優化包含指定硬件優化、自動調優技術以及內核庫優化。指定硬件優化指的是將底層IR轉化為LLVM IR，利用LLVM基礎結構生存適配CPU或者GPU的代碼，提高邏輯執行效率，這個過程叫做Codegen。

自動調優技術指的是合理設定每個計算單元在硬件層面的參數，因為每個計算都需要設定共享內存、寄存器大小等參數。利用自動調優技術，建立Cost model，可以自動判定參數，提升計算效率。

優化內核庫技術指的是把數據自動轉換成適配英特爾的DNNL（以前是MKL-DNN）、NVIDIA的cuDNN和AMD的MLOpen的數據結構和格式，使得代碼可以直接使用到這些底層的優化內核庫。

3、MLIR[15]

最后介紹下什么是MLIR，前面說了IR是計算邏輯在編譯層面的一種抽象表示。深度學習在編譯過程中，需要將high-level的IR轉換成low-level的IR去適配不同硬件環境。如下圖：

這種做法就造成了系統的復雜性，MLIR希望提供一種框架標準，使得不同IR可以模塊化，并且統一標準去適配各種硬件條件，從而減少系統復雜性。

總結

本文是我個人在學習深度學習優化技術過程中的一個總結，覆蓋了分布式策略、顯存、通信、編譯各個維度的優化基礎概念，這些概念的學習來自于下方引用部分的小伙伴的分享，感謝大家。有不準確的地方望大家指正，我會修改。整個深度學習優化內容負責，種類繁多，后續需要不斷學習和進步。

引用:

[1]分布式計算策略：PAI視頻內容

[2]算子拆分：大規模分類的分布式訓練（算子拆分） - 機器學習PAI - 阿里云

[3]流水并行：BertLarge分布式訓練（流水并行） - 機器學習PAI - 阿里云

[4]DeepSpeed：https://github.com/microsoft/DeepSpeed

[5]PS&RingAllReduce[深度學習] 分布式模式介紹（一）_小墨魚的專欄-CSDN博客

[6]PS:https://zhuanlan.zhihu.com/p/50116885

[7]Ring-Allreduce:ring allreduce和tree allreduce的具體區別是什么？ - 知乎

[8]通信內容：分布式深度學習訓練中的通信優化有哪些主流的研究方向？ - 知乎

[9]Parquet：https://zhuanlan.zhihu.com/p/45364584

[10]顯存優化：科普帖：深度學習中GPU和顯存分析 - 知乎

[11]分時復用、共享復用：It - 收藏夾 - 知乎

[12]編譯器的必要性：深度學習編譯技術的現狀和未來 - 知乎

[13]編譯器概述：https://zhuanlan.zhihu.com/p/139552817

[14]The Deep Learning Compiler: A Comprehensive Survey翻譯：翻譯《The Deep Learning Compiler: A Comprehensive Survey》綜述翻譯 - 簡書

[15]MLIR說明：https://zhuanlan.zhihu.com/p/101879367

總結

以上是生活随笔為你收集整理的深度学习分布式策略优化、显存优化、通信优化、编译优化综述的全部內容，希望文章能夠幫你解決所遇到的問題。

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