運用 NVIDIA NeMo AutoModel 大幅加速 Transformer 模型微調

HuggingFace Transformers 已成為開源 AI 生態系統的基石，而近期發布的 Transformers v5 版本，透過對 Mixture-of-Experts (MoE) 模型的一流支援，進一步強化了其地位，MoE 模型目前已是前沿模型的主流架構。

v5 提供了 MoE 的基礎設施：專家後端、動態權重載入和分散式執行，使 MoE 易於擴展和建構。

NVIDIA NeMo AutoModel 是 NVIDIA NeMo 框架的一部分，這是一個用於大規模建構客製化生成式 AI 模型的開源函式庫。NeMo AutoModel 在 v5 基礎上，增加了 Expert Parallelism、DeepEP 融合式 all-to-all 分派和 TransformerEngine 核心，並利用 v5 的動態權重載入，將這些優化應用於廣泛且不斷增長的模型家族。

其成果是，在微調 MoE 模型時，訓練吞吐量比原生 Transformers v5 高出 3.4-3.7 倍，GPU 記憶體使用減少 29-32%，且僅需使用相同的 from_pretrained() API，無需修改其他程式碼。

這篇部落格文章詳細介紹了這種組合如何運作，以及使用者如何在不改變 API 的情況下，更快地微調 MoE 模型。

背景介紹

MoE 模型的興起為高效訓練帶來了新挑戰：在數百個專家之間路由 tokens、將專家矩陣乘法融合到單一核心、在 GPU 之間分片權重，以及重疊通訊與計算，這些都需要通用函式庫無法直接提供的基礎設施。

Transformers v5（簡稱 v5）引入了對 MoE 的一流支援，例如專家後端、動態權重載入和用於分散式執行的張量平行計畫。此外，v5 還透過將 PyTorch 的 DeviceMesh 直接整合到 from_pretrained() 中，使分散式訓練成為一流功能。

NeMo AutoModel 透過繼承 AutoModelForCausalLM 並增加 Expert Parallelism (EP)、DeepEP 融合式 all-to-all 分派和 TransformerEngine 核心，在 v5 基礎上進行建構。DeepEP 是 v5 尚未具備的部分：它能將通訊與專家計算重疊。

由於 NeMo AutoModel 透過 v5 的可逆權重轉換來載入每個模型，因此它可以將工程重心放在這些可重複使用的核心操作上，而不是針對每個模型的檢查點處理，同時 save_pretrained() 仍會產生標準的 HF 檢查點，供 vLLM 和 SGLang 等工具載入。

下一節將介紹兩者如何協同運作，以及我們測量到的效能提升，從跨 16 個節點對 NVIDIA Nemotron 3 Ultra 550B A55B 進行完整微調，到單節點模型如 Qwen3-30B-A3B 和 Nemotron 3 Nano 30B A3B 的微調。

NeMo AutoModel：相同的 API，更高的效能

NeMo AutoModel 的目標之一是與 HuggingFace Transformers 保持 API 相容性，以促進開源社群的發展。NeMoAutoModelForCausalLM 繼承自 AutoModelForCausalLM，因此任何適用於 HF 模型的程式碼也適用於 AutoModel。

以下是兩種載入模型的方式，只有匯入語句不同。

僅僅一個匯入語句就完成了許多工作。對於 Qwen3、NVIDIA Nemotron、GPT-OSS 和 DeepSeek V3 等流行的 MoE 架構，NeMo AutoModel 提供了手動優化的實作，包含 TransformerEngine attention、融合式線性層和客製化專家核心。

對於其他模型，它會回退到原生的 HF 實作，同時仍應用 Liger 核心修補等優化。無論採用哪種路徑，產生的模型都已準備好進行擴展：傳入 device_mesh 即可進行多 GPU 訓練，無需進一步重寫程式碼。

NeMo AutoModel 真正出色之處在於將 MoE 模型擴展到多 GPU 訓練。要在 8 個 GPU 上使用 Expert Parallelism 訓練 Nemotron 3 Nano 30B A3B，只需添加分散式網格配置：

```python

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：跨 16 個節點對一個前沿規模的 550B 模型進行完整微調，以及在單一節點上訓練兩個 30B MoE 模型。550B 的結果展示了 Expert Parallelism 在大規模應用中的重要性；30B 的結果則量化了相對於 Transformers v5 的每個 GPU 加速。

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

Nemotron 3 Ultra 550B A55B 是一個 550B 參數的混合模型，包含 Mamba2、LatentMoE 和 Multi-Token Prediction (MTP)。我們對其進行了完整微調基準測試：每個參數都經過更新，並且 Adam 優化器狀態被實體化，在這種規模下，這需要跨越 16 個 H100 節點（128 個 GPU）。

方法：

| Parameter | Value |

|---|---|

| Hardware | 16x H100 80GB (128 GPUs) |

| Expert Parallelism | EP=64 |

| Local batch size | 2 |

| Sequence length | 4,096 |

| Features | MTP, activation checkpointing, fused linear cross-entropy |

| Kernels | DeepEP dispatch + torch_mm experts + TransformerEngine |

| Metric | NeMo AutoModel (EP=64) |

|---|---|

| TPS/GPU (avg) | 815 |

| TFLOP/s/GPU | ~293 |

| Peak Memory | 58.2 GiB |

為何沒有 Transformers v5 的欄位？Transformers v5 在這種規模下會記憶體不足，因此沒有 v5 的數據可供報告。AutoModel 的 Expert Parallelism 將專家分片到不同的 GPU 上，將記憶體佔用控制在預算內，這使得完整微調得以運行。

下面的 30B 比較也顯示了在 v5 能夠運行的情況下，AutoModel 具有相同的優勢。

單節點 30B MoE 基準測試

我們在一個配備 8 個 H100 80GB GPU 的單一節點上，對三種方法進行了基準測試：HF Transformers v4（hub 程式碼）、HF Transformers v5（採用最佳可用優化）和 NeMo AutoModel（EP=8 + 客製化核心）。

方法：

| Parameter | Value |

|---|---|

| Hardware | 8x H100 80GB (single node) |

| Sequence length | 4,096 |

| Local batch size | 1 |

關於路由閘道 (routing gate) 的說明。下面的 NeMo AutoModel 數據使用了平衡路由閘道，它強制 tokens 均勻分佈到各個專家。這模擬了 MoE 訓練的理想操作點：一個訓練良好的模型，其負載平衡損失會使專家利用率接近均勻，因此平衡路由反映了實際工作負載收斂到的穩定狀態（並消除了隨機虛擬 tokens 可能注入到專家平行中的雜訊）。

v4/v5 則在相同的虛擬 tokens 上運行其原生路由器。因此，平衡閘道測量的是 NeMo AutoModel 在其目標 MoE 操作點下的表現，而 v4/v5 欄位則反映了它們開箱即用的行為。

Qwen3-30B-A3B

| Metric | v4 | v5 (FA2 + grouped_mm) | NeMo AutoModel (EP=8) | v5 → NeMo AutoModel |

|---|---|---|---|---|

| TPS/GPU (avg) | deadlock | 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 |

為何 v4 會死鎖：Transformers v4 將 Qwen3 MoE 專家儲存為 128 個獨立 MLP 模組的 ModuleList，每個模組都單獨進行 FSDP 包裝。前向傳播使用數據依賴的迴圈，只迭代接收到 tokens 的專家。

由於每個 rank 的數據不同，不同的 rank 會跳過不同的專家，導致 FSDP AllGather/ReduceScatter 集合操作不匹配，進而導致無限期掛起。Transformers v5 透過將專家儲存為融合的 3D 參數張量（沒有每個專家的模組，也沒有每個專家的 FSDP 集合操作）來解決此問題。

Nemotron 3 Nano 30B A3B

| Metric | 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 建模程式碼）。hub 程式碼的專家迴圈是 FSDP 安全的（無論 token 分配如何，都會迭代所有專家），因此它不會像 Qwen3 v4 那樣死鎖。

加速的來源

NeMo AutoModel 相較於 Transformers v5 實現 3.4-3.7 倍加速，主要來自三個方面：

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%），為更大的批次大小或更長的序列騰出空間。

DeepEP 融合通訊與計算。DeepEP 不再為專家路由使用單獨的 AllGather/ReduceScatter 集合操作，而是融合 token 分派並將其組合到優化的 GPU 核心中，從而實現通訊與專家計算的重疊。

TransformerEngine 核心加速核心操作。TransformerEngine 的融合式 attention、線性層和 RMSNorm 實作，在所有層類型（不僅限於 MoE 層）上都比 PyTorch/Flash Attention 的等效實作提供一致的加速。

HuggingFace AutoModel 利用的 Transformers v5 功能

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

| Backend | Description | Best for |

|---|---|---|

| eager | For-loop over selected experts | Debugging, compatibility, and correctness. Also available for v4. |

| batched_mm | Duplicates expert params, single batched GEMM via torch.bmm | Small inputs, fast with torch.compile. Added for v5 |

| grouped_mm | Orders tokens by expert, single grouped GEMM via torch.nn.functional.grouped_mm | Training (memory efficient, no param duplication). Added for v5. |

grouped_mm 後端是關鍵的訓練優化：它不是逐一迴圈處理專家，而是根據分配的專家對 tokens 進行排序，並執行單一融合的 grouped matrix multiplication。

NeMo AutoModel 更進一步。對於客製化實作的模型，它使用 DeepEP 融合式 all-to-all 分派結合 grouped GEMM 核心和 TransformerEngine 線性層。其演進路徑如下：v4 (eager for-loop) → v5 (grouped_mm) → NeMo AutoModel (DeepEP + GMM + TE)。

在 NeMo AutoModel 中，專家後端透過 BackendConfig 進行配置：

```python

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 也提供了 Expert Parallelism 路徑。它將專家權重分片到多個 GPU 上。GroupedGemmParallel 樣式只載入每個設備的本地專家，而 RouterParallel 則路由 tokens 並透過 all_reduce 組合結果。

它巧妙地建立在 v5 現有的張量平行機制之上。啟用它會使模型的 tp_plan 返回其專家計畫，因此專家平行與數據平行共享設備預算（ep × dp = world_size）。對於此處的單節點 30B 基準測試，我們發現純數據平行 v5 (dp=8, ep=1) 是最快的 v5 配置，因此這是我們報告的 v5 設定。

NeMo AutoModel 採用了一種互補的方法，專為多 GPU MoE 訓練而優化。它將 EP 設為自己的平行維度，一個專用的 moe_mesh，與數據平行網格並行（而非從中劃分），使用 PyTorch 的 DTensor 搭配 Shard(0)。

由於專家網格與數據平行正交，兩者可以在相同的設備上組合。在 8 個 GPU 上，NeMo AutoModel 同時運行 ep=8 和 dp=8，因此每個 GPU 在處理自己的數據分片時，只持有 1/8 的專家。專家權重沿著專家維度物理分片到 GPU 上。