新智元報道
新智元報道
編輯: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


