一、引言

自 Transformer 架構問世以來，大模型領域進展迅速，短短幾年間模型參數(shù)規(guī)模攀升至天文數(shù)字，輕松跨過萬億大關。面對如此龐然大物，傳統(tǒng)的單機單 GPU 訓練方式顯然力不從心。因此，單機多 GPU、多機多 GPU 的分布式訓練方案應運而生，成為駕馭超大規(guī)模模型的必要手段。然而，這一轉變并非一帆風順，各種挑戰(zhàn)隨之而來，如：顯存要求高、訓練周期長、模型吞吐要求高。

為了攻克這些難點，一系列以提高效率為目標的并行計算策略和內存優(yōu)化技術被引入到分布式訓練過程中，編織了一張廣闊的技術網絡。智能工程部結合集群算力調度，充分考慮任務特點和硬件環(huán)境，融入并行優(yōu)化技術，精心調整參數(shù)配置，最終實現(xiàn)訓練效率和資源利用率的雙重飛躍。

本文將簡單介紹這些技術在分布式訓練中的作用，并用實際案例展示如何利用這些技術提高大模型訓練和微調的效率。

2. 相關技術原理

2.1 并行技術

隨著大模型訓練性能的要求，各種并行技術被應用于分布式訓練，主要有：

1、數(shù)據(jù)并行（DP）：將訓練數(shù)據(jù)分成多個batch，每個batch分配給不同的設備進行訓練，每個設備都有完整的模型，在本地計算loss和梯度，然后通過同步所有設備的梯度來更新模型參數(shù)。

2、流水線并行（PP）：將模型按層劃分為多個階段，每個階段由不同的設備處理。當每個設備處理完一個微批次時，可以并行處理其他微批次，從而提高設備利用率，減少等待時間。

3. 張量并行（TP）：將模型中的張量分成更小的塊，分配給不同的計算節(jié)點，每個節(jié)點只計算部分張量，節(jié)點之間需要進行通信將結果合并。

4.序列并行（SP）：主要用于訓練長序列數(shù)據(jù)，將長序列拆分成多個子序列并同時在不同的設備上處理這些子序列。

5、上下文并行（CP）：可以看作是SP的加強版，它將輸入輸出激活全部在序列維度上拆分，分別在不同的設備上進行計算，可以有效提升長上下文相關任務的訓練效率。

下面將詳細介紹各種并行技術的原理和特點。

2.1.1 數(shù)據(jù)并行

圖2.1 數(shù)據(jù)并行圖

數(shù)據(jù)并行在每個 worker 上復制一個模型，并將數(shù)據(jù)集劃分到多個 worker 上；每個 worker 接收不同的批次數(shù)據(jù)并執(zhí)行前向傳播，在反向，worker 定期匯總其梯度，以確保所有 worker 看到的權重版本一致。這種方法可以有效提高模型的訓練速度。

2.1.2 模型并行

根據(jù)拆分方式不同，模型并行分為流水線并行和張量并行。

2.1.2.1 流水線并行

流水線并行就是將模型逐層切分[1]，將模型的不同層劃分為不同的stage，放在不同的worker上。如圖2.2所示，前幾層放在一個worker上，后幾層放在另一個worker上。每部分計算完成后，數(shù)據(jù)再傳遞到下一個stage，直至整個模型計算完成。

圖2.2 層分割示意圖

但這種劃分方式會在訓練過程中產生大量的時間泡沫，如圖2.3所示，由于在前一個worker還未訓練完成時，下一個worker就一直在等待，極大的浪費了GPU的資源。

圖2.3 簡單流水線圖

對于這種 F-then-B（Forward then Backward）方案，Megatron-LM 采用的是 1F1B（One Forward pass following one Backward pass）策略 [2]。如圖 2.4 所示，一批訓練數(shù)據(jù)被分成多個微批次，每一批數(shù)據(jù)訓練完成后，立即傳遞給下一個 worker，當前 worker 立即開始訓練下一個微批次。這樣，bubble 時間就縮短了，GPU 利用率也提高了。如果采用交錯調度，bubble 時間還可以進一步縮短。

圖2.4 1F1B管線示意圖

2.1.2.2 張量并行

張量并行是模型的層內切分，如圖2.5所示，主要用于加速矩陣運算。將輸入矩陣和參數(shù)矩陣合理分塊計算，按照矩陣切分方式可分為行并行和列并行。

圖2.5 層內分割示意圖

1. 行并行

如圖2.6所示，若A按行分為A1和A2，則需要按列分為X1和X2，即：

圖2.6 矩陣按行拆分操作示意圖

這樣就將矩陣分成了塊，放在兩張卡上計算，計算完成后再通過All-Reduce獲取其他卡上的計算結果，得到最終的Y。

從正向和反向的角度看行并行過程，f算子在正向過程中會把X拆分成X1和X2，分別放在不同的GPU上，在反向過程中會通過All-Gather把返回的梯度拼接起來，f算子在正向過程中會通過All-Reduce把結果累加起來，在反向過程中則分別計算梯度。

圖2.7 并行化前逆向過程示意圖

2. 列平行度

如圖2.8所示，若將A按列拆分為A1和A2，則X不需要拆分，即：

圖2.8 矩陣列拆分操作示意圖

從正向和反向的角度看，列并行過程就是，在正向過程中，f算子會把X分散到兩張卡上，在反向過程中，通過All-Reduce把返回的梯度累加起來；在正向過程中，f算子會通過All-Gather把結果連接起來，在反向過程中，會把梯度矩陣拆分成兩部分然后返回。

圖2.9 柱并行化前逆過程示意圖

3. Transformer 中的張量并行

Transformer結構由Attention模塊和MLP模塊組成，如圖2.10所示，Attention模塊由SelfAttention層+Dropout層組成；MLP模塊有兩個GEMM，第一個GEMM將維度從H改成4H，第二個GEMM將維度從4H改回H，中間使用了GeLU激活函數(shù)；各層的連接上也使用了殘差連接。

圖2.10 Transformer結構圖

對于 Attention 模塊，可以利用多頭 Attention 操作固有的并行性，拆分自注意力模塊，對查詢（Query，Q）、鍵（Key，K）、值（Value，V）矩陣進行列并行優(yōu)化，同時輸出線性層可以直接對 Attention 層的輸出進行操作。

圖 2.11 自注意力機制

對于 MLP 模塊中的第一個 GEMM，相應的操作可以表示為 Y = GeLU (XA)。如果將 A 拆分為行，則

但由于 GeLU 是一個非線性函數(shù)，

因此，如果使用行并行，則必須在 GeLU 操作之前添加 All-Reduce 同步點，以允許不同 GPU 之間交換信息。另一種并行方法是列并行。如果使用列并行，則可以將 GeLU 函數(shù)獨立應用于每個 GEMM 的輸出，對應于

這種方法不需要同步點。因此，第一個 GEMM 沿列并行劃分，第二個 GEMM 沿行劃分。GeLU 的輸出可以直接用作下一個 GeLU 的輸入，無需任何通信。最終的整體并行方法如圖 2.12 所示：

圖 2.12 MLP

2.1.3 順序并行

序列并行采用均分序列維度的策略來降低模型內存占用。Megatron-LM 在張量并行的基礎上，將 Transformer 層中的 LayerNorm 和 Dropout 操作的輸入按照序列長度維度進行拆分[3]，使得每個設備只需要處理一部分序列的 Dropout 和 LayerNorm 操作。

當然，額外的劃分會導致通信模式的變化。在Transformer層，TP的通信模式包括兩個正向All-Reduce和兩個反向All-Reduce。但是由于SP劃分了序列維度，傳統(tǒng)的All-Reduce已經不再適用。為了將SP生成的結果在各個設備上聚集，滿足SP與TP之間的數(shù)據(jù)傳輸需求，需要引入All-Gather；而為了將TP生成的結果傳遞到序列并行層，需要引入Reduce-Scatter。因此如圖2.13所示，引入了兩個算子和，分別代表正向All-Gather和反向Reduce-Scatterkaiyun體育登錄網頁入口，和代表正向Reduce-Scatter和反向All-Gather。由于一個All-Reduce相當于一個Reduce-Scatter和一個All-Gather，因此在開啟TP時開啟SP并不會增加額外的通信。 相反，在反向實現(xiàn)中，Reduce-Scatter和權重梯度的計算是重疊的，進一步減少了通信占用的時間。

圖 2.13 啟用 TP 和 SP 的 Transformer 層

2.1.4 語境平行

上下文并行也是一種序列維度上的并行化方案[4]，它對長文本處理任務有很好的優(yōu)化。與 SP 不同的是，SP 只在序列維度上拆分了 LayerNorm 和 Dropout 輸出的激活值，而 CP 則在序列維度上拆分了所有輸入的輸入和所有輸出的激活值。CP 可以看作是 SP 的增強版。

為了減少激活內存的使用，每個 GPU 在前向傳遞中只存儲一個序列塊的 KV，在后向傳遞中再次收集 KV。KV 通信發(fā)生在一個 GPU 和其他 TP 組對應 GPU 之間。在底層，這些 All-Gather 和 Reduce-Scatter 操作轉換為環(huán)形拓撲中的點對點通信。此外，還可以通過使用多查詢注意或分組查詢注意來交換 KV 以減少通信。

以圖2.14中的TP2-CP2 Transformer網絡為例，Attention前面的算子為CP的通信算子，其他的為TP的通信算子。AG代表All-Gather，RS代表Reduce-Scatter，AG/RS代表Forward All-Gather Reverse Reduce-Scatter，RS/AG代表Forward Reduce-Scatter Reverse All-Gather。

圖 2.14 啟用 TP 和 CP 的 Transformer 層

2.2 顯存優(yōu)化

2.2.1 訓練過程中的顯存分析

在大模型的訓練過程中，GPU的瓶頸往往不是它的算力，而是顯存的容量限制，占用顯存的部分主要有模型權重、梯度、優(yōu)化器、激活值等。

1. 模型重量

在模型訓練時，一般有fp32、fp16、bf16、int8幾種模型保存格式。fp32使用32位保存一個模型參數(shù)，int8使用8位，fp16和bf16使用16位保存。假設模型參數(shù)個數(shù)為Ψ，訓練時模型以bf16精度保存，其權重本身占用的顯存約為2Ψ(Bytes)(如無特殊說明，顯存單位也是Bytes)。

2.漸變

梯度數(shù)據(jù)類型通常和模型數(shù)據(jù)類型匹配，在混合精度訓練中，梯度類型一般為fp16，對于fp32訓練，梯度占用顯存4Ψ，對于fp16/混合精度訓練，梯度占用顯存2Ψ。

3. 優(yōu)化器

目前大模型訓練普遍使用Adam優(yōu)化器，訓練時需要用到動量和二階動量，如果優(yōu)化器使用fp32格式保存參數(shù)，其占用的顯存為12Ψ。

4. 激活值

在大型模型訓練中，Embedding 層和最后的 LayerNorm 和輸出層占用的內存與 Transformer 塊產生的內存相比可以忽略不計。

假設激活值的數(shù)據(jù)類型為fp16，輸入長度為s，batch size為b，hidden的維度為h，attention head的數(shù)量為a，模型層數(shù)為L，則：

各層Transformer block產生的激活值的內存占用為：

如圖2.15所示，開啟TP后，Transformer層部分模塊的顯存可以在不同設備間共享，不能共享的顯存主要由兩個LayerNorm block和兩個Dropout block產生，共計2個。

圖 2.15 啟用 TP 的 Transformer 層

假設張量并行度為t，Transformer每層顯存占用為：

如果再次開啟SP，Transformer層中的LayerNorm和Dropout塊也會進行分裂，在序列維度上對張量進行分裂，分裂次數(shù)等于張量并行度的大小。那么每層激活值的內存占用為：

若采用流水線并行，由于第一階段需要存儲p個微批次的激活值，因此無論并行規(guī)模p為多少，存儲第一階段激活值所需的顯存為：

如果使用交叉調度，則總活動視頻內存將根據(jù) (1+(p-1)/pm) 進行縮放，其中 m 是交叉管道的大小。

2.2.2 顯存優(yōu)化技術

顯存優(yōu)化技術主要通過減少數(shù)據(jù)冗余、用計算代替存儲、壓縮數(shù)據(jù)表示等方式來降低上述部分的顯存占用率。下面介紹兩種常用的節(jié)省顯存的方法。

1. 重新計算

激活重新計算技術，它僅存儲每層的輸入激活值，并在需要時重新計算其他激活值。

如果選擇根據(jù)TP重新計算，則顯存使用量為

如果再次打開 SP，視頻內存將

如果重新計算全部量，則視頻內存

雖然全量重新計算可以大幅降低激活值的內存占用，但是會導致30%-40%的計算性能損失，訓練時需要根據(jù)實際情況決定采用哪種重新計算方式。

2. ZeRO

ZeRO[5] 是零冗余優(yōu)化器的縮寫，是 DeepSpeed 項目為解決數(shù)據(jù)并行訓練中常見的內存冗余問題而提出的一種創(chuàng)新方法。在傳統(tǒng)的數(shù)據(jù)并行訓練中，模型參數(shù)、梯度和優(yōu)化器狀態(tài)在每個 GPU 上都是完全復制的，從而導致內存冗余。

ZeRO 技術通過將這些數(shù)據(jù)分布并共享到不同的 GPU 上，從而降低了每個 GPU 的內存需求。具體來說，ZeRO 提出了三種策略，分別對應分別存儲優(yōu)化器參數(shù)、分別存儲優(yōu)化器和梯度以及分別存儲三個參數(shù)。具體效果可見圖 2.16：

圖 2.16 ZeRO 的三個層級

假設模型參數(shù)Ψ=7.5B，GPU數(shù)量Nd=64，優(yōu)化器狀態(tài)的顯存乘數(shù)為K，對于Adam來說為K=12。對于Baseline（DP），每個GPU都要保存所有的參數(shù)、梯度和優(yōu)化器狀態(tài)，??假設使用fp16保存模型，則總顯存為（2+2+K）Ψ，也就是120GB。對于，將KΨ顯存分成Nd個GPU，因此每個GPU上的顯存為2Ψ+2Ψ+KΨ/Nd，也就是31.4GB。對于，將KΨ+2Ψ優(yōu)化器狀態(tài)和梯度分成Nd個GPU，每個GPU為2Ψ+（2+k）Ψ/Nd，也就是16.6GB。最后將全部除以2，因此每個GPU為（2+2+K）Ψ/Nd，也就是1.9GB。

與更高級別的 ZeRO 相比，保留了模型參數(shù)的完整復制，避免了計算過程中頻繁的數(shù)據(jù)交換，保持了較高的計算效率；與和相比，僅涉及優(yōu)化器狀態(tài)的并行化，而不涉及模型權重和梯度的分裂，實現(xiàn)相對簡單，這使得它在許多情況下成為實用的選擇。

3. QLM 加速實踐經驗

QLM，全稱奇虎語言模型，是智能工程部基于目前主流大模型加速框架 Megatron-LM 開發(fā)的一體化框架，支持 Huggingface、Megatron 多種模型格式的轉換、預訓練、評估、微調、分析等，結合 TAI 平臺可有效支持 Qwen、zhinao 等千卡訓練、長文本微調等各類任務。

圖3.1 QLM框架及功能圖

下面演示如何使用QLM框架支撐業(yè)務部門128k長文本微調提速實踐，通過以上優(yōu)化提速技術，長文本微調速率從120s/sample提升到35.5s/sample，提升3.4倍。

基礎版本在四臺H800機器上訓練，由于資源有限，優(yōu)化在四臺A100上進行云開·全站apply體育官方平臺，H800與A100的性能如下圖所示：

圖3.2 不同GPU性能參數(shù)對比

并行提速調優(yōu)過程中，可以先嘗試開啟DP，對輸入數(shù)據(jù)進行切片，提高數(shù)據(jù)處理速率?；A版本以DP1-TP4-PP1-SP-CP8的配置實現(xiàn)，基礎版本不開啟DP，由于需要GPU總數(shù)等于DP*TP*PP*CP，所以需要相應減少CP數(shù)量。開啟DP的同時，使用--overlap-grad-reduce --overlap-param-gather兩個參數(shù)進行DP通信時的重疊，可以提升12%的性能。如果要開啟MP，則優(yōu)先開啟TP，因為TP是機器節(jié)點內部通信，可以通過NVLINK加速，相比PP通信效率更高。另外由于模型特性，TP最大值固定為4，因此無需調整。

由于任務上下文較長，開啟SP和CP是較好的選擇，但CP值也不是越大越好，最佳值需要經過仔細測試和綜合考慮后才能確定，最終調優(yōu)結果如表3.1所示：

表 3.1 不同優(yōu)化參數(shù)下的結果

v2 版本在優(yōu)化過程中遭遇了 OOM，這促使我們在 v3 版本中采取了激進的顯存碎片清理策略，力爭在每個計算周期結束后釋放不再需要的顯存空間。雖然這一舉措對訓練速度有一定影響，但相較于 v1 版本的速度和 v2 版本的 OOM 問題kaiyun官方網app下載app，v3 版本最終以 35.5s/sample 的速度實現(xiàn)了穩(wěn)定運行，成功平衡了性能與資源利用率的關系。

值得注意的是，v3 及之前版本均采用了全量重新計算的方式，這將帶來約 30% 的性能損失。鑒于此，v4 版本為了提升效率嘗試降低重新計算的層級，但遺憾的是，此次調整再次觸發(fā)了 OOM 異常。此時尚未開啟 PP 策略。v5 開啟 PP 策略，雖然在速度上超越了基礎版本，但與 V3 版本的峰值性能相比仍有不小差距。因此如圖 3.2 所示，本次調優(yōu)達到的最優(yōu)效率為 35.5s/sample，相比基礎版本提升了 3.4 倍，模型推理結果達到了預期效果。

圖3.2 調優(yōu)前后速率對比示意圖

四、結論

本文講解了大型模型分布式訓練過程中并行優(yōu)化技術的基本原理，分析了模型訓練過程中的顯存占用情況，以及如何更好的節(jié)省顯存，并結合實際應用場景展示了 QLM 如何利用這些技術來提升長文本任務的訓練效率。

隨著大模型的快速演進，未來的訓練場景必然會遭遇“三重壁壘”——內存墻、通信墻、計算墻，這三大瓶頸將帶來嚴峻的挑戰(zhàn)。面對這種情況，如何精簡顯存使用量，降低跨節(jié)點通信成本，提升計算性能和效率，成為了QLM團隊不懈的追求，我們正在全力以赴探索創(chuàng)新解決方案，構建更高效、可持續(xù)的大模型生態(tài)。

參考：

[1] MegatronLM序列模型并行訓練（Sequence Parallel）詳解，

[2]Narayanan D, Shoeybi M, Casper J 等人。使用 megatron-lm 在 GPU 集群上進行高效的大規(guī)模語言模型訓練[C]//國際高性能計算、網絡、存儲和分析會議論文集。2021：1-15。

[3]Korthikanti VA, Casper J, Lym S 等人。減少大型 Transformer 模型中的激活重新計算[J]。機器學習與系統(tǒng)學報，2023 年，5：341-353。

[4]上下文并行概述，#context-parallelism-overview

[5] Rajbhandari S, Rasley J, Ruwase O 等人。零：面向訓練萬億參數(shù)模型的內存優(yōu)化[C]//SC20：國際高性能計算、網絡、存儲和分析會議。IEEE，2020：1-16。