文 | 立交橋跳水冠軍

源 | 知乎

---

大規模神經網絡訓練一般會涉及到幾百個分布式節點同時工作，模型的參數量以及運算量往往很大，作者認為在這個task下當前的工作主要歸結為以下三種：對通信本身的優化，神經網絡訓練通信的優化，大規模下如何保持精度。

之前一段時間接觸了大規模神經網絡訓練，看了不少優秀的工作，在這里當做筆記記下來。同時也希望可以拋磚引玉，和各位大佬交流一下這方面的現有工作以及未來的方向（1）大規模訓練工作的幾種類型大規模訓練和普通分布式訓練還是有區別的，主要體現在大這個字上面。一般來說會涉及到幾百個分布式節點同時工作，模型的參數量以及運算量往往很大（比如BERT，GPT3等等）我認為在這個task下當前的工作主要歸結為以下三種：

對通信本身的優化

神經網絡訓練通信的優化

大規模下如何保持精度

其中1主要是通信庫的優化，嚴格來說和神經網絡本身并沒有關系，這里面比較優秀的工作有經典的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部分的工作都針對于如何在神經網絡這個訓練模式下做通信優化。這方面的思路很廣，比如商湯提出的稀疏通信：

https://arxiv.org/abs/1902.06855

杜克大學提出的TernGrad (TernGrad: Ternary Gradients to Reduce Communication in Distributed Deep Learning：

https://arxiv.org/abs/1705.07878

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

對通信本身的優化

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

神經網絡的通信優化

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

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

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

下面這張圖摘自英偉達的Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism:
https://arxiv.org/abs/1909.08053

在這里演示了如何用兩個節點去算連續的兩個矩陣乘法。

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

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

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

展開來講，數據并行和模型并行也可以細分。數據并行可以分為：

• 同步式數據并行

• 異步式數據并行

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

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

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

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

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

• 粗粒度并行

• 細粒度并行

它們的區別在于并行的層級：粗粒度每個節點會算不同的layer，而細粒度會將layer也做拆。

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

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

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

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

后臺回復關鍵詞【入群】

加入賣萌屋NLP/IR/Rec與求職討論群

后臺回復關鍵詞【頂會】

獲取ACL、CIKM等各大頂會論文集！

總結

以上是生活随笔為你收集整理的谈谈神经网络的大规模训练优化的全部內容，希望文章能夠幫你解決所遇到的問題。

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