LLM分布式預訓練淺析
1）模型參數過大，如GLM 130B模型參數需要520GB（130B*4bytes）的顯存，參數梯度為520GB，優化器狀態需要1040GB，共計2096GB的顯存，即使顯存最大的GPU也放不下大模型的模型參數；
2）即使可以將模型參數放在一張GPU上，對海量數據進行預訓練也需要耗費很長時間，單卡如A100訓練130B參數量的模型大約需要232年。分布式預訓練技術，就是在這樣的背景下應運而生，它允許模型在多個計算節點上并行訓練，大幅度縮短了訓練時間，本文總結了目前常見的分布式訓練的常見方法。
一、分布式訓練
1.數據并行
數據并行指的是數據集切分為多份，每張GPU分配不同的數據進行訓練，每個進程運行完整模型，在每個進程上獨立運行前向和反向過程，最后通過AllReduce操作，將多張卡上的梯度進行同步。
▲ 圖1 數據并行
2.張量并行
張量并行是指將一個獨立的層劃分到不同的GPU上，例如用兩塊GPU實現一個并行的GEMM矩陣乘法運算Y=X*W，可以將X和W分別沿著列方向和行方向切分成兩份，通過矩陣乘法的分塊規則，分別在每個GPU上做子矩陣乘法，然后將兩個子矩陣乘法結果相加得到最終結果
針對transformer模型MLP的張量并行實現方式圖2所示，第一個GEMM輸入X不做切分，直接通過f操作拷貝到兩個GPU上，權重A沿列方向進行切分，每個GPU上保存一半的運算結果，通過GELU運算得到Y1和Y2。第二個GEMM由于Y1和Y2已經分別存在于兩個GPU上，相當于已經做了沿列方向的切分，所以對權重B做行方向的切分，分別進行矩陣乘法得到Z1和Z2，g操作是一個AllReduce操作，用來實現公式中的加法，AllReduce之后，每個GPU上都有一份完整的Z矩陣。
▲ 圖2 MLP張量并行
針對transformer模型Atttention的張量并行如圖3所示，QKV的三個矩陣的權重shape均為是[dim, head_hum*head_dim]，考慮到多頭注意力天然的并行性，可以將三個權重沿著列方向進行切分，每張卡上權重的shape為[dim, head_hum*head_dim/N]，每張卡上不同部分head的softmax等計算式互相獨立，分別得到Y1和Y2，再對輸出的權重 B沿著行方向切分，得到計算結果Z1和Z2，再通過g做AllReduce操作得到最終的Z。可以看到，每張卡上做的仍然是一個完整的Attention運算，只是head number變為了原來的1/N。
▲ 圖3 Attention 張量并行
3.流水線并行
流水線并行是指將不同的layer劃分到不同的GPU上，為了提升節點的利用率，將一個batch，繼續拆分為 M 個micro batch，當節點1計算完第一個micro batch, 節點2開始計算后，節點 1可以繼續計算第二個micro batch，對于P個節點，M個micro batch的流水線并行，流水線氣泡占比為(P-1)/M，為減少氣泡占比，需保證M>>P。
▲ 圖4 流水線并行
4.序列并行
序列并行是在 張量的基礎上，將 Transformer 層中的 LayerNorm 以及 Dropout 的輸入按輸入長度（Sequence Length）維度進行了切分，使得各個設備上面只需要做一部分的 Dropout 和 LayerNorm，LayerNorm 和 Dropout 的計算被平攤到了各個設備上，減少了計算資源的浪費，進一步降低了顯存開銷。
▲ 圖5 序列并行
4.激活重計算
激活重計算是指在前向計算時，只把檢查點節點保存在顯存中，檢查點以外的中間值全部被丟棄，在方向傳播計算梯度時，從最近的檢查點開始，其余節點重新計算一次前向，激活重計算需要額外的計算，是一個以時間換空間的策略。
5.ZeRO零冗余優化器
數據并行在通信和計算效率上表現良好，但存在內存冗余問題，ZeRO通過分區參數和狀態來減少冗余，使每個GPU只保留部分數據。圖6展示了啟用內存優化后各設備的內存消耗，其中Ψ代表模型大小，K是優化器狀態的內存乘數，Nd是數據并行度。
▲ 圖6 ZeRO零冗余優化
假設模型大小為Ψ=75億，基于Adam優化器的混合精度訓練，數據并行度為Nd=64，K=12，我們在數據并行的基礎上，分三步對顯存進行優化
Pos（優化器狀態優化）
我們在每個GPU中保存全部的參數和梯度，但是只保存1/Nd的優化器狀態，這會導致總的顯存消耗變為2Ψ+2Ψ+KΨ/Nd , 在圖示假設下為顯存占用31.4GB。
Pos+g（優化器狀態和梯度優化）
優化器狀態優化的基礎上增加對梯度內存的優化，即每個GPU中只保存1/Nd的梯度，這會導致總的顯存消耗變為2Ψ+（2Ψ+KΨ）/Nd , 在圖示假設下為顯存占用為16.6GB。
Pos+g+p（優化器狀態、梯度以及參數優化）
優化器狀態和梯度優化的基礎上增加對參數顯存的優化，即每個GPU中只保存1/Nd的參數，這將導致總的內存消耗變為（2Ψ+2Ψ+KΨ）/Nd ,在圖示假設下為顯存占用為1.9GB。
二、討論與展望
未來的LLM可能會采用更先進的并行策略，以進一步減少通信開銷并提高計算效率。例如探索使用分層并行策略，將模型的不同層次分布在不同的節點上進行訓練，以減少節點間的通信需求。此外，還可以結合使用數據并行和模型并行的方法，根據硬件環境和模型規模靈活調整并行策略。分布式預訓練技術的發展為大規模智算平臺提供了有效支撐，但仍然面臨非常大挑戰，通過精心設計解決方案并不斷探索新的并行策略，我們可以進一步提高分布式預訓練的效率和可擴展性，為構建更大規模的模型提供支持。
參考文獻
[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.
