1）模型參數(shù)過大，如GLM 130B模型參數(shù)需要520GB（130B*4bytes）的顯存，參數(shù)梯度為520GB，優(yōu)化器狀態(tài)需要1040GB，共計(jì)2096GB的顯存，即使顯存最大的GPU也放不下大模型的模型參數(shù)；

2）即使可以將模型參數(shù)放在一張GPU上，對(duì)海量數(shù)據(jù)進(jìn)行預(yù)訓(xùn)練也需要耗費(fèi)很長(zhǎng)時(shí)間，單卡如A100訓(xùn)練130B參數(shù)量的模型大約需要232年。分布式預(yù)訓(xùn)練技術(shù)，就是在這樣的背景下應(yīng)運(yùn)而生，它允許模型在多個(gè)計(jì)算節(jié)點(diǎn)上并行訓(xùn)練，大幅度縮短了訓(xùn)練時(shí)間，本文總結(jié)了目前常見的分布式訓(xùn)練的常見方法。

一、分布式訓(xùn)練

1.數(shù)據(jù)并行

數(shù)據(jù)并行指的是數(shù)據(jù)集切分為多份，每張GPU分配不同的數(shù)據(jù)進(jìn)行訓(xùn)練，每個(gè)進(jìn)程運(yùn)行完整模型，在每個(gè)進(jìn)程上獨(dú)立運(yùn)行前向和反向過程，最后通過AllReduce操作，將多張卡上的梯度進(jìn)行同步。

▲ 圖1 數(shù)據(jù)并行

2.張量并行

張量并行是指將一個(gè)獨(dú)立的層劃分到不同的GPU上，例如用兩塊GPU實(shí)現(xiàn)一個(gè)并行的GEMM矩陣乘法運(yùn)算Y=X*W，可以將X和W分別沿著列方向和行方向切分成兩份，通過矩陣乘法的分塊規(guī)則，分別在每個(gè)GPU上做子矩陣乘法，然后將兩個(gè)子矩陣乘法結(jié)果相加得到最終結(jié)果

針對(duì)transformer模型MLP的張量并行實(shí)現(xiàn)方式圖2所示，第一個(gè)GEMM輸入X不做切分，直接通過f操作拷貝到兩個(gè)GPU上，權(quán)重A沿列方向進(jìn)行切分，每個(gè)GPU上保存一半的運(yùn)算結(jié)果，通過GELU運(yùn)算得到Y(jié)1和Y2。第二個(gè)GEMM由于Y1和Y2已經(jīng)分別存在于兩個(gè)GPU上，相當(dāng)于已經(jīng)做了沿列方向的切分，所以對(duì)權(quán)重B做行方向的切分，分別進(jìn)行矩陣乘法得到Z1和Z2，g操作是一個(gè)AllReduce操作，用來實(shí)現(xiàn)公式中的加法，AllReduce之后，每個(gè)GPU上都有一份完整的Z矩陣。

▲ 圖2 MLP張量并行

針對(duì)transformer模型Atttention的張量并行如圖3所示，QKV的三個(gè)矩陣的權(quán)重shape均為是[dim, head_hum*head_dim]，考慮到多頭注意力天然的并行性，可以將三個(gè)權(quán)重沿著列方向進(jìn)行切分，每張卡上權(quán)重的shape為[dim, head_hum*head_dim/N]，每張卡上不同部分head的softmax等計(jì)算式互相獨(dú)立，分別得到Y(jié)1和Y2，再對(duì)輸出的權(quán)重 B沿著行方向切分，得到計(jì)算結(jié)果Z1和Z2，再通過g做AllReduce操作得到最終的Z?？梢钥吹剑繌埧ㄉ献龅娜匀皇且粋€(gè)完整的Attention運(yùn)算，只是head number變?yōu)榱嗽瓉淼?/N。

▲ 圖3 Attention 張量并行

3.流水線并行

流水線并行是指將不同的layer劃分到不同的GPU上，為了提升節(jié)點(diǎn)的利用率，將一個(gè)batch，繼續(xù)拆分為 M 個(gè)micro batch，當(dāng)節(jié)點(diǎn)1計(jì)算完第一個(gè)micro batch, 節(jié)點(diǎn)2開始計(jì)算后，節(jié)點(diǎn) 1可以繼續(xù)計(jì)算第二個(gè)micro batch，對(duì)于P個(gè)節(jié)點(diǎn)，M個(gè)micro batch的流水線并行，流水線氣泡占比為(P-1)/M，為減少氣泡占比，需保證M>>P。

▲ 圖4 流水線并行

4.序列并行

序列并行是在 張量的基礎(chǔ)上，將 Transformer 層中的 LayerNorm 以及 Dropout 的輸入按輸入長(zhǎng)度（Sequence Length）維度進(jìn)行了切分，使得各個(gè)設(shè)備上面只需要做一部分的 Dropout 和 LayerNorm，LayerNorm 和 Dropout 的計(jì)算被平攤到了各個(gè)設(shè)備上，減少了計(jì)算資源的浪費(fèi)，進(jìn)一步降低了顯存開銷。

▲ 圖5 序列并行

4.激活重計(jì)算

激活重計(jì)算是指在前向計(jì)算時(shí)，只把檢查點(diǎn)節(jié)點(diǎn)保存在顯存中，檢查點(diǎn)以外的中間值全部被丟棄，在方向傳播計(jì)算梯度時(shí)，從最近的檢查點(diǎn)開始，其余節(jié)點(diǎn)重新計(jì)算一次前向，激活重計(jì)算需要額外的計(jì)算，是一個(gè)以時(shí)間換空間的策略。

5.ZeRO零冗余優(yōu)化器

數(shù)據(jù)并行在通信和計(jì)算效率上表現(xiàn)良好，但存在內(nèi)存冗余問題，ZeRO通過分區(qū)參數(shù)和狀態(tài)來減少冗余，使每個(gè)GPU只保留部分?jǐn)?shù)據(jù)。圖6展示了啟用內(nèi)存優(yōu)化后各設(shè)備的內(nèi)存消耗，其中Ψ代表模型大小，K是優(yōu)化器狀態(tài)的內(nèi)存乘數(shù)，Nd是數(shù)據(jù)并行度。

▲ 圖6 ZeRO零冗余優(yōu)化

假設(shè)模型大小為Ψ=75億，基于Adam優(yōu)化器的混合精度訓(xùn)練，數(shù)據(jù)并行度為Nd=64，K=12，我們?cè)跀?shù)據(jù)并行的基礎(chǔ)上，分三步對(duì)顯存進(jìn)行優(yōu)化

Pos（優(yōu)化器狀態(tài)優(yōu)化）

我們?cè)诿總€(gè)GPU中保存全部的參數(shù)和梯度，但是只保存1/Nd的優(yōu)化器狀態(tài)，這會(huì)導(dǎo)致總的顯存消耗變?yōu)?Ψ+2Ψ+KΨ/Nd , 在圖示假設(shè)下為顯存占用31.4GB。

Pos+g（優(yōu)化器狀態(tài)和梯度優(yōu)化）

優(yōu)化器狀態(tài)優(yōu)化的基礎(chǔ)上增加對(duì)梯度內(nèi)存的優(yōu)化，即每個(gè)GPU中只保存1/Nd的梯度，這會(huì)導(dǎo)致總的顯存消耗變?yōu)?Ψ+（2Ψ+KΨ）/Nd , 在圖示假設(shè)下為顯存占用為16.6GB。

Pos+g+p（優(yōu)化器狀態(tài)、梯度以及參數(shù)優(yōu)化）

優(yōu)化器狀態(tài)和梯度優(yōu)化的基礎(chǔ)上增加對(duì)參數(shù)顯存的優(yōu)化，即每個(gè)GPU中只保存1/Nd的參數(shù)，這將導(dǎo)致總的內(nèi)存消耗變?yōu)椋?Ψ+2Ψ+KΨ）/Nd ,在圖示假設(shè)下為顯存占用為1.9GB。

二、討論與展望

未來的LLM可能會(huì)采用更先進(jìn)的并行策略，以進(jìn)一步減少通信開銷并提高計(jì)算效率。例如探索使用分層并行策略，將模型的不同層次分布在不同的節(jié)點(diǎn)上進(jìn)行訓(xùn)練，以減少節(jié)點(diǎn)間的通信需求。此外，還可以結(jié)合使用數(shù)據(jù)并行和模型并行的方法，根據(jù)硬件環(huán)境和模型規(guī)模靈活調(diào)整并行策略。分布式預(yù)訓(xùn)練技術(shù)的發(fā)展為大規(guī)模智算平臺(tái)提供了有效支撐，但仍然面臨非常大挑戰(zhàn)，通過精心設(shè)計(jì)解決方案并不斷探索新的并行策略，我們可以進(jìn)一步提高分布式預(yù)訓(xùn)練的效率和可擴(kuò)展性，為構(gòu)建更大規(guī)模的模型提供支持。

參考文獻(xiàn)

[1] Korthikanti, Vijay, Jared Casper, Sangkug Lym, Lawrence McAfee, Michael Andersch, Mohammad Shoeybi, Bryan Catanzaro. 《Reducing Activation Recomputation in Large Transformer Models》.

[2] Narayanan, Deepak, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Korthikanti, Dmitri Vainbrand. 《Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM》. Proceedings of the International Conference for High Performance Computing, Networking, Storage and Analysis, 1–15. St. Louis Missouri: ACM, 2021. https://doi.org/10.1145/3458817.3476209.

[3] Shoeybi, Mohammad, Mostofa Patwary, Raul Puri, Patrick LeGresley, Jared Casper, Bryan Catanzaro. 《Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism》.

[4] Rajbhandari, Samyam, Jeff Rasley, Olatunji Ruwase和Yuxiong He. 《ZeRO: Memory Optimizations Toward Training Trillion Parameter Models》. arXiv, 2020年5月13日. http://arxiv.org/abs/1910.02054.

本文轉(zhuǎn)載自??AI遇見云??，作者： 何武 ????