NVIDIA NeMo AutoModel が Mixture-of-Experts モデルのファインチューニングを加速
NVIDIA NeMo AutoModel が Mixture-of-Experts モデルのファインチューニングを加速
背景
Mixture-of-Experts モデルの台頭により、トレーニング上の課題が生じました。Transformers v5 はエキスパートバックエンド、動的重みロード、分散実行でこれらに対処しましたが、融合通信カーネルは依然として欠けていました。
NeMo AutoModel: 同じ API、より高いパフォーマンス
NeMo AutoModel は AutoModelForCausalLM をサブクラス化しているため、ユーザーはインポート行を 1 行変更するだけで Hugging Face Transformers から NeMo AutoModel に切り替えられ、エキスパート並列、DeepEP 融合 all‑to‑all ディスパッチ、TransformerEngine カーネルをトレーニングコードを変更せずに利用できます。
モデルのロードは Hugging Face バージョンと同様ですが、インポートが異なります:
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 アテンション、融合線形層、カスタムエキスパートカーネルを備えた手調整実装を提供します。その他のモデルについては、最適化（Liger カーネルパッチなど）を適用しつつ、標準の Hugging Face にフォールバックします。結果として得られるモデルは、device_mesh を渡すことでマルチ GPU トレーニングにスケールします。
Nemotron 3 Nano 30B A3B を 8 GPU にまたがるエキスパート並列でトレーニングするには、分散メッシュ設定を追加します:
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()
これにより、FSDP2、エキスパート並列、TransformerEngine カーネル、DeepEP ディスパッチから得られる速度、スケーラビリティ、メモリ最適化がすべて単一の from_pretrained() 呼び出しで実現します。
パフォーマンス比較
ベンチマークは、NeMo AutoModel が 550B モデル（16 ノード）と 2 つの 30B MoE モデル（単一ノード）において、GPU あたりトークン/秒を 3.4‑3.7 倍に向上させ、ピークメモリを 29‑32% 低減できることを示しています。
Nemotron 3 Ultra 550B A55B（フルファインチューニング、マルチノード）
方法論:
- ハードウェア: 16x H100 80GB (128 GPU)
- Expert Parallelism: EP=64
- ローカルバッチサイズ: 2
- シーケンス長: 4,096
- 機能: MTP、アクティベーションチェックポイント、融合線形クロスエントロピー
- カーネル: DeepEP ディスパッチ + torch_mm エキスパート + TransformerEngine
結果:
| 指標 | NeMo AutoModel (EP=64) |
|---|---|
| TPS/GPU（平均） | 815 |
| TFLOP/s/GPU | ~293 |
| ピークメモリ | 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% 削減します。
方法論:
- ハードウェア: 8x H100 80GB（シングルノード）
- シーケンス長: 4,096
- ローカルバッチサイズ: 1
Qwen3-30B-A3B
| 指標 | v4 | v5 (FA2 + grouped_mm) | NeMo AutoModel (EP=8) | v5 → NeMo AutoModel |
|---|---|---|---|---|
| TPS/GPU（平均） | deadlock | 3,075 | 11,340 | 3.69x |
| ピークメモリ | — | 68.2 GiB | 48.1 GiB | -29% |
| 平均 Forward+Loss | — | 582 ms | 194 ms | 3.00x |
| 平均 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（平均） | 1,807 | 4,583 | 15,421 | 3.36x |
| ピークメモリ | 61.9 GiB | 62.1 GiB | 42.5 GiB | -32% |
| 平均 Forward+Loss | 1,024 ms | 283 ms | 109 ms | 2.60x |
| 平均 Backward | 1,246 ms | 611 ms | 157 ms | 3.89x |
v4 は trust_remote_code=True（NVIDIA のハブモデリングコード）を使用し、トークン割り当てに関係なくすべてのエキスパートを走査するループを持つため、Qwen3 v4 で見られたデッドロックを回避します。
スピードアップの要因
3.4‑3.7 倍のスピードアップは、エキスパート並列によるエキスパート重みのシャーディング、DeepEP によるトークンルーティングと計算の融合、そして TransformerEngine カーネルによるアテンションと線形層の高速化に起因します。
- エキスパート並列はメモリ圧迫を軽減: EP=8 はエキスパート重みを GPU 間で分散し、GPU あたりの MoE フットプリントを 8 倍削減します。Qwen3 ではピークメモリが 68.2 GiB から 48.1 GiB（‑29%）に、Nemotron Nano では 62.1 GiB から 42.5 GiB（‑32%）に減少し、より大きなバッチサイズや長いシーケンスの余裕が生まれます。
- DeepEP は通信と計算を融合: エキスパートルーティング用の個別 AllGather/ReduceScatter の代わりに、DeepEP はトークンディスパッチを融合し、最適化された GPU カーネルに統合して通信とエキスパート計算をオーバーラップさせます。
- TransformerEngine カーネルはコア演算を加速: TE の融合アテンション、線形層、RMSNorm 実装は、PyTorch/Flash Attention の同等実装に対して一貫した高速化を提供し、MoE 層に限らずすべての層で効果を発揮します。
HuggingFace AutoModel が活用する Transformers v5 の機能
エキスパートバックエンド
Transformers v5 の最もインパクトのある機能の一つは experts_implementation パラメータで、以下の 3 つのエキスパートバックエンドを提供します:
| バックエンド | 説明 | 適用対象 |
|---|---|---|
| eager | 選択されたエキスパートに対する for‑loop | デバッグ、互換性、正確性。v4 でも利用可能。 |
| batched_mm | エキスパートパラメータを複製し、torch.bmm による単一のバッチ GEMM | 小規模入力、torch.compile で高速。v5 用に追加。 |
| grouped_mm | エキスパート別にトークンを並べ替え、torch.nn.functional.grouped_mm による単一のグループ化 GEMM | トレーニング（メモリ効率が高く、パラメータ複製なし）。v5 用に追加。 |
grouped_mm バックエンドはトレーニング最適化の鍵です。エキスパートを一つずつループする代わりに、割り当てられたエキスパートごとにトークンをソートし、単一の融合グループ化行列乗算を実行します。
NeMo AutoModel はこれをさらに進めます。カスタム実装を持つモデルでは、DeepEP 融合 all‑to‑all ディスパッチと grouped 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
)
エキスパート並列と DeepEP
Transformers v5 もエキスパート並列パスを提供し、エキスパート重みを GPU 間でシャーディングします。NeMo AutoModel はマルチ GPU MoE トレーニング向けに調整された補完的アプローチを採用し、ep_size をデータ並列メッシュとは別の専用 moe_mesh として扱い、PyTorch の DTensor と Shard(0) を使用します。エキスパートメッシュはデータ並列と直交しているため、同一デバイス上で両者が合成されます。8 GPU 環境では ep=8 と dp=8 を同時に走らせ、各 GPU は自分のデータシャードでトレーニングしつつ、エキスパートの 1/8 しか保持しません。エキスパート重みはエキスパート次元に沿って物理的にシャーディングされます。
# 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)])
ep_size=8 の 8 GPU 環境では、各 GPU がエキスパートパラメータの 1/8 しか保持しません。たとえば Nemotron‑3‑Nano‑30B‑A3B のエキスパート重みが約 55 GiB の場合、EP により GPU あたりのエキスパートフットプリントは約 55 GiB から約 6.8 GiB に削減され、FSDP のみの手法がメモリ不足で失敗するケースでもトレーニングが可能になります。
EP に加えて、NeMo AutoModel は DeepEP を統合し、トークンルーティングを最適化された GPU カーネルに融合させ、グループ化 GEMM と組み合わせることで大幅な速度向上を実現します。
動的重みロード
Transformers v5 は WeightConverter と WeightRenaming を通じて動的重みロードシステムを導入しました。これにより MoE チェックポイントを融合 3D テンソルとして保存でき、実行時の効率が向上します。WeightConverter は from_pretrained() 中にチェックポイントテンソルをオンザフライで変換する合成可能な操作を適用します。
NeMo AutoModel はこの v5 API を直接利用しています。MODELS_REQUIRING_TENSOR_MERGING を介して 20 以上のモデルタイプ（Mixtral、Qwen2 MoE、Qwen3 MoE、DeepSeek V2/V3、OLMoE など）がこの仕組みを使用します。変換は完全に可逆で、save_pretrained() は標準 HF 形式のチェックポイントを生成し、任意の下流ツールがロード可能です。
はじめに
ユーザーは NVIDIA のドキュメントから NeMo AutoModel をインストールし、インポートを 1 行変更するだけで同じ from_pretrained() API を利用できます。
詳細は以下をご参照ください:
- NeMo AutoModel HuggingFace API Compatibility Guide
- NeMo AutoModel Model Coverage
- NeMo AutoModel Performance Summary
- NeMo AutoModel on HuggingFace
結論
NeMo AutoModel は、MoE トレーニングをスケールさせる Hugging Face ユーザーに対し、フリクションゼロのアップグレードパスを提供し、マルチ倍の速度向上とメモリ削減を実現します。また、vLLM や SGLang と互換性のある標準 Hugging Face チェックポイントを生成します。
コード、設定ファイル、ベンチマークスクリプトはすべて NeMo AutoModel リポジトリで入手可能です。