NVIDIA NeMo AutoModel 加速 Mixture-of-Experts 模型的微調

NVIDIA NeMo AutoModel 加速 Mixture-of-Experts 模型的微調

背景

Mixture-of-Experts 模型的興起帶來了訓練挑戰，Transformers v5 透過專家後端、動態權重載入與分散式執行來解決這些問題，但仍缺乏融合的通信核心。

NeMo AutoModel：相同 API，更佳效能

NeMo AutoModel 繼承自 AutoModelForCausalLM，使用者只需更改一行 import，即可從 Hugging Face Transformers 切換至 NeMo AutoModel，並在不修改訓練程式碼的情況下獲得 Expert Parallelism、DeepEP 融合的 all-to-all 分派以及 TransformerEngine 核心。

使用 NeMo AutoModel 載入模型的方式與 Hugging Face 版本相同，唯一差別在於 import：

from nemo_automodel import NeMoAutoModelForCausalLM
model = NeMoAutoModelForCausalLM.from_pretrained(
    "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16",
    dtype=torch.bfloat16,
)

對於 Qwen3、NVIDIA Nemotron、GPT-OSS 與 DeepSeek V3 等流行的 MoE 架構，NeMo AutoModel 提供手動調校的實作，包含 TransformerEngine 注意力、融合的線性層以及自訂的專家核心。對於其他模型則回退至原生 Hugging Face，同時仍套用如 Liger 核心修補等最佳化。最終模型可透過傳入 device_mesh 進行多 GPU 訓練。

若要在 8 GPU 上使用 Expert Parallelism 訓練 Nemotron 3 Nano 30B A3B，使用者需加入分散式 mesh 設定：

import os
import torch
import torch.distributed as dist
from nemo_automodel import NeMoAutoModelForCausalLM
from nemo_automodel.recipes._dist_utils import create_distributed_setup_from_config

dist.init_process_group(backend="nccl
torch.manual_seed(0)
torch.cuda.set_device(int(os.environ.get("LOCAL_RANK", 0)))

dist_setup = create_distributed_setup_from_config({
    "strategy": "fsdp2",
    "ep_size": 8,
})

model = NeMoAutoModelForCausalLM.from_pretrained(
    "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16",
    dtype=torch.bfloat16,
    distributed_setup=dist_setup,
)

dist.destroy_process_group()

這樣即可從單一 from_pretrained() 呼叫中獲得來自 FSDP2、Expert Parallelism、TransformerEngine 核心與 DeepEP 分派的速度、可擴展性與記憶體最佳化。

效能比較

基準測試顯示，NeMo AutoModel 在每 GPU 的 tokens‑per‑second 上提升 3.4‑3.7 倍，峰值記憶體降低 29‑32%，相較於最佳的 Transformers v5 設定，無論是在 16 個節點的 550B 模型，或是單節點的兩個 30B MoE 模型上皆如此。

Nemotron 3 Ultra 550B A55B（完整微調，多節點）

Expert Parallelism 使得在 16 個 H100 節點上完整微調 550B 參數的 Nemotron 3 Ultra 模型成為可能，而在此規模下 Transformers v5 會因記憶體不足而失敗。

方法論：

  • 硬體：16 台 H100 80GB（共 128 顆 GPU）
  • Expert Parallelism：EP=64
  • 本地 batch 大小：2
  • 序列長度：4,096
  • 特性：MTP、激活檢查點、融合線性交叉熵
  • 核心：DeepEP 分派 + torch_mm 專家 + TransformerEngine

結果：

指標 NeMo AutoModel (EP=64)
TPS/GPU (avg) 815
TFLOP/s/GPU ~293
Peak Memory 58.2 GiB

Transformers v5 在此規模下記憶體不足，故未報告 v5 數值。

單節點 30B MoE 基準測試

在單一 8 GPU H100 節點上，NeMo AutoModel 在吞吐量上比 Transformers v5 快 3.36‑3.69 倍，且對 Qwen3‑30B‑A3B 與 Nemotron 3 Nano 30B‑A3B 的峰值記憶體降低 29‑32%。

  • 硬體：8 台 H100 80GB（單節點）
  • 序列長度：4,096
  • 本地 batch 大小：1

Qwen3-30B-A3B

指標 v4 v5 (FA2 + grouped_mm) NeMo AutoModel (EP=8) v5 → NeMo AutoModel
TPS/GPU (avg) 死鎖 3,075 11,340 3.69x
Peak Memory — 68.2 GiB 48.1 GiB -29%
Avg Forward+Loss — 582 ms 194 ms 3.00x
Avg Backward — 758 ms 178 ms 4.26x

Transformers v4 會死鎖，原因是它將 Qwen3 MoE 專家儲存為 128 個獨立 MLP 模組的 ModuleList，且每個模組分別被 FSDP 包裝，導致集合操作不匹配。Transformers v5 透過將專家儲存為融合的 3D 參數張量來解決此問題。

Nemotron 3 Nano 30B A3B

指標 v4 (hub code) v5 (FA2 + grouped_mm + Mamba CUDA) NeMo AutoModel (EP=8) v5 → NeMo AutoModel
TPS/GPU (avg) 1,807 4,583 15,421 3.36x
Peak Memory 61.9 GiB 62.1 GiB 42.5 GiB -32%
Avg Forward+Loss 1,024 ms 283 ms 109 ms 2.60x
Avg Backward 1,246 ms 611 ms 157 ms 3.89x

v4 使用 trust_remote_code=True（NVIDIA 的 hub 建模程式碼），其專家迴圈會遍歷所有專家，不論 token 分配為何，從而避免了 Qwen3 v4 中的死鎖。

加速來源

3.4‑3.7 倍的加速來源於 Expert Parallelism 對專家權重的分片、DeepEP 將 token 路由與計算融合，以及 TransformerEngine 核心加速注意力與線性層。

  1. Expert Parallelism 減少記憶體壓力：EP=8 將專家權重分散至多個 GPU，使每個 GPU 的 MoE 佔用空間縮小 8 倍。對於 Qwen3，峰值記憶體從 68.2 GiB 降至 48.1 GiB（‑29%）。對於 Nemotron Nano，則從 62.1 GiB 降至 42.5 GiB（‑32%），為更大的 batch 或更長的序列提供空間。
  2. DeepEP 將通信與計算融合：不再使用獨立的 AllGather/ReduceScatter 來進行專家路由，DeepEP 將 token 分派融合至最佳化的 GPU 核心，實現通信與專家計算的重疊。
  3. TransformerEngine 核心加速核心運算：TE 的融合注意力、線性層與 RMSNorm 實作，在所有層型（不僅限於 MoE 層）相較於 PyTorch/Flash Attention 等等皆提供穩定的加速。

Transformers v5 功能在 HuggingFace AutoModel 中的應用

NeMo AutoModel 基於 Transformers v5 的專家後端（尤其是 grouped_mm）、動態權重載入與張量平行的 Expert Parallelism，並在此基礎上加入 DeepEP 與 TransformerEngine。

專家後端

Transformers v5 中最具影響力的功能之一是 experts_implementation 參數，包含三種專家後端：

後端 描述 最適用於
eager 對選定的專家進行 for‑loop 除錯、相容性與正確性。亦可於 v4 使用。
batched_mm 複製專家參數，透過 torch.bmm 執行單一批次 GEMM 輸入較小、使用 torch.compile 時快速。v5 新增
grouped_mm 依專家排序 token，透過 torch.nn.functional.grouped_mm 執行單一分組 GEMM 訓練（記憶體效率高，無參數複製）。v5 新增

grouped_mm 後端是關鍵的訓練最佳化：它不會逐一迴圈專家，而是依分配的專家排序 token，並執行單一融合的分組矩陣乘法。

NeMo AutoModel 更進一步。對於具有自訂實作的模型，它結合 DeepEP 融合的 all‑to‑all 分派、分組 GEMM 核心與 TransformerEngine 線性層。其演進如下：

v4 (eager for‑loop) → v5 (grouped_mm) → NeMo AutoModel (DeepEP + GMM + TE)

在 NeMo AutoModel 中，專家後端透過 BackendConfig 設定：

from nemo_automodel.components.models.common.utils import BackendConfig

backend = BackendConfig(
    attn="te",           # TransformerEngine attention
    linear="te",         # TransformerEngine linear layers
    experts="torch_mm",  # Grouped expert matmul
    dispatcher="deepep", # DeepEP fused all-to-all
)

Expert Parallelism 與 DeepEP

Transformers v5 也提供了將專家權重在 GPU 之間分片的 Expert Parallelism 路徑。NeMo AutoModel 採取互補的方式，針對多 GPU MoE 訓練進行調校。它將 EP 作為獨立的平行維度，建立專屬的 moe_mesh，與資料平行 mesh 並行，使用 PyTorch 的 DTensor 與 Shard(0)。由於專家 mesh 與資料平行是正交的，兩者可在相同裝置上組合。在 8 GPU 上，NeMo AutoModel 同時執行 ep=8 與 dp=8，使每個 GPU 在訓練自己的資料分片時，只持有 1/8 的專家。專家權重在 GPU 上沿專家維度實際分片。

# From nemo_automodel/components/moe/parallelizer.py
from torch.distributed.tensor import Shard, distribute_tensor

# Each GPU holds only 1/ep_size of the expert weights
distribute_tensor(param, device_mesh, [Shard(0)])

在 8 GPU 上使用 ep_size=8 時，每個 GPU 只持有 1/8 的專家參數。以 Nemotron‑3‑Nano‑30B‑A3B（約 55 GiB 專家權重）為例，EP 可將每 GPU 的專家佔用空間從約 55 GiB 降至約 6.8 GiB，使得僅使用 FSDP 的方法在記憶體不足的情況下仍能訓練。

在 EP 之上，NeMo AutoModel 整合 DeepEP，將 token 路由融合至最佳化的 GPU 核心，並在結合分組 GEMM 進行分組專家計算時提供顯著的加速。

動態權重載入

Transformers v5 亦透過 WeightConverter 與 WeightRenaming 引入了動態權重載入系統。此機制允許將 MoE 檢查點儲存為融合的 3D 張量，以提升執行效率。WeightConverter 在 from_pretrained() 時即時對檢查點張量套用可組合的轉換操作。

NeMo AutoModel 直接使用此 v5 API。超過 20 種模型類型透過 MODELS_REQUIRING_TENSOR_MERGING 使用此機制，包括 Mixtral、Qwen2 MoE、Qwen3 MoE、DeepSeek V2/V3、OLMoE 等。這些轉換是完全可逆的：save_pretrained() 會產生標準 HF 格式的檢查點，任何下游工具皆可載入。

快速入門

使用者可依照 NVIDIA 的文件安裝 NeMo AutoModel，並僅透過一次 import 變更，即可使用相同的 from_pretrained() API。

欲了解更多細節，請參考：

  • NeMo AutoModel HuggingFace API 相容性指南
  • NeMo AutoModel 模型覆蓋
  • NeMo AutoModel 效能摘要
  • NeMo AutoModel 在 HuggingFace 上

結論

NeMo AutoModel 為擴展 MoE 訓練的 Hugging Face 使用者提供零摩擦的升級路徑，帶來多倍的加速與記憶體節省，同時產生與 vLLM 與 SGLang 相容的標準 Hugging Face 檢查點。

程式碼、設定檔與基準測試腳本皆可在 NeMo AutoModel 倉庫中取得。

Sources