大規(guī)模神經網絡的訓練優(yōu)化入門

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

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

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

我認為在這個task下當前的工作主要歸結為以下三種

1. 對通信本身的優(yōu)化
2. 神經網絡訓練通信的優(yōu)化
3. 大規(guī)模下如何保持精度

其中1主要是通信庫的優(yōu)化，嚴格來說和神經網絡本身并沒有關系，這里面比較優(yōu)秀的工作有經典的ring-base all-reduce（最先在百度的工作中被用于神經網絡訓練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部分的工作都針對于如何在神經網絡這個訓練模式下做通信優(yōu)化。這方面的思路很廣，比如商湯提出的稀疏通信(https://arxiv.org/abs/1902.06855)，杜克大學提出的TernGrad (TernGrad: Ternary Gradients to Reduce Communication in Distributed Deep Learning（https://arxiv.org/abs/1705.07878)

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

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

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

（懶得寫了，偷個懶）

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

II：神經網絡的通信優(yōu)化

分布式神經網絡訓練目前主要有兩種模式：數(shù)據并行和模型并行。

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

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

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

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

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

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

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

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

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

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

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

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

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

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

模型并行可以分為

• 粗粒度并行
• 細粒度并行

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

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

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

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

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