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

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

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

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

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

除了完全不重算的情況之外，為了降低動(dòng)態(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 類模型的模塊。我們通過(guò)實(shí)現(xiàn)一個(gè)基類 MemEstimator 并基于此基類派生出所有需要的模塊類，根據(jù)每個(gè)模塊的顯存占用特點(diǎn)分別計(jì)算其中的參數(shù)量和激活量。然后復(fù)用 Megatron 中本身構(gòu)建模型的代碼，實(shí)現(xiàn)一個(gè) Megatron 模擬器，并可以展示出個(gè)層次的模塊數(shù)據(jù)量。

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

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

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

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

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

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

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