大規(guī)模神經(jīng)網(wǎng)絡(luò)的訓(xùn)練優(yōu)化入門

之前一段時間接觸了大規(guī)模神經(jīng)網(wǎng)絡(luò)訓(xùn)練，看了不少優(yōu)秀的工作，在這里當(dāng)做筆記記下來。同時也希望可以拋磚引玉，和各位大佬交流一下這方面的現(xiàn)有工作以及未來的方向

（1）大規(guī)模訓(xùn)練工作的幾種類型

大規(guī)模訓(xùn)練和普通分布式訓(xùn)練還是有區(qū)別的，主要體現(xiàn)在大這個字上面。一般來說會涉及到幾百個分布式節(jié)點(diǎn)同時工作，模型的參數(shù)量以及運(yùn)算量往往很大（比如BERT，GPT3等等）

我認(rèn)為在這個task下當(dāng)前的工作主要?dú)w結(jié)為以下三種

1. 對通信本身的優(yōu)化
2. 神經(jīng)網(wǎng)絡(luò)訓(xùn)練通信的優(yōu)化
3. 大規(guī)模下如何保持精度

其中1主要是通信庫的優(yōu)化，嚴(yán)格來說和神經(jīng)網(wǎng)絡(luò)本身并沒有關(guān)系，這里面比較優(yōu)秀的工作有經(jīng)典的ring-base all-reduce（最先在百度的工作中被用于神經(jīng)網(wǎng)絡(luò)訓(xùn)練baidu-research/baidu-allreduce（https://github.com/baidu-research/baidu-allreduce）,騰訊的分層通信（https://arxiv.org/abs/1807.11205），以及sony的2D all-reduce（Massively Distributed SGD: ImageNet/ResNet-50 Training in a Flash（https://arxiv.org/abs/1811.05233）。

而第2部分的工作都針對于如何在神經(jīng)網(wǎng)絡(luò)這個訓(xùn)練模式下做通信優(yōu)化。這方面的思路很廣，比如商湯提出的稀疏通信(https://arxiv.org/abs/1902.06855)，杜克大學(xué)提出的TernGrad (TernGrad: Ternary Gradients to Reduce Communication in Distributed Deep Learning（https://arxiv.org/abs/1705.07878)

第三部分和前兩個不同，主要關(guān)注點(diǎn)在于精度而非性能。在大規(guī)模訓(xùn)練的情況下，一種常見的做法是做數(shù)據(jù)并行，即把batch size設(shè)的很大，那么原來跑90個epoch需要迭代1000次的話，把batch size擴(kuò)大10倍，就只需要迭代100次，即參數(shù)的更新次數(shù)減少了很多。如何在這種情況下收斂到小batch size也是一個棘手的問題。在這個領(lǐng)域比較好的工作有face book的線性倍增學(xué)習(xí)率（https://arxiv.org/pdf/1706.02677.pdf）以及伯克利尤洋的LAR算法（https://arxiv.org/pdf/1709.05011.pdf）

打一個廣告：尤洋已經(jīng)從伯克利畢業(yè)，目前在新加坡國立大學(xué)成立了實(shí)驗(yàn)室。目前正在火熱招生，有興趣的同學(xué)歡迎投遞簡歷(郵箱：youy (at) comp (dot) nus (dot) edu (dot) sg)。如果害羞不好意思直接聯(lián)系大牛，可以私聊我，我?guī)湍戕D(zhuǎn)發(fā)~

I：對通信本身的優(yōu)化

（懶得寫了，偷個懶）

我對這方面了解十分有限，推薦大家讀騰訊團(tuán)隊(duì)寫的介紹（蘭瑞Frank：騰訊機(jī)智團(tuán)隊(duì)分享--AllReduce算法的前世今生（https://zhuanlan.zhihu.com/p/79030485）

II：神經(jīng)網(wǎng)絡(luò)的通信優(yōu)化

分布式神經(jīng)網(wǎng)絡(luò)訓(xùn)練目前主要有兩種模式：數(shù)據(jù)并行和模型并行。

數(shù)據(jù)并行比較簡單，下面這張圖是經(jīng)典的數(shù)據(jù)并行的同步訓(xùn)練的場景：所有節(jié)點(diǎn)（即圖中的GPU0-GPU3）都保存整個模型(粉色的Params)，每次迭代，不同的節(jié)點(diǎn)會得到不同的數(shù)據(jù)，每個節(jié)點(diǎn)用得到的數(shù)據(jù)做正向和反向計(jì)算，得到每個參數(shù)的梯度。之后整個分布式系統(tǒng)會同步所有節(jié)點(diǎn)的梯度，即每個節(jié)點(diǎn)的local gradient做一次all reduce操作，得到全局的global gradient（最下面藍(lán)色的Gradients）。每個節(jié)點(diǎn)用這個global gradient更新參數(shù)。

顯而易見，數(shù)據(jù)并行基于一個假設(shè)：每個節(jié)點(diǎn)都可以放下整個模型。這個假設(shè)在如今某些模型上（說的就是你，GPT3?。。。┦遣缓侠淼?，因此我們還需要模型并行，即不同節(jié)點(diǎn)負(fù)責(zé)計(jì)算神經(jīng)網(wǎng)絡(luò)模型的不同部分（比如有一個100層的網(wǎng)絡(luò)，那么我們可以讓第一個節(jié)點(diǎn)存儲前50層的參數(shù)，并負(fù)責(zé)計(jì)算前50層，另一個網(wǎng)絡(luò)則負(fù)責(zé)后面50層）

下面這張圖摘自英偉達(dá)的Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism（https://arxiv.org/abs/1909.08053）。在這里演示了如何用兩個節(jié)點(diǎn)去算連續(xù)的兩個矩陣乘法。

我們要做的操作是首先算出Y=GeLU(XA)，再算Z=Dropoug(YB)。其中，X，A，B都是矩陣，而且矩陣規(guī)模都很大

假設(shè)我們希望用兩個分布式節(jié)點(diǎn)完成這個計(jì)算，那么我們可以把矩陣A按colum切成A1，A2兩份，分別存到節(jié)點(diǎn)0和節(jié)點(diǎn)1中。同時我們也把矩陣B按行切成B1，B2兩份，分別存到節(jié)點(diǎn)0和節(jié)點(diǎn)1中。然后我們將X做一個broadcast（圖中f部分），分別發(fā)送到兩個節(jié)點(diǎn)上，算得Z1和Z2，在做一次all reduce（圖中g(shù)部）將Z1和Z2相加，得到最終的Z。

這里面有一個很巧（也很繞）的地方，那就是為什么A要按列切，B要按行切？我們可不可以把它們反過來？答案是：最好不要，因?yàn)槿绻催^來，的確計(jì)算上可行，但是我們就會增加一次通信（即算Y=XA的時候我們就要做一次通信），這樣顯然速度會變慢。

展開來講，數(shù)據(jù)并行和模型并行也可以細(xì)分

數(shù)據(jù)并行可以分為

• 同步式數(shù)據(jù)并行
• 異步式數(shù)據(jù)并行

同步式比較簡單，就是我最上面那張圖演示的。

異步式復(fù)雜一些：我們很容易發(fā)現(xiàn)，最后全局all reduce gradient的時候會耗時比較多，分布式系統(tǒng)越大，消耗越大，而且這樣做還有一個隱藏的假設(shè)：分布式系統(tǒng)是homogeneous的，即每個分布式節(jié)點(diǎn)不會差的很多。舉個例子，如果每個節(jié)點(diǎn)實(shí)力相當(dāng)，那么都會算10s就可以結(jié)束一個iteration，那么我們10s之后就可以開始一次通信。然而如果有一個節(jié)點(diǎn)（害群之馬）需要算100s，那么其他節(jié)點(diǎn)算完之后就得干等它90s才能做通信，那么是對資源的極大浪費(fèi)

想想看，你的老板絕對不允許你（打工人）干坐著什么事都不干，只因?yàn)槟愕倪M(jìn)度被別的同事block了。研究員也是如此，于是為了解決上面的問題，引入了異步式通信。簡單來說就是如果遭遇了上面的情況，快的節(jié)點(diǎn)等一會兒就不等了，他們之間做一次通信然后接著算下一輪。這個節(jié)點(diǎn)什么時候算好什么時候再和其他人一起all reduce梯度。

這樣做快是快了，但引入了另一個問題，那就是每個人的參數(shù)都不一樣了，那么他們根據(jù)不同的參數(shù)算得的梯度再去做all reduce就有一些不合理，就會導(dǎo)致神經(jīng)網(wǎng)絡(luò)精度受損。

有很多工作嘗試解決異步并行帶來的精度損失，不過據(jù)我所知并沒有特別general的方法，因此異步并行如今也很少被使用了

模型并行可以分為

• 粗粒度并行
• 細(xì)粒度并行

它們的區(qū)別在于并行的層級：粗粒度每個節(jié)點(diǎn)會算不同的layer，而細(xì)粒度會將layer也做拆分

粗粒度并行比較優(yōu)秀的工作有g(shù)oogle的GPipe(https://arxiv.org/pdf/1811.06965.pdf)

在粗粒度并行中，每個節(jié)點(diǎn)負(fù)責(zé)不同的layer，但是layer之間是存在數(shù)據(jù)依賴的，這就導(dǎo)致在之前的節(jié)點(diǎn)算的時候，后面的節(jié)點(diǎn)干等著。GPipe提出把數(shù)據(jù)按照batch緯度做切分得到多個micro batch，這樣第一個節(jié)點(diǎn)先算第一個micro batch（圖中F[0,0]），把算到的結(jié)果發(fā)給第二個節(jié)點(diǎn)去算，于是下一個時刻第二個節(jié)點(diǎn)在算第一個micro batch(F[1,0])，而第一個節(jié)點(diǎn)開始算第二個micro batch(F[0,1])

細(xì)粒度并行比較好的工作除了我之前介紹的Megatron之外，還有GShard(GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding（https://arxiv.org/abs/2006.16668)

這個工作主要的貢獻(xiàn)在于提供了一套原語，允許最高層的開發(fā)者（寫python的人）通過簡單的方式指導(dǎo)代碼生成（即編譯器）生成對應(yīng)的模型并行的代碼