NVIDIA NeMo AutoModel Accelerates Fine-Tuning of Mixture-of-Experts Models
NVIDIA NeMo AutoModel Accelerates Fine-Tuning of Mixture-of-Experts Models
NVIDIA NeMo AutoModel accelerates fine-tuning of Mixture-of-Experts models, delivering 3.4-3.7x higher training throughput and 29-32% lower GPU memory usage compared to Hugging Face Transformers v5, while requiring only a one-line import change.
Background
The rise of Mixture-of-Experts models created training challenges that Transformers v5 addressed with expert backends, dynamic weight loading, and distributed execution, but still lacked fused communication kernels.
NeMo AutoModel: Same API, More Performance
NeMo AutoModel subclasses AutoModelForCausalLM, so users can switch from Hugging Face Transformers to NeMo AutoModel by changing a single import line and gain Expert Parallelism, DeepEP fused all-to-all dispatch, and TransformerEngine kernels without altering training code.
Loading a model with NeMo AutoModel looks like the Hugging Face version except for the import:
from nemo_automodel import NeMoAutoModelForCausalLM
model = NeMoAutoModelForCausalLM.from_pretrained(
"nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16",
dtype=torch.bfloat16,
)
For popular MoE architectures such as Qwen3, NVIDIA Nemotron, GPT-OSS, and DeepSeek V3, NeMo AutoModel provides hand‑tuned implementations with TransformerEngine attention, fused linear layers, and custom expert kernels. For other models it falls back to vanilla Hugging Face while still applying optimizations like Liger kernel patching. The resulting model scales to multi‑GPU training by passing a device_mesh.
To train Nemotron 3 Nano 30B A3B with Expert Parallelism across 8 GPUs, users add a distributed mesh configuration:
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()
This yields speed, scalability and memory optimizations from FSDP2, Expert Parallelism, TransformerEngine kernels and DeepEP dispatch, all from a single from_pretrained() call.
Performance Comparison
Benchmarks show NeMo AutoModel achieves 3.4-3.7x higher tokens-per-second per GPU and 29-32% lower peak memory than the best Transformers v5 configuration on both a 550B model across 16 nodes and two 30B MoE models on a single node.
Nemotron 3 Ultra 550B A55B (full fine-tune, multi-node)
Expert Parallelism enables full fine‑tuning of the 550B‑parameter Nemotron 3 Ultra model on 16 H100 nodes, a scale at which Transformers v5 runs out of memory.
Methodology:
- 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
Results:
| Metric | NeMo AutoModel (EP=64) |
|---|---|
| TPS/GPU (avg) | 815 |
| TFLOP/s/GPU | ~293 |
| Peak Memory | 58.2 GiB |
Transformers v5 runs out of memory at this scale, so no v5 number is reported.
Single-node 30B MoE benchmarks
On a single 8‑GPU H100 node, NeMo AutoModel outperforms Transformers v5 by 3.36‑3.69x in throughput and reduces peak memory by 29‑32% for Qwen3‑30B‑A3B and Nemotron 3 Nano 30B‑A3B.
Methodology:
- Hardware: 8x H100 80GB (single node)
- Sequence length: 4,096
- Local batch size: 1
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 |
Transformers v4 deadlocks because it stores Qwen3 MoE experts as a ModuleList of 128 individual MLP modules, each separately FSDP‑wrapped, causing mismatched collectives. Transformers v5 fixes this by storing experts as fused 3D parameter tensors.
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 uses trust_remote_code=True (NVIDIA's hub modeling code) whose expert loop iterates all experts regardless of token assignment, avoiding the deadlock seen in Qwen3 v4.
Where the speedup comes from
The 3.4‑3.7x speedup stems from Expert Parallelism sharding expert weights, DeepEP fusing token routing with computation, and TransformerEngine kernels accelerating attention and linear layers.
- Expert Parallelism reduces memory pressure: EP=8 distributes expert weights across GPUs, cutting the per‑GPU MoE footprint by 8×. For Qwen3 this drops peak memory from 68.2 GiB to 48.1 GiB (‑29%). For Nemotron Nano it drops from 62.1 GiB to 42.5 GiB (‑32%), freeing headroom for larger batch sizes or longer sequences.
- DeepEP fuses communication with computation: instead of separate AllGather/ReduceScatter collectives for expert routing, DeepEP fuses token dispatch and combines into optimized GPU kernels, overlapping communication with expert computation.
- TransformerEngine kernels accelerate core operations: TE’s fused attention, linear layers, and RMSNorm implementations provide consistent speedups over their PyTorch/Flash Attention equivalents across all layer types, not just MoE layers.
Transformers v5 Features Leveraged by HuggingFace AutoModel
NeMo AutoModel builds on Transformers v5’s expert backends (especially grouped_mm), dynamic weight loading, and tensor‑parallel Expert Parallelism, adding DeepEP and TransformerEngine on top.
Expert Backends
One of the most impactful features in Transformers v5 is the experts_implementation parameter, which includes three expert backends:
| 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 |
The grouped_mm backend is the key training optimization: instead of looping over experts one by one, it sorts tokens by their assigned expert and executes a single fused grouped matrix multiplication.
NeMo AutoModel takes this further. For models with custom implementations, it uses DeepEP fused all‑to‑all dispatch combined with grouped GEMM kernels and TransformerEngine linear layers. The progression looks like:
v4 (eager for‑loop) → v5 (grouped_mm) → NeMo AutoModel (DeepEP + GMM + TE)
In NeMo AutoModel, the expert backend is configured through 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 and DeepEP
Transformers v5 also ships an Expert Parallelism path that shards expert weights across GPUs. NeMo AutoModel takes a complementary approach tuned for multi‑GPU MoE training. It makes EP its own parallelism dimension, a dedicated moe_mesh alongside (rather than carved from) the data‑parallel mesh, using PyTorch’s DTensor with Shard(0). Because the expert mesh is orthogonal to data parallelism, the two compose on the same devices. On 8 GPUs NeMo AutoModel runs ep=8 and dp=8 together, so every GPU trains on its own data shard while holding only 1/8 of the experts. Expert weights are physically sharded across GPUs along the expert dimension.
# 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)])
With ep_size=8 on 8 GPUs, each GPU holds only 1/8 of the expert parameters. For a model like Nemotron‑3‑Nano‑30B‑A3B with ~55 GiB of expert weights, EP reduces the per‑GPU expert footprint from ~55 GiB to ~6.8 GiB, making training possible where FSDP‑only approaches run out of memory.
On top of EP, NeMo AutoModel integrates DeepEP that fuses the token routing into optimized GPU kernels, and delivers significant speedups when combined with grouped GEMM for grouped expert computation.
Dynamic Weight Loading
Transformers v5 also introduced a dynamic weight loading system through WeightConverter and WeightRenaming. This enables MoE checkpoints to be stored in fused 3D tensors for more efficient execution. The WeightConverter applies composable operations to transform checkpoint tensors on‑the‑fly during from_pretrained().
NeMo AutoModel is a direct consumer of this v5 API. Over 20 model types use this mechanism through MODELS_REQUIRING_TENSOR_MERGING, including Mixtral, Qwen2 MoE, Qwen3 MoE, DeepSeek V2/V3, OLMoE, and more. The conversions are fully reversible: save_pretrained() produces standard HF‑format checkpoints that any downstream tool can load.
Getting Started
Users can install NeMo AutoModel from NVIDIA’s documentation and use the same from_pretrained() API with a single import change.
For more details, see:
- NeMo AutoModel HuggingFace API Compatibility Guide
- NeMo AutoModel Model Coverage
- NeMo AutoModel Performance Summary
- NeMo AutoModel on HuggingFace
Conclusion
NeMo AutoModel offers a zero‑friction upgrade path for Hugging Face users scaling MoE training, delivering multi‑x speedups and memory savings while producing standard Hugging Face checkpoints compatible with vLLM and SGLang.
The code, configs, and benchmark scripts are all available in the NeMo AutoModel repository.