探索在大模型訓(xùn)練中使用 Megatron-Core 訓(xùn)練框架提高顯存使用效率

隨著模型規(guī)模邁入百億、千億甚至萬億參數(shù)級別，如何在有限顯存中“塞下”訓(xùn)練任務(wù)，對研發(fā)和運(yùn)維團(tuán)隊(duì)都是巨大挑戰(zhàn)。NVIDIA Megatron-Core 作為流行的大模型訓(xùn)練框架，提供了靈活高效的并行化策略；理解這些策略對顯存的影響，才能更好地規(guī)劃訓(xùn)練超參數(shù)，在不 OOM (out of memory) 的情況下盡可能提升硬件使用效率。

• 顯存的組成與衡量方法：通過 torch 的顯存可視化工具捕捉一個典型的模型訓(xùn)練中的顯存占用。靜態(tài)顯存主要組成部分包括模型參數(shù)、梯度和優(yōu)化器的所占用的空間，及一些其他的系統(tǒng)開銷。設(shè)定 R 為參數(shù)重復(fù)次數(shù)，則對 bf16 訓(xùn)練來說每個參數(shù)占用的字節(jié)數(shù)為 6+12/R。對于 Mixture of Experts (MoE) 模型來說，由于 Megatron 支持 parallel folding，模型的模型會分為稠密部分和 MoE 部分，其中稠密部分的 R 為 DP*CP，MoE 部分的 R 為 EDP=n_GPU/PP/EP/ETP。

動態(tài)顯存則是模型前向傳播過程中暫存的中間結(jié)果，用于反向傳播時計(jì)算梯度，通常被稱為激活 (Activation)，絕大部分為 bf16 數(shù)據(jù)類型。

• 對顯存影響的關(guān)鍵超參數(shù)：Megatron-Core 支持以下并行、重算維度，組合后可覆蓋當(dāng)下主流大模型訓(xùn)練需求。

約束關(guān)系：`n_GPU / PP = TP×CP×DP = EP×ETP×EDP`，其中 `EDP` 為專家數(shù)據(jù)并行度。

除了完全不重算的情況之外，為了降低動態(tài)顯存，Megatron-Core 0.14 提供完全重算 (full) 和細(xì)粒度重算 (selective) 這兩檔重算。

• 顯存估計(jì)器的設(shè)計(jì)：當(dāng)前 Megatron 基于 torch 實(shí)現(xiàn)，所有模塊均派生自 torch.nn.Module，構(gòu)成訓(xùn)練 GPT 類模型的模塊。我們通過實(shí)現(xiàn)一個基類 MemEstimator 并基于此基類派生出所有需要的模塊類，根據(jù)每個模塊的顯存占用特點(diǎn)分別計(jì)算其中的參數(shù)量和激活量。然后復(fù)用 Megatron 中本身構(gòu)建模型的代碼，實(shí)現(xiàn)一個 Megatron 模擬器，并可以展示出個層次的模塊數(shù)據(jù)量。

• 關(guān)鍵結(jié)論：選取 Qwen3 235B 和 DeepSeek v3 兩個時下流行的大模型，使用流行的配置開啟訓(xùn)練，并對比顯存估計(jì)的結(jié)果與真實(shí)的顯存占用。兩個模型的實(shí)際峰值與估計(jì)峰值相差均小于 2GB。

通過對動態(tài)顯存分析，只有 TP 和 CP 能降低激活量，EP 和 ETP 只會改變集群內(nèi)激活值的分布，無法降低激活量，PP 和 VPP 由于 1f1b 的流水線預(yù)熱機(jī)制，無法有效降低峰值激活量。對每一部分激活量，可以通過卸載到 CPU 或者重算的方式來降低顯存。Megatron-Core 0.13 當(dāng)前對卸載的支持還在開發(fā)中，但重算已經(jīng)支持。

Megatron-Core 0.13 現(xiàn)已支持通過 CPU 分擔(dān) optimizer 的顯存占用，并可以通過超參數(shù)設(shè)置卸載到 CPU 的比例，每個參數(shù)的 6 字節(jié) (bf16 參數(shù)，fp32 梯度) 無法卸載，其余可以卸載。

• 用例分析：用戶目標(biāo)在 32 張 80GB 顯存的 GPU 上實(shí)現(xiàn) Qwen3-30B-A3B 的強(qiáng)化學(xué)習(xí)訓(xùn)練，序列長度是 10k，用戶使用顯存估計(jì)器對并行配置進(jìn)行摸底。

Megatron 開發(fā)者可以通過顯存分析工具的 breakdown 視角，詳細(xì)察看每個模塊的激活量，通過權(quán)衡激活量和計(jì)算量尋找性價比高（激活量 / 計(jì)算量）的模塊的激活為其開發(fā)進(jìn)行重算或卸載功能。

以上為摘要內(nèi)容，點(diǎn)擊鏈接閱讀完整內(nèi)容：

探索在大模型訓(xùn)練中使用 Megatron-Core 訓(xùn)練框架提高顯存使用效率 - NVIDIA 技術(shù)博客