LLM並行訓練7-混合並行總結

SunStriKE發表於2024-07-29
並行

概述

根據前面的系列文章, 對預訓練大模型裡用到的主要並行加速技術做了一系列拆分. 但是在實際的訓練裡往往是多種並行混合訓練. 我們要怎麼配置這些並行策略才能讓訓練框架儘可能的減少通訊瓶頸, 提升GPU計算利用率呢? 這裡的變數太多了, 以最簡單的3D並行為例:

  • 硬體層面有: 單臺機器的卡數/卡間頻寬/網路卡頻寬, 機器間通訊時的網路拓撲構建.
  • 並行策略上有: 張量並行數/流水線並行數/資料並行數
  • 訓練超參有: batch_size / AttnHeads / seq_len / hidden_size

如果靠腦補來調整這些引數, 會存在一個非常巨大的搜尋空間, 很難找到最優於計算效率的方法, 所以需要先透過理論分析確定各個引數的大致範圍. 最後再透過有限次嘗試找到較優的方案. 本章參考nvidia的調參實踐GTC演講, 結合GPT3訓練例子對如何調整並行策略進行總結

並行方法適用場景分析

後文的標記備註:

  • \((p, t, d)\) : 3D 並行維度. \(p\) 代表流水並行數， \(t\) 代表張量並行數， \(d\) 代表資料並行數
  • \(n\): 總共的 GPU 數量. 要求 \(p\cdot t \cdot d = n\).
  • \(B\): Global batch size.
  • \(b\): Microbatch size.
  • \(b^{'}\): 一個流水線要處理的 batch size 大小， 等於 \(B/d\).
  • \(m = \frac{1}{b} \cdot \frac{B}{d}\): 一個 batch 在每個 pipeline 的 microbatch 的數量.
  • \(s\): seq_len.
  • \(h\): hidden_size = emb_size * attention_head_size
  • \(a\): attention_head_size

張量並行(TP)

TP開銷

模式 Normal ColParallel ratio
flops （n次乘法 + n次加法）* n^2 = 2n^3 2n^3/t 1/t
Bandwidth (n^2)【n*n 矩陣的讀或寫】 * 2(fp16)) * 3(讀 X、讀A，寫 Y) = 6n^2 2n^2 + 4n^2/t(A,Y切分) (1+2/t)/3
Intensity(flops/bandwidth) n/3 n/(2+p) 3/(2+t)

當並行度\(t\)增長的時候, 可以看到intensity也處於一個增長的趨勢. 需要權衡通訊和計算成本的平衡, 由於TP需要在結束時進行一次啟用的AllReduce, 在多機通訊上會導致較高的通訊成本. 所以TP一般只考慮在單機卡間通訊時使用. TP在LLM裡主要有兩個使用場景:

  1. MLP先列再行, 這塊前後一般會和SP結合進行將AllReduce拆分為allGather和reduceScatter
image-20240727180417038
  1. attention處多頭切分並行 每個頭之間的計算各自獨立, 所以可以進行切分計算.
image-20240727180501261

流水線並行(PP)

流水線主要是將一個batch的資料切分為多個mirco-batch, 在micro-batch之間做非同步並行. 因為通訊內容只包含切分stage的輸出, 而且是點對點通訊, 不需要多點集合通訊. 通訊資料量小, 因此比較適合在多臺機器間通訊的場景. LLM裡一般把一個transformLayer作為一個stage, 在多個stage之間構建pipeline, 如下圖:

image-20240727182819634

混合並行

當網路結構確定後, 一般TP和PP就能估算到比較合理的區間, 最後根據視訊記憶體容量的計算來估計DP需設定的值.

image-20240727195029105

TP與PP的策略分析

資料並行度\(d=1\)時, \(p * t = n\), 會有以下計算公式:

  • 流水線bubble_time: \(\frac{(p-1)}{m}=\frac{n/t-1}{m}\), 提高TP並行度時會減少氣泡佔比, 但會增大單機內部的通訊量, tp內部一個microbatch需要4個allReduce(fp/bp各兩個)

  • 單機單次allReduce通訊量: \(2bsh(\frac{t-1}{t})\), (layer啟用為\(bsh\), allReduce通訊量為資料量2倍)

  • 流水線並行時單個micro-batch機器間通訊量為: \(C_{inter} = 2bsh\) (fp/bp各一次)

設一個pipeline內有\(l^{stage}\)個transformLayer, 則在1F1B非交錯排程的情況下單個stage單機內部通訊量為:

\[C_{inra} = l^{stage}\cdot4\cdot2bsh(\frac{t-1}{t}) = l^{stage}\cdot4\cdot2bsh(1-\frac{1}{t}) \]

所以機器間和機器內的通訊量關係為:

\[C_{intra}= l^{stage}\cdot4\cdot(1-\frac{1}{t}) \cdot C_{inter} \]

因為機器間通訊速率遠小(IB 200GB/s)於卡間通訊(NVLink 600GB/s), 所以我們如果希望最佳化吞吐, 那麼應該儘量降低機器間通訊比率.

[!TIP]

也就是在不會導致TP產生機器間通訊的前提下讓t儘可能的大. 如果這樣還放不下模型，再使用流水線並行來切分模型。

image-20240727212003374
micro-batch設定

在固定其他引數的前提下. 只調整micro_batch數, 單個batch的執行時間: \((\frac{b^{'}}{b}+(p-1))\cdot(t_{f} + t_{b})\) , 如果增大b, 單個pipeline內數量減少但執行時間會變長, 計算耗時和b是非線性的關係. 而且調整micro-batch後, 通訊耗時也會變化, 所以mirco-batch調整需要實驗嘗試才能找到最優解. megatron在論文中嘗試gpt訓練的mirco-batch設定4比較合適

image-20240727220854686

DP的策略分析

便於分析設\(t=1, d * p = n\), 在這種情況下的流水線bubble佔比為 \(\frac{p-1}{m} = \frac{n/d - 1}{B/b/d} = \frac{b(n - d)}{B}\)

PP和DP關係: 對於d單調遞減, 也從下圖可以看到, 當流水線並行的數量越小, 資料並行度越大的時候訓練速度越快. 所以我們可以在PP滿足視訊記憶體佔用的情況下儘可能的提升DP並行度.

和Batch_size關係: bubble和B成反比, B越大吞吐越高. 但是過大的B和資料並行度會導致模型不收斂. 需要在不影響效果的前提下調整B

image-20240727213002586

DP和TP關係: 在TP內每個batch 需要進行4次allReduce, 而DP只需要對梯度做一次allReduce, 另外在TP的時候如果W比較小也會影響矩陣乘法計算效率. 下圖可以看到TP並行度越小, DP並行度越大吞吐越高. 調整策略是隻要TP滿足視訊記憶體需求, 就儘可能的增大DP來提升吞吐.

image-20240727220023520

[!TIP]

如果模型比較大，需要先組合模型並行和流水並行，\(M=t \cdot p\) 的組合用來滿足模型和模型相關的資料的視訊記憶體需求，但是要讓 M 儘量小。之後使用資料並行來擴充套件訓練規模（擴大資料並行度、擴大 Global batch size）

GPT-3例子分析

以如下的超參GPT-3訓練為例:

image-20240729183643039

視訊記憶體分析

ModelMemory

單卡儲存模型引數主要包含4個部分(由於流水線並行, 單卡一般只儲存1-2個transformLayer): attention引數 / FC引數 / token_emb / Positional encoding

image-20240729184709396

以\(N_p\)代表一份全量引數, 則單卡上包含的引數量如下:

\[\frac{N_p}{n} = h * \frac{h}{t}*3(QKV引數)+h*\frac{h}{t}(multihead拼接後經過的fc) + h*\frac{h}{t}*4*2(fc1+fc2引數) \\ + \frac{v}{t} *h(token) + s*\frac{h}{t}(positional) \approx 1.73B\approx \frac{175B}{p*t} (單卡上單份DP引數量) \]

在混合精度訓練中, 總共的資料量包擴1份fp16的w和grad, 1份fp32的optimizer_state(\(w+grad+momentum+variance\))

\[N_{storage} = 2Bytes * N_p + 2Bytes * N_p + (4+4+4+4)Bytes*N_p = 20N_P = 27.4GB \]

Activation

在nvidia分享裡, 看著activation只存了過token前的emb啟用和進fc前的啟用, 剩下的全部都是bp時重計算的..因為也沒使用SP, 這裡每張卡的啟用都存了TP並行數的冗餘資料

\[M_{act}^{emb} = 2𝐵𝑦𝑡𝑒𝑠 * 𝑏 * 𝑠 * \frac{v}{t} \]

\[M_{transformer}^{emb} = 2𝐵𝑦𝑡𝑒𝑠 * 𝑏 * 𝑠 * h * \frac{n}{p} \]

Extra

包含在fp時所需分配的臨時視訊記憶體 & 通訊需要的臨時視訊記憶體 & allocator導致的視訊記憶體碎片

這塊在上一章-啟用最佳化裡其實已經分析過, 這裡忽略不表

\[\begin{gathered} M_{\text {extra }}^{\text {embed }}=4 \text { Bytes } \times b \times s \times \frac{v}{t}=0.05 G B \\ M_{\text {extra }}^{\text {allreduce }}=(2+4) \text { Bytes } \times N_p \\ M_{\text {extra }}^{q k v}=2 \text { Bytes } \times\left(s \times \frac{h}{t} \times b \times 3+b \times s \times s \times \frac{a h}{t} \times 2\right) \\ M_{\text {extra }}^{m l p}=2 \text { Bytes } \times\left(b \times s \times \frac{4 h}{t}+b \times s \times \frac{4 h}{t}+b \times s \times \frac{h}{t}\right) \end{gathered} \]

視訊記憶體峰值

視訊記憶體使用可能出現峰值的地方有三個地方:

  • fp完成時, 這裡主要變化量是存在大量的extra: \(M_1=M_{p a r a}+M_{o p t}+p \times M_{a c t}^{\text {transformer }}+M_{a c t}^{e m b e d}+M_{e x t r a}^{e m b e d}\) , 視訊記憶體消耗32.2GB
  • bp完成時, 這裡因為把大量的act釋放後重計算,相對消耗不多: \(M_2 =M_{p a r a}+M_{o p t}+p \times M_{a c t}^{t r a n s f o r m e r}+M_{g r a d}-M_{a c t}^{t r a n s f o r m e r} +M_{e x t r a}^{q k v}+M_{e x t r a}^{m l p}\), 視訊記憶體消耗25.0GB
  • 更新optimizer_state時, 這是由於很多臨時視訊記憶體用於梯度allReduce, 所以出現視訊記憶體峰值: \(M_3= M_{p a r a}+M_{o p t}+p \times M_{a c t}^{\text {transformer }}+M_{g r a d}-M_{a c t}^{\text {transformer }} +M_{\text {extra }}^{\text {allreduce }}\), 視訊記憶體消耗34.3GB

通訊分析

BW: bus bandwidth(單次通訊的資料長度)

TP: 每個mlp和attention 在fp /bp / bp時的fp重計算 三個階段各需要一次allReduce

\[T_{tp} = \frac{2bsh}{BW} * \frac{2(t-1)}{t}[allReduce通訊次數*單次通訊長度]*(3+3)[mlp和attention\ allReduce次數] *\frac{n}{p}[layer數] * \frac{B}{bd}[minibatch數] \\ \]

DP: 在optimizer更新時需要對各個資料副本進行一次allReduce

\[T_{dp} = \frac{N_p}{BW} * \frac{2(d-1)}{d} \]

PP(1F1B交錯式): 在機器間通訊的點對點方式和在機器內通訊的allGather(TODO: 這裡沒太看懂)

\[T_{p p}=\underbrace{\left(2 \frac{B}{b d}+2(p-1)\right) \times \frac{\frac{\text { message size }}{t}}{B W_{\text {inter }}}}_{\text {P2P }}+\underbrace{\left(\frac{B}{b d}+(p-1)\right) \times \frac{\text { message size }}{B W_{\text {intra }}} \times \frac{t-1}{t}}_{\text {Allgather }} \]

從實際實驗上也能看到TP佔了主要的通訊成本.

image-20240729211705948

總結

3d並行的調優經驗:

  1. 如果模型比較大，需要先組合模型並行和流水並行，\(M=t \cdot p\) 的組合用來滿足模型和模型相關的資料的視訊記憶體需求，但是要讓 M 儘量小。之後使用資料並行來擴充套件訓練規模（擴大資料並行度、擴大 Global batch size）
  2. 在不會導致TP產生機器間通訊的前提下讓t儘可能的大. 如果這樣還放不下模型，再使用流水線並行來切分模型。

參考

megatron-LM複雜度分析論文: https://arxiv.org/pdf/2104.04473

nvidia GTC演講: https://developer.nvidia.com/gtc/2020/video/s21496

[nvidia GTC GPT-3調參分析](連結: https://pan.baidu.com/s/190TFeOI9SALaaH9CVMWH7Q?pwd=chux 提取碼: chux)

megatron分析部落格: https://www.cnblogs.com/rossiXYZ/p/15876714.html

相關文章