MLNLP社群是國內外知名的機器學習與自然語言處理社群，受眾覆蓋國內外NLP碩博生、高校老師以及企業研究人員。

社群的願景是促進國內外自然語言處理，機器學習學術界、產業界和廣大愛好者之間的交流和進步，特別是初學者同學們的進步。

來源 | 知乎

作者 | 哈密瓜

一、前言

單個GPU卡記憶體有限，從而引發的各種問題，可分為兩個維度來看：

• 資料維度：對於小模型而言，雖然單卡能放下一個模型，但沒法放下一個batch所要的全部資料，所以資料得切分到不同卡上計算最終彙總，從而引入資料並行DP
• 模型維度：對於大模型而言，單卡無法放下一個模型，必須將模型切分到多張卡上，由此引入了除DP之外的各種並行（TP/PP/EP/SP），可統稱為模型並行

換句話而言，如果模型又小，單個batch資料不多，一張卡跑就夠了，也不需要各種並行。

如果模型一張卡能放下，但單個batch的資料放不下，那隻需要DP，這樣計算效率最高。在這種情況下再加入其他並行策略反而拖累訓練效率。

只有在模型單卡也放不下的時候，這時才必須考慮DP之外的並行。不同的並行策略根據模型結構的不同，採取的策略也不一樣。

DP/TP/PP/EP/SP都可以同時開啟（EP依賴於DP，只有DP開啟的情況下，才可能開啟EP；SP依賴於TP，只有在TP開啟的情況下才能用SP）。

二、各種並行方案

2.1 資料並行（Data Parallelism）

資料並行最簡單，每張卡複製相同的模型結構，僅對資料做切分。每張卡計算完的梯度也是針對各自資料的，需要做一次allreduce，然後使用最佳化器更新模型，進入下一次迭代。其數學基礎如下：

資料並行的優點是在所有並行中效率最高，除了一次allreduce之外，沒有任何額外通訊開銷，其通訊量為 2Mdtype ，M是模型大小，dtype是資料型別，2是因為allreduce的特性（allreuce=reducescatter+allgather，他們各自通訊量引數都是M）。

資料並行的缺點是：只能放微小模型，以llama-7B的bf16混合訓練為例，視訊記憶體佔用至少是7B*(2+2+(4+4+4))=224G，其中模型和梯度各佔2byte，adam最佳化器佔12byte。目前主流的H100視訊記憶體為80G，顯然一張卡放不下，因此需要引入TP/PP等模型並行，將一個模型分攤在多張卡上。

由於DP效率最高，DP為所有並行中首選且必帶的。同時只要一張卡放不下一個batch中的資料（以及其產生的啟用值），DP也是必選的。

2.2 張量並行（Tensor Parallelism）

張量並行是指對模型內部的引數矩陣切分，然後利用分塊矩陣乘進行計算得到正確結果，由於引數矩陣是以tensor表示，故叫張量並行

張量並行的優點是能分攤模型到多張卡上，缺點是帶來了不小的通訊開銷，影響訓練效率

以transfomer為例，闡述如何做TP，並分析其通訊開銷（可參考李沐老師的影片^[1]）

embeding部分：假設embeding的大小為 vocab*hidden ，GPU卡為N，每張卡分攤 vocab/N*hidden ；由於embeding的計算是查詢，每張卡拿到所有的input_ids並查詢input_id是否在embeding中，沒有設定為0。查詢完畢後做一次allreduce，通訊開銷為 2*b*s*h ，b是batch，s是seqlen，h是hidden_size

attetion部分：

• mha：直接按heads進行tp_size切分，heads必須能被tp_size整除
• gqa：對於query group部分進行tp_size切分，query_groups必須能被tp_size整除
• output_layer：將output_layer按行劃分，mha/gqa的結果經過tp切分後本質是按列劃分，和對應的output_layer按行劃分後，相乘的結果最後allreduce。
• 額外帶來的通訊開銷是： 2*b*s*h

MLP部分：

• 輸入不變，第一個矩陣（h*4h）按列切，每張卡運算的結果為 b*s*4h/tp ，其值進行啟用運算，第二個矩陣(4h*h)按照行切，乘法後得到的結果是 b*s*h ，對這個結果做allreduce，其通訊開銷是 2*b*s*h
• 為什麼不將第一個矩陣按行切第二個矩陣按列切？因為這樣第一個矩陣的結果就是 bs4h ，在啟用之前還得allreduce，即沒有節省空間，還有額外通訊。

output_layer部分：

• output_layer是 h*volcab_size ，輸入是 b*s*h 。
• 將output_layer按列劃分，每張卡的結果為 (b*s)*(volcab_size/tp) ，由於最後結果算loss時需要交叉熵，要對每一行做softmax，因此還需要再做一次allreduce。由於vocab_size很大(>100k)，這個通訊量就很誇張了，一種最佳化方式是先將每張卡上的結果按行sum，得到的結果是 b*s ，對這個求和結果做allreduce，其通訊量僅為 b*s*tp

張量並行帶來的額外通訊開銷是 o(b*s*h*l) ，其中l是layers

張量並行一般僅在單機8卡內做，因為通訊量太大了，只有單機內nvlink這種高頻寬才能支援（H100 900GB/s卡間通訊），一旦跨機即便是IB也不夠。

在實現過程中，tp有一些要注意的點：

• 將輸入送到列並行之前，需要將經過一個特殊的函式，前向是identiy，反向是allreduce。反向之所以是allreduce而不是其它操作（如reducescatter），是因為鏈式法則：，中間的加好就是allreduce
• 在行並行的結果allreduce時，它的反向是一個identity，這是因為sum的偏導就是1：

TP過程示意圖

2.3 流水線並行（Pipeline Parallelism）

上面說的張量並行是對矩陣進行切分，缺點是通訊量很大。

另一種思路是按layer進行切分，例如一個模型是80層，切分到8張卡上，每一張卡放10層。

通訊量僅為 b*s*h*pp ，小於TP的通訊量

帶pp會帶來一個嚴重的問題，那就是氣泡率較高，另外沒有使用梯度累加，導致視訊記憶體消耗比較高，和micro-batchs成正比。

可以將GBS切分為多個micro_batch，儘量將pipeline填滿，從而一定程度減緩氣泡率問題。它的氣泡率為 ,其中p為pipeline stage，m為micro-batchs，，分別是前向反向的耗時。簡單來說就是數device-4的灰色格子就好，前面有(p-1)個氣泡，後面有(p-1)個氣泡

透過1F1B可以減少視訊記憶體消耗，氣泡率不變，1F1B的思想很簡單，只要有一個mb完成了，立馬進行反向，從而減少了視訊記憶體佔用，從之前的m下降至p。但是不會減少氣泡率，因為流水線的結構其實沒有變，只不過優先處理了已經完成forward的micro-batch

為了進一步降低氣泡率，提出了interleaved 1F1B schedule，又可以叫virtual-pipeline。相當於把流水線虛擬化了，本來4個gpu只能做4層流水，把它變成了8層甚至更多層。具體操作如下，例如2個gpu，每個gpu分攤4層layer（gpu1是0~3 layers，gpu2是4~7 layers），那麼流水線虛擬化後就變為：gpu1是0~1，gpu2是2~3，緊接著又是gpu1為4~5，gpu2為6~7，這樣就硬生生的用2個gpu實現了4層流水，這樣氣泡率也就降低了2倍。假設v為virtual stage，那麼virtual pipeline將的氣泡率為，這是因為每個pipeline stage的fwd和bwd時間降低至原有的。

24年的zero-bubble^[2]進一步降低了氣泡率，它的思路很簡單，將計算粒度進一步拆細，是將bwd進一步拆細，拆分為引數的梯度和輸入的梯度，算完後立即傳輸給下一個stage。為什麼將計算粒度變細能降低氣泡率？因為氣泡的存在本質是因為要等待其他裝置的資料，如果能將資料切分至無限小，並且傳輸延遲是無限大的話，那氣泡就沒有了。這篇文章中有提供了2種方法：

• 第一種方法，就是上述所說，將bwd拆細，視訊記憶體和1F1B一樣，氣泡率降低為1F1B的1/3，只存在最初的warmup階段的氣泡，BWD的氣泡被完全填滿。但是每個step都會有warmup的氣泡。第二個方法解決了這一點。
• 第二種方法，除了將bwd拆細外，在warmup階段送入兩倍的資料，同時我們知道在optimizer.step之前，一般為了數值穩定新需要allreduce全域性梯度，做一些nan check和clip（例如發現有nan就直接跳過這次更新）。但這種方法不會做全域性的allreduce，每個stage做自己的nan check和clip，然後直接進入下一個迭代（一旦某個stage發現有nan或者clip，就通知所有stage回滾）。這樣視訊記憶體翻倍，但氣泡率幾乎為0了

DeepSeek V3結合zero-bubble pipeline和Chimera 雙流排程提出了Dual Pipe，還沒看明白。後面補充。

2.4 序列並行（Sequence Parallelism）

這裡介紹的序列並行是Megatron-LM中實現的版本，與處理長序列的SP貌似不是一回事，後者的代表是DeepSpeed-Ulysses和Ring Attention。

TP主要處理的是attention和FFN維度的並行，主要是對引數矩陣進行切分，而輸入還是有冗餘（的每個TP在attention和MLP中都有相同的輸入）。SP則是對輸入矩陣按sequence維度進行切分，體現在layernom和dropout層（layernorm和droupout對序列長度是無感的，SP不會影響計算結果），收益是減少了啟用值。

由於attention和mlp都需要完整的序列，所以在進入attention和mlp需要做allgather，反向時需要做reduce-scatter（從計算圖中能看出是reduce-scatter）；從attention和mlp出來時做reduce-scatter，返現做allgather。

SP一般是基於TP的，一旦開啟了SP，TP的演算法也要有所改變，主要是通訊操作，如圖4所示

SP相比於TP不會帶來更多的通訊開銷，這是因為TP的通訊操作是allreduce，SP+TP的通訊操作是reduscatter+allgather，兩者開銷相同。

2.5 專家並行（Expert Parallelism）

專家並行主要解決的是MoE模型的問題，moe模型結構如下圖：

由於每個token只會分發給topk個專家，導致每個專家（FFN）只處理一部分token，計算效率很低。為了提升計算效率，會將專家分攤在不同的DP上。在之前的各種並行中，每個DP擁有完整的模型，在FWD和BWD時也是各自處理各自的資料，不會有任何通訊（只會在梯度allreduce時涉及DP間通訊）。而在EP中，專家被分攤到不同的DP裡面，在前向和反向時DP之間便會有資料的通訊，如圖：

一般會透過all2all進行資料的分發，同時並行組也會複雜不少。

引用連結

[1] 李沐老師的影片: https://www.bilibili.com/video/BV1nB4y1R7Yz/?spm_id_from=333.1387.search.video_card.click[2] zero-bubble: https://arxiv.org/pdf/2401.10241