之前一段時間接觸了大規模神經網路訓練，看了不少優秀的工作，在這裡當做筆記記下來。同時也希望可以拋磚引玉，和各位大佬交流一下這方面的現有工作以及未來的方向
（1）大規模訓練工作的幾種型別
大規模訓練和普通分散式訓練還是有區別的，主要體現在
大
這個字上面。一般來說會涉及到幾百個分散式節點同時工作，模型的引數量以及運算量往往很大（比如BERT，GPT3等等）
我認為在這個task下當前的工作主要歸結為以下三種
對通訊本身的最佳化
神經網路訓練通訊的最佳化
大規模下如何保持精度
其中1主要是通訊庫的最佳化，嚴格來說和神經網路本身並沒有關係，這裡面比較優秀的工作有經典的ring-base all-reduce（最先在百度的工作中被用於神經網路訓練baidu-research/baidu-allreduce），騰訊的分層通訊（
https：//
arxiv。org/abs/1807。1120
5
），以及sony的2D all-reduce（Massively Distributed SGD： ImageNet/ResNet-50 Training in a Flash）。
而第2部分的工作都針對於如何在神經網路這個訓練模式下做通訊最佳化。這方面的思路很廣，比如商湯提出的稀疏通訊（
https：//
arxiv。org/abs/1902。0685
5
），杜克大學提出的TernGrad （TernGrad： Ternary Gradients to Reduce Communication in Distributed Deep Learning）
第三部分和前兩個不同，主要關注點在於精度而非效能。在大規模訓練的情況下，一種常見的做法是做資料並行，即把batch size設的很大，那麼原來跑90個epoch需要迭代1000次的話，把batch size擴大10倍，就只需要迭代100次，即引數的更新次數減少了很多。如何在這種情況下收斂到小batch size也是一個棘手的問題。在這個領域比較好的工作有face book的線性倍增學習率（
https：//
arxiv。org/pdf/1706。0267
7。pdf
）以及伯克利尤洋的LAR演算法（
https：//
arxiv。org/pdf/1709。0501
1。pdf
）
打一個廣告：尤洋已經從伯克利畢業，目前在新加坡國立大學成立了實驗室。目前正在火熱招生，有興趣的同學歡迎投遞簡歷(郵箱：youy (at) comp (dot) nus (dot) edu (dot) sg)。如果害羞不好意思直接聯絡大牛，可以私聊我，我幫你轉發~
I：對通訊本身的最佳化
（懶得寫了，偷個懶）
我對這方面瞭解十分有限，推薦大家讀騰訊團隊寫的介紹（蘭瑞Frank：騰訊機智團隊分享——AllReduce演算法的前世今生）
II：神經網路的通訊最佳化
分散式神經網路訓練目前主要有兩種模式：資料並行和模型並行。
資料並行比較簡單，下面這張圖是經典的資料並行的同步訓練的場景：所有節點（即圖中的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。在這裡演示瞭如何用兩個節點去算連續的兩個矩陣乘法。
我們要做的操作是首先算出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。0696
5。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）
這個工作主要的貢獻在於提供了一套原語，允許最高層的開發者（寫python的人）透過簡單的方式指導程式碼生成（即編譯器）生成對應的模型並行的程式碼
