綜述

因為我個人最近在從事可能是AI領(lǐng)域?qū)π阅芴魬?zhàn)最大的方向，自動駕駛領(lǐng)域，所以對整個深度學(xué)習(xí)訓(xùn)練的優(yōu)化尤為關(guān)注，最近一直在學(xué)習(xí)相關(guān)內(nèi)容，謹以此篇文章做一個總結(jié)。

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

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

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

具體結(jié)構(gòu)如下：

分布式策略優(yōu)化

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

1、數(shù)據(jù)并行[1]

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

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

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

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

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

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

3、流水并行[3]

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

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

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

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

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

通信優(yōu)化

通信優(yōu)化有很多方面，比如通信系統(tǒng)優(yōu)化、通信模式優(yōu)化和通信內(nèi)容優(yōu)化。接下來逐一介紹下這些通信優(yōu)化領(lǐng)域的基礎(chǔ)概念。

1、通信系統(tǒng)優(yōu)化[5]

通信系統(tǒng)可以分為兩個大的類別，一種是Allreduce結(jié)構(gòu)，目前更常用的是Ring-Allreduce。另一種是PS，也就是Parameter Server參數(shù)服務(wù)器結(jié)構(gòu)。

1.1、PS框架[6]

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

但是隨著模型逐漸增大，導(dǎo)致每次PS節(jié)點和Worker節(jié)點需要傳輸?shù)臄?shù)據(jù)越來越多，PS的通信結(jié)構(gòu)使得PS節(jié)點的帶寬成為通信瓶頸，并不一定適合于超大規(guī)模的分布式深度學(xué)習(xí)訓(xùn)練，于是便有了下文提到的，Ring Allreduce通信結(jié)構(gòu)。

1.2 、Ring-Allreduce框架[7]

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

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

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

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

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

因為每張卡上面的網(wǎng)絡(luò)結(jié)構(gòu)是固定的，所以里面的參數(shù)結(jié)構(gòu)相同。每次通信的過程中，只將參數(shù)send到右手邊的卡，然后從左手邊的卡receive數(shù)據(jù)。經(jīng)過不斷地迭代，就會實現(xiàn)整個參數(shù)的同步，也就是reduce。形成以下這張圖的樣式：

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

通過上圖不難發(fā)現(xiàn)，Ring-Allreduce通信框架在數(shù)據(jù)并行方面有著是否優(yōu)異的效果。

2、通信模式

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

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

3、通信內(nèi)容優(yōu)化[8][9]

以上介紹的都是通信的架構(gòu)和模式，那么通信管道的內(nèi)容也是有優(yōu)化空間的。比如稠密通信可以做信道合并，稀疏通信可以數(shù)據(jù)存儲和格式轉(zhuǎn)換的優(yōu)化。

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

以下是行存數(shù)據(jù)結(jié)構(gòu)：

下圖是列存的數(shù)據(jù)結(jié)構(gòu)：

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

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

顯存優(yōu)化

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

另外在系統(tǒng)工程層面也有顯存優(yōu)化的方案，可以分為分時復(fù)用和合并共享兩種模式。

1、分時復(fù)用[11]

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

2、合并共享[11]

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

整體上MPS的內(nèi)存合并共享機制架構(gòu)圖如下：

編譯優(yōu)化

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

編譯器將深度學(xué)習(xí)框架描述的模型在各種硬件平臺上生成有效的代碼實現(xiàn)，其完成的模型定義到特定代碼實現(xiàn)的轉(zhuǎn)換將針對模型規(guī)范和硬件體系結(jié)構(gòu)高度優(yōu)化。目前主流的編譯器對于環(huán)境的支持情況如下：

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

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

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

1、前端圖優(yōu)化

前端圖優(yōu)化總體上是通過輸入的神經(jīng)網(wǎng)絡(luò)圖結(jié)構(gòu)，找到其中可以優(yōu)化的模塊并重新構(gòu)造圖。圖優(yōu)化可以分為節(jié)點優(yōu)化、塊優(yōu)化、數(shù)據(jù)優(yōu)化。

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

2、后端優(yōu)化

后端優(yōu)化包含指定硬件優(yōu)化、自動調(diào)優(yōu)技術(shù)以及內(nèi)核庫優(yōu)化。指定硬件優(yōu)化指的是將底層IR轉(zhuǎn)化為LLVM IR，利用LLVM基礎(chǔ)結(jié)構(gòu)生存適配CPU或者GPU的代碼，提高邏輯執(zhí)行效率，這個過程叫做Codegen。

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

優(yōu)化內(nèi)核庫技術(shù)指的是把數(shù)據(jù)自動轉(zhuǎn)換成適配英特爾的DNNL（以前是MKL-DNN）、NVIDIA的cuDNN和AMD的MLOpen的數(shù)據(jù)結(jié)構(gòu)和格式，使得代碼可以直接使用到這些底層的優(yōu)化內(nèi)核庫。

3、MLIR[15]

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

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

總結(jié)

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

引用:

[1]分布式計算策略：PAI視頻內(nèi)容

[2]算子拆分：大規(guī)模分類的分布式訓(xùn)練（算子拆分） - 機器學(xué)習(xí)PAI - 阿里云

[3]流水并行：BertLarge分布式訓(xùn)練（流水并行） - 機器學(xué)習(xí)PAI - 阿里云

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

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

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

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

[8]通信內(nèi)容：分布式深度學(xué)習(xí)訓(xùn)練中的通信優(yōu)化有哪些主流的研究方向？ - 知乎

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

[10]顯存優(yōu)化：科普帖：深度學(xué)習(xí)中GPU和顯存分析 - 知乎

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

[12]編譯器的必要性：深度學(xué)習(xí)編譯技術(shù)的現(xiàn)狀和未來 - 知乎

[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

總結(jié)

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

如果覺得生活随笔網(wǎng)站內(nèi)容還不錯，歡迎將生活随笔推薦給好友。