最近，Hugging Face發布了一個「超大規模訓練手冊」，教我們如何在GPU集群上訓練LLM。

這項震撼的研究，在512個GPU上做了超過4000個Scaling實驗，并測量了吞吐量（標記的大小）和GPU利用率（標記的顏色）。

圖片

HuggingFace聯創兼CEO Clement表示，自己的理想，就是能生活在這樣一個世界：所有的公司和組織無論大小，無論是否富有，都能訓練自己的AI。

未來的世界，就應該是這個樣子。

因此，當發布這份大規模訓練手冊時，他感到十分自豪。

圖片

網友們紛紛表示，這就是自己愛Hugging Face的原因。

圖片

未來我們應該擁有的就是民主化的AI和去中心化數據，讓全球大腦共同工作。從此，AI不再是為1%的人服務，而是讓99%的人參與其中！

圖片

本文將從基礎入手，介紹如何將LLM訓練規模從一塊GPU擴展到數十塊、數百塊甚至數千塊GPU。

隨著訓練集群規模不斷擴大，數據并行、張量并行、流水線并行、上下文并行，以及ZeRO和內核融合等技術相繼被提出，來確保GPU的高效利用。

在深入代碼和實驗之前，先要理解每種方法的工作原理、優缺點和適用場景。例如，LLM在訓練過程中的哪些部分占用了最多的顯存，訓練過程中的哪個階段會成為瓶頸。

同時，還會介紹如何通過并行來解決顯存限制，提高吞吐量。

這樣一來，就能理解下面這個用于計算Transformer模型顯存占用的小工具是如何工作的。

圖片

除了理論分析，還提供了一個工具，用于預測訓練過程中顯存的實際使用情況：

圖片

本文運行了4100多次分布式實驗，用了512塊GPU，以探索可能的分布式訓練架構和模型大小的影響。

圖片

概覽

本文內容廣泛，作者將內容總結如下：

圖片

訓練過程有三個關鍵挑戰：

1. 顯存占用：如果某個訓練步驟超出了顯存容量，訓練就無法繼續。
2. 計算效率：希望GPU大部分時間都在執行計算，而不是浪費在數據傳輸或等待其他GPU執行任務。
3. 通信開銷：減少通信開銷，以免GPU處于空閑狀態。需充分利用節點內部和節點之間的帶寬，盡量讓通信和計算過程重疊進行，以提高訓練效率。

在很多情況下，可以在計算、通信和顯存中進行取舍，如通過重計算或張量并行，找到合適的平衡點。

在單個GPU上訓練模型時，通常包含三個步驟：前向傳播、反向傳播和優化步驟。

在預訓練中，批大小通常以token為單位。這使得訓練的計算量通常與輸入序列長度無關。

近期LLM的批大小通常在4M到60M個token。Llama 1的訓練批大小約為4M個token，DeepSeek的訓練批大小約為60M個token。

第一個挑戰已經出現：「顯存不足」。

訓練時，顯存需存儲以下內容：模型權重、模型梯度、優化器狀態和計算梯度所需的激活值。

如何根據這些變量，快速確定顯存使用情況呢？一個簡單的方法是通過實驗測量。

分析顯存使用情況

用PyTorch分析器，可以了解訓練過程中顯存的分配方式。顯存利用率在訓練過程中，會有很大的變化。

圖片

接下來，探討如何在擴展訓練規模的過程中，最大化計算效率，同時確保激活值、參數、梯度和優化器狀態的顯存需求在限制范圍內。

對于一個簡單的Transformer LLM，參數量可以用以下公式計算：

圖片

顯存需求可通過總參數乘以每個參數占用的字節數來計算。出于穩定性及效率的考慮，通常采用混合精度。

接下來，估算模型的顯存需求：

圖片

一旦模型參數達到7B，權重和優化器狀態的顯存需求就開始大幅增加，并超過典型GPU顯存容量，例如H100的80G。

相比于權重、梯度和優化器狀態，激活值的計算更加復雜，部分原因是它取決于模型的輸入。

現在，介紹第一種技術——激活值重計算。

激活值重計算

激活值重計算的基本思想是在前向傳播中丟棄一些激活值，以節省顯存。并通過增加一些計算量，在反向傳播中動態計算這些激活值。

使用重計算時，通常只在模型架構的幾個關鍵點存儲激活值，丟棄其余的激活值，并在反向傳播中從最近保存的激活值開始重新計算它們。

圖片

選擇要存儲的關鍵激活值有全量和選擇性等策略。接下來看到，重計算如何減少顯存占用，以及如何在節省顯存和增加計算成本之間取得良好的平衡。

圖片

對于規模較小的模型，長序列的激活值產生的影響更大，因此重計算的效果更顯著。

如今，大多數訓練框架都使用FlashAttention，原生集成了激活值重計算的優化策略。

激活值重計算會稍微增加浮點運算次數，但它減少了內存訪問開銷。這種權衡在GPU上特別有利，計算速度通常更快，同時降低了顯存占用。

然而，激活值仍然與批大小呈線性相關，當批大小增加時，激活值的顯存可能又會成為問題。

還有第二個方法，梯度累積來救場！

梯度累積

梯度累積是一種避免顯存爆炸的方法，原理是將批量數據拆分為多個微批次，依次進行前向傳播和反向傳播。

通過梯度累積，全局批大小可通過以下公式計算：

圖片

梯度累積還可以與激活重計算相結合，進一步減少顯存占用。

圖片

梯度累積能通過僅計算部分微批次，來減少激活值占用的顯存。每個微批次的前向和反向傳播可以并行運行，看來是時候考慮多個GPU了！

在擴展到多個GPU前，介紹分布式訓練工具箱中最有用的工具之一：分析器。

PyTorch分析器

分析器能精確追蹤和可視化訓練過程中的情況，展示了：

• CPU線程異步啟動內核到GPU。
• 多個CUDA流并行處理計算和通信任務。
• 內核執行時間和內存分配。

首先介紹數據并行技術，它是梯度累積的并行版本。

數據并行

數據并行的核心思想是在多個GPU上運行，并在每個GPU上并行處理不同微批次的數據。

圖片

每個GPU上的梯度是不同的，為了讓不同GPU上的模型保持同步，用all-reduce操作對模型的梯度進行平均。

圖片

由于不希望GPU處于空閑狀態，應盡可能地讓通信和計算同時進行。這里有三種優化方法：將梯度同步與后向傳播重疊進行、梯度分桶和與梯度累積相結合。

重新審視全局批大小

結合新引入的數據并行和梯度累積參數來更新批大小：

圖片

給定一個目標全局批大小，可以通過調整梯度累積步數和并行進程數來加快訓練速度。

當GPU數量超過限制時，吞吐量開始顯著下降。

圖片

當數據并行達到一定規模后，會出現明顯的通信開銷瓶頸。對于更大的模型或者批大小，還有其他選擇嗎？

幸運的是，確實有解決方案。這些方案將一些張量移到CPU上，或將權重、梯度、優化器等張量拆分到多個GPU上。

拆分主要有兩種方法：并行化（張量并行、上下文并向或流水線并行）和共享（如DeepSpeed Zero或PyTorch FSDP）。兩種方法相互獨立，也可以結合使用！

共享模式與數據并行密切相關，首先來研究ZeRO方法。

ZeRO（零冗余優化器）

DeepSpeed ZeRO是一種旨在減少LLM訓練中內存冗余的優化技術。

數據并行是一種高效的方法，但在每個實例上簡單復制優化器狀態、梯度和參數會引入大量的內存冗余。

ZeRO通過在數據并行維度上對優化器狀態、梯度和參數進行分區，消除了內存冗余。

ZeRO有三個優化階段：

• ZeRO-1：優化器狀態分區
• ZeRO-2：優化器狀態+梯度分區
• ZeRO-3（也稱為FSDP）：優化器狀態+梯度+參數分區

ZeRO的理念是在數據并行進程之間，對這些對象進行分區，每個節點僅存儲一部分，需要時進行重建。

圖片

ZeRO-1：優化器狀態分區

在ZeRO-1中，優化器狀態被分成N_d等份，N_d是數據并行度。在優化步驟中，只有1/N_d的權重會被更新。

對梯度執行reduce-scatter操作。

圖片

與傳統的數據并行相比，ZeRO-1將all-reduce梯度通信替換為reduce-scatter操作，并在優化器后添加了一個全參數的all-gather操作。

圖片

ZeRO-2：增加梯度分區

在反向傳播中，不再對梯度執行all-reduce操作，而是僅執行reduce-scatter操作。只傳播1/N_d的梯度，與ZeRO-1相比，節省了更多內存。

梯度分區后，內存有所減少。隨著N_d的增加，最多可節省8倍內存。在通信方面，與ZeRO-1相比，唯一的區別在于可以動態釋放內存。

圖片

ZeRO-3：增加參數分區

在第三階段，將之前對優化器狀態和梯度進行分區的做法，擴展到模型參數。

如果模型的所有部分都被分布式存儲，在前向傳播中，具體操作如下：

圖片

在進行前向傳播時，遍歷各層，獲取所需的參數，并在不再需要時將其從內存中清除。反向傳播與之類似，只是流程相反：

圖片

由于需要在前向傳播和反向傳播中持續進行all-gather操作，與ZeRO-2相比，每個訓練步驟會增加（2×層數-1）次額外的操作。

圖片

借助ZeRO技術，可以在數據并行中，將參數、梯度和優化器狀態進行分區。

然而，這里有一個限制，ZeRO無法對激活值內存進行處理。這部分內存會隨著序列長度和批大小的增加而增加。

圖片

為克服這些問題，是時候探索一種新的并行方式了——張量并行。與嚴重依賴參數通信的ZeRO方法不同，張量并行提出將參數、梯度、優化器狀態和激活值分布到多個設備上，而無需在各GPU之間進行模型參數的通信。

張量并行

當激活值內存占用超過預算時，ZeRO就會遇到瓶頸。張量并行是在ZeRO基礎上，針對激活內存超預算問題的優化技術。

利用矩陣乘法特性，通過按列或按行分區，將張量分布到多個GPU上計算。列線性需廣播輸入矩陣、分割權重矩陣列，用all-reduce組合結果；行線性要分割權重矩陣行和輸入，分別使用scatter和all-reduce操作。

圖片

圖片

張量并行能減少矩陣乘法激活內存，在多GPU間分布模型參數、梯度、優化器狀態，使7B參數模型可在單節點8個GPU上運行。

缺點是跨節點通信慢，當張量并行度超過8個GPU時，通信開銷明顯，從TP=8到TP=16、TP=16到TP=32性能顯著下降。層歸一化和隨機失活等操作仍需收集完整激活值。

圖片

序列并行

為解決層歸一化和隨機失活需完整激活值的問題，引入序列并行技術。

圖片

序列并行的優勢是減少最大激活值存儲大小，僅使用張量并行時需存儲形狀為 (b,s,h) 的激活值，使用序列并行后可減少到

圖片

這有助于節省激活值內存，能增大批大小和序列長度，如在70B參數模型中，使用TP/SP=16時可處理16k token的序列長度，優于單純使用張量并行的情況。

圖片

隨著張量并行度增加，計算效率和顯存容量需要權衡。從TP=8提升到TP=16時性能下降明顯，因為涉及節點內到節點間的通信轉變。

圖片

更高的并行度雖能減少激活內存以處理更大批次，但會降低單GPU吞吐量，超過單節點GPU數量對應的閾值時更為明顯。

序列長度增加時，TP區域的激活值內存仍會激增；模型過大時，TP=8也無法適配，會因節點間連接導致速度大幅下降。可分別用上下文并行和流水線并行解決這些問題。

上下文并行

圖片

借鑒序列并行按序列長度拆分的思路，對已應用張量并行的模塊沿序列長度和另一個維度進行拆分，在整個模型上應用序列拆分，而非僅在模型的序列并行區域。

這種方式對多數模塊無影響，因為這些模塊中每個token獨立處理，且無需像張量并行那樣進行高成本的通信，僅分割輸入，不拆分權重矩陣。計算梯度后，通過all-reduce操作同步上下文并行組內的梯度。

注意力模塊中每個token需訪問其他所有token的鍵/值對。由于上下文并行按序列維度拆分輸入，注意力模塊需在GPU間進行全面通信以交換鍵/值數據。

為高效處理這種通信，引入了環形注意力（Ring Attention）技術。

圖片

以4個GPU和4個token的輸入為例，每個GPU先異步將自身的鍵/值對發送給其他GPU，在等待時計算已有數據的注意力分數。理想狀態下，計算完成前能收到下一個鍵/值對，可立即開始下一輪計算。

每個GPU在每個時間步執行三個操作：非阻塞式發送當前鍵和值（最后一步除外）、本地計算注意力分數、等待接收前一個GPU的鍵/值后循環執行。

圖片

環形注意力的簡單實現會因因果注意力矩陣的形狀導致GPU計算負載不均衡。

SoftMax按行計算，GPU1因起始就有完整行的token數據可立即計算，且無需接收其他GPU信息；GPU2則需等待第二輪接收更多數據才能計算，GPU1計算量明顯少于其他GPU。

Zig-Zag環形注意力機制

當前需更好的方式分配輸入序列，Zig-Zag注意力摒棄順序分配token至GPU，而是采用混合排序，使每個GPU上都有早期和晚期token，實現計算在各GPU上的平衡分布。

圖片

由于每個GPU完成計算需要其他GPU的信息，存在兩種常見方式來重疊計算和通信：AllGather和All-to-All（環形）實現。

張量并行可在單個節點上拆分模型處理大型模型，上下文并行則用于解決長序列導致的激活量激增問題。但張量并行在跨節點擴展時效果不佳。

那么，如果模型權重無法輕松地在一個節點上存儲，該怎么辦呢？

這時，流水線并行（Pipeline Parallelism）就派上用場了！

流水線并行

張量并行擴展到超過單個節點的GPU數量（一般4或8個）時，會受到節點間連接低帶寬網絡影響，性能下降明顯。

對于70B參數以上的模型，單節點4-8個GPU難以承載其權重規模，因此需要流水線并行技術。

將模型的各層分布到多個GPU上，如8個GPU時，可把第1-4層放于GPU1，第5-8層放于GPU2等。這樣每個GPU僅需存儲和處理部分模型層，減少了單個GPU的內存需求。

但由于每個GPU仍需處理完整批次數據，激活內存不會因層的劃分而減少，且激活張量需在GPU間按流水線順序傳遞。流水線并行中的數據處理具有順序性，GPU利用率不高。

全前向全反向（AFAB）調度

由于計算是順序進行的，存在效率不高的問題。GPU存在空閑時間，會降低利用率。通過公式推導可知，空閑時間與流水線并行度相關，并行度越高，空閑時間越長，利用率越低。

圖片

將批次數據拆分成更小的微批次進行并行處理。AFAB調度先進行所有前向傳播，再進行所有反向傳播，保留了模型訓練代碼的總體結構，易于實現。計算表明，增加微批次數量可減小空閑時間占比，提高效率。

圖片

由于反向傳播前需保存所有激活值，這種方法會導致內存迅速耗盡。

因此需要探索新方法，如在進行前向計算時就開始反向傳播，以盡早丟棄反向傳播所需的激活值，避免顯存爆炸。

One-forward-one-backward調度

交替執行一次前向和一次反向傳播，盡早開始反向傳播。該方式雖未顯著提升訓練效率，但能減少激活內存占用，且可增加微批次以減小空閑時間。

圖片

微批次數量等于或小于流水線并行度-1時，性能低且隨并行度增加而下降。使用更多微批次有助于提升低并行度性能，但高并行度時仍受限。

受全局批大小限制，不能隨意增加微批次，并行度增加時空閑時間會相應增加。

圖片

微批次數量較少時，跨節點擴展性能下降僅14%，優于張量并行，在跨節點分布式訓練中有吸引力。

交錯階段技術

不同于簡單按模型深度劃分，交錯階段如將奇數層和偶數層分別置于不同GPU，形成「循環流水線」。微批次前向傳播時在GPU間循環。

圖片

雖增加了通信量，但每次前向和反向傳播時間因v（每個GPU的階段數或模型塊數）而減少，可通過增加微批次和交錯階段減小空閑時間。

圖片

交錯階段使調度更復雜，需在特定時刻決定優先策略，即「深度優先」（盡快通過模型）或「廣度優先」（盡可能填滿流水線）。

圖片

Llama 3.1的流水線并行方法采用了帶交錯階段的單前向單反向設置，深度優先和廣度優先的優先級設置可調節。

零氣泡和雙管道技術

為減少空閑時間提出了更復雜方法，關鍵是細粒度拆分操作并交錯執行。如DeepSeek V3/R1的DualPipe。

ZeroBubble發現矩陣乘法反向傳遞中，輸入反向操作（B）和權重反向操作（W）可分離，W可在對應B之后靈活安排，用于填補流水線空閑時間。

圖片

DeepSeek的DualPipe在V3技術報告中擴展了分解方法，針對從PP維度兩端傳播的兩個數據流交錯，以進一步減少 GPU空閑時間。

圖片

專家并行

MoE模型近年來因GPT-4、Mixtral、DeepSeek-V3/R1等模型受到關注。其基本思想是每一層不采用單個前饋模塊，而是設置多個并行模塊，對token進行不同處理。

MoE層的設計使專家并行易于實現，因前饋層完全獨立。與張量并行（TP）相比，專家并行更輕量，無需拆分矩陣乘法，只需將token隱藏狀態路由到合適專家。

實際中，專家并行（EP）常與其他并行方式結合使用。因EP僅影響MoE層，不分片輸入token，若僅用EP，GPU處理非MoE模塊時會有冗余計算。EP高效運行的技巧與模型設計緊密相關。

更多詳細內容請查看原文！

參考資料：https://huggingface.co/spaces/nanotron/ultrascale-playbook