歷時6個月，Hugging Face開源LLM「超大規(guī)模實(shí)戰(zhàn)手冊」！200頁3萬字4000次訓(xùn)練

作者：新智元 2025-03-03 13:19:21

Hugging Face發(fā)布了「超大規(guī)模實(shí)戰(zhàn)手冊」，在512個GPU上進(jìn)行超過4000個scaling實(shí)驗(yàn)。聯(lián)創(chuàng)兼CEO Clement對此感到十分自豪。

最近，Hugging Face發(fā)布了一個「超大規(guī)模訓(xùn)練手冊」，教我們?nèi)绾卧贕PU集群上訓(xùn)練LLM。

這項(xiàng)震撼的研究，在512個GPU上做了超過4000個Scaling實(shí)驗(yàn)，并測量了吞吐量（標(biāo)記的大小）和GPU利用率（標(biāo)記的顏色）。

圖片

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

未來的世界，就應(yīng)該是這個樣子。

因此，當(dāng)發(fā)布這份大規(guī)模訓(xùn)練手冊時，他感到十分自豪。

圖片

網(wǎng)友們紛紛表示，這就是自己愛Hugging Face的原因。

圖片

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

圖片

本文將從基礎(chǔ)入手，介紹如何將LLM訓(xùn)練規(guī)模從一塊GPU擴(kuò)展到數(shù)十塊、數(shù)百塊甚至數(shù)千塊GPU。

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

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

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

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

圖片

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

圖片

本文運(yùn)行了4100多次分布式實(shí)驗(yàn)，用了512塊GPU，以探索可能的分布式訓(xùn)練架構(gòu)和模型大小的影響。

圖片

概覽

本文內(nèi)容廣泛，作者將內(nèi)容總結(jié)如下：

圖片

訓(xùn)練過程有三個關(guān)鍵挑戰(zhàn)：

顯存占用：如果某個訓(xùn)練步驟超出了顯存容量，訓(xùn)練就無法繼續(xù)。
計算效率：希望GPU大部分時間都在執(zhí)行計算，而不是浪費(fèi)在數(shù)據(jù)傳輸或等待其他GPU執(zhí)行任務(wù)。
通信開銷：減少通信開銷，以免GPU處于空閑狀態(tài)。需充分利用節(jié)點(diǎn)內(nèi)部和節(jié)點(diǎn)之間的帶寬，盡量讓通信和計算過程重疊進(jìn)行，以提高訓(xùn)練效率。

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

在單個GPU上訓(xùn)練模型時，通常包含三個步驟：前向傳播、反向傳播和優(yōu)化步驟。

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

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

第一個挑戰(zhàn)已經(jīng)出現(xiàn)：「顯存不足」。

訓(xùn)練時，顯存需存儲以下內(nèi)容：模型權(quán)重、模型梯度、優(yōu)化器狀態(tài)和計算梯度所需的激活值。

如何根據(jù)這些變量，快速確定顯存使用情況呢？一個簡單的方法是通過實(shí)驗(yàn)測量。

分析顯存使用情況

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

圖片

接下來，探討如何在擴(kuò)展訓(xùn)練規(guī)模的過程中，最大化計算效率，同時確保激活值、參數(shù)、梯度和優(yōu)化器狀態(tài)的顯存需求在限制范圍內(nèi)。

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

圖片

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

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

圖片

一旦模型參數(shù)達(dá)到7B，權(quán)重和優(yōu)化器狀態(tài)的顯存需求就開始大幅增加，并超過典型GPU顯存容量，例如H100的80G。

相比于權(quán)重、梯度和優(yōu)化器狀態(tài)，激活值的計算更加復(fù)雜，部分原因是它取決于模型的輸入。

現(xiàn)在，介紹第一種技術(shù)——激活值重計算。

激活值重計算

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

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

圖片

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

圖片

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

如今，大多數(shù)訓(xùn)練框架都使用FlashAttention，原生集成了激活值重計算的優(yōu)化策略。

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

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

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

梯度累積

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

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

圖片

梯度累積還可以與激活重計算相結(jié)合，進(jìn)一步減少顯存占用。

圖片

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

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

PyTorch分析器

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

CPU線程異步啟動內(nèi)核到GPU。
多個CUDA流并行處理計算和通信任務(wù)。
內(nèi)核執(zhí)行時間和內(nèi)存分配。

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

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

數(shù)據(jù)并行的核心思想是在多個GPU上運(yùn)行，并在每個GPU上并行處理不同微批次的數(shù)據(jù)。

圖片

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

圖片

由于不希望GPU處于空閑狀態(tài)，應(yīng)盡可能地讓通信和計算同時進(jìn)行。這里有三種優(yōu)化方法：將梯度同步與后向傳播重疊進(jìn)行、梯度分桶和與梯度累積相結(jié)合。

重新審視全局批大小

結(jié)合新引入的數(shù)據(jù)并行和梯度累積參數(shù)來更新批大小：

圖片

給定一個目標(biāo)全局批大小，可以通過調(diào)整梯度累積步數(shù)和并行進(jìn)程數(shù)來加快訓(xùn)練速度。

當(dāng)GPU數(shù)量超過限制時，吞吐量開始顯著下降。

圖片

當(dāng)數(shù)據(jù)并行達(dá)到一定規(guī)模后，會出現(xiàn)明顯的通信開銷瓶頸。對于更大的模型或者批大小，還有其他選擇嗎？

幸運(yùn)的是，確實(shí)有解決方案。這些方案將一些張量移到CPU上，或?qū)?quán)重、梯度、優(yōu)化器等張量拆分到多個GPU上。

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

共享模式與數(shù)據(jù)并行密切相關(guān)，首先來研究ZeRO方法。

ZeRO（零冗余優(yōu)化器）

DeepSpeed ZeRO是一種旨在減少LLM訓(xùn)練中內(nèi)存冗余的優(yōu)化技術(shù)。

數(shù)據(jù)并行是一種高效的方法，但在每個實(shí)例上簡單復(fù)制優(yōu)化器狀態(tài)、梯度和參數(shù)會引入大量的內(nèi)存冗余。

ZeRO通過在數(shù)據(jù)并行維度上對優(yōu)化器狀態(tài)、梯度和參數(shù)進(jìn)行分區(qū)，消除了內(nèi)存冗余。

ZeRO有三個優(yōu)化階段：

ZeRO-1：優(yōu)化器狀態(tài)分區(qū)
ZeRO-2：優(yōu)化器狀態(tài)+梯度分區(qū)
ZeRO-3（也稱為FSDP）：優(yōu)化器狀態(tài)+梯度+參數(shù)分區(qū)

ZeRO的理念是在數(shù)據(jù)并行進(jìn)程之間，對這些對象進(jìn)行分區(qū)，每個節(jié)點(diǎn)僅存儲一部分，需要時進(jìn)行重建。

圖片

ZeRO-1：優(yōu)化器狀態(tài)分區(qū)

在ZeRO-1中，優(yōu)化器狀態(tài)被分成N_d等份，N_d是數(shù)據(jù)并行度。在優(yōu)化步驟中，只有1/N_d的權(quán)重會被更新。

對梯度執(zhí)行reduce-scatter操作。

圖片

與傳統(tǒng)的數(shù)據(jù)并行相比，ZeRO-1將all-reduce梯度通信替換為reduce-scatter操作，并在優(yōu)化器后添加了一個全參數(shù)的all-gather操作。

圖片

ZeRO-2：增加梯度分區(qū)

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

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

圖片

ZeRO-3：增加參數(shù)分區(qū)

在第三階段，將之前對優(yōu)化器狀態(tài)和梯度進(jìn)行分區(qū)的做法，擴(kuò)展到模型參數(shù)。

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

圖片

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

圖片

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

圖片

借助ZeRO技術(shù)，可以在數(shù)據(jù)并行中，將參數(shù)、梯度和優(yōu)化器狀態(tài)進(jìn)行分區(qū)。

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

圖片

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

張量并行

當(dāng)激活值內(nèi)存占用超過預(yù)算時，ZeRO就會遇到瓶頸。張量并行是在ZeRO基礎(chǔ)上，針對激活內(nèi)存超預(yù)算問題的優(yōu)化技術(shù)。

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

圖片

張量并行能減少矩陣乘法激活內(nèi)存，在多GPU間分布模型參數(shù)、梯度、優(yōu)化器狀態(tài)，使7B參數(shù)模型可在單節(jié)點(diǎn)8個GPU上運(yùn)行。

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

圖片

序列并行

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

圖片

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

圖片

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

圖片

隨著張量并行度增加，計算效率和顯存容量需要權(quán)衡。從TP=8提升到TP=16時性能下降明顯，因?yàn)樯婕肮?jié)點(diǎn)內(nèi)到節(jié)點(diǎn)間的通信轉(zhuǎn)變。

圖片

更高的并行度雖能減少激活內(nèi)存以處理更大批次，但會降低單GPU吞吐量，超過單節(jié)點(diǎn)GPU數(shù)量對應(yīng)的閾值時更為明顯。

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

上下文并行

圖片

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

這種方式對多數(shù)模塊無影響，因?yàn)檫@些模塊中每個token獨(dú)立處理，且無需像張量并行那樣進(jìn)行高成本的通信，僅分割輸入，不拆分權(quán)重矩陣。計算梯度后，通過all-reduce操作同步上下文并行組內(nèi)的梯度。

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

為高效處理這種通信，引入了環(huán)形注意力（Ring Attention）技術(shù)。

圖片

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

每個GPU在每個時間步執(zhí)行三個操作：非阻塞式發(fā)送當(dāng)前鍵和值（最后一步除外）、本地計算注意力分?jǐn)?shù)、等待接收前一個GPU的鍵/值后循環(huán)執(zhí)行。

圖片

環(huán)形注意力的簡單實(shí)現(xiàn)會因因果注意力矩陣的形狀導(dǎo)致GPU計算負(fù)載不均衡。

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

Zig-Zag環(huán)形注意力機(jī)制

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

圖片

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

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

那么，如果模型權(quán)重?zé)o法輕松地在一個節(jié)點(diǎn)上存儲，該怎么辦呢？

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

流水線并行

張量并行擴(kuò)展到超過單個節(jié)點(diǎn)的GPU數(shù)量（一般4或8個）時，會受到節(jié)點(diǎn)間連接低帶寬網(wǎng)絡(luò)影響，性能下降明顯。

對于70B參數(shù)以上的模型，單節(jié)點(diǎn)4-8個GPU難以承載其權(quán)重規(guī)模，因此需要流水線并行技術(shù)。

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

但由于每個GPU仍需處理完整批次數(shù)據(jù)，激活內(nèi)存不會因?qū)拥膭澐侄鴾p少，且激活張量需在GPU間按流水線順序傳遞。流水線并行中的數(shù)據(jù)處理具有順序性，GPU利用率不高。

全前向全反向（AFAB）調(diào)度

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

圖片

將批次數(shù)據(jù)拆分成更小的微批次進(jìn)行并行處理。AFAB調(diào)度先進(jìn)行所有前向傳播，再進(jìn)行所有反向傳播，保留了模型訓(xùn)練代碼的總體結(jié)構(gòu)，易于實(shí)現(xiàn)。計算表明，增加微批次數(shù)量可減小空閑時間占比，提高效率。

圖片