新智元報道
新智元報道
編輯：Aeneas 英智
【新智元導讀】Hugging Face釋出了「超大規模實戰手冊」，在512個GPU上進行超過4000個scaling實驗。聯創兼CEO Clement對此感到十分自豪。
最近，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，以探索可能的分散式訓練架構和模型大小的影響。
概覽
本文內容廣泛，作者將內容總結如下：
訓練過程有三個關鍵挑戰：
-
視訊記憶體佔用：如果某個訓練步驟超出了視訊記憶體容量，訓練就無法繼續。
-
計算效率：希望GPU大部分時間都在執行計算，而不是浪費在資料傳輸或等待其他GPU執行任務。
-
通訊開銷：減少通訊開銷，以免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