NVIDIA NeMo AutoModel Delivers 3.7x Faster MoE Fine-Tuning via Expert Parallelism and DeepEP

NVIDIA has rolled out NeMo AutoModel, an open-source library inside the NeMo framework that brings sharp performance gains to fine-tuning mixture-of-experts (MoE) models. Built directly on Hugging Face Transformers v5, the library delivers 3.4 to 3.7 times higher training throughput and cuts GPU memory consumption by 29 to 32 percent compared with the best available Transformers v5 configuration, all while keeping the same from_pretrained() API. Users only have to change a single import line to claim those benefits.

Background: The MoE Efficiency Challenge

The rise of MoE models has introduced new hurdles for efficient training. Routing tokens across hundreds of experts, fusing expert matrix multiplications into single kernels, sharding weights across GPUs, and overlapping communication with computation all require infrastructure that general-purpose libraries typically don't provide out of the box.

Transformers v5 added first-class MoE support, including expert backends, dynamic weight loading, and tensor parallel plans for distributed execution. But key performance pieces, such as efficient communication overlap, remained unaddressed. NeMo AutoModel plugs that gap.

How NeMo AutoModel Works

NeMo AutoModel subclasses AutoModelForCausalLM and layers in three major optimizations: Expert Parallelism (EP), DeepEP fused all-to-all dispatch, and TransformerEngine kernels. DeepEP overlaps communication with expert computation, a capability Transformers v5 lacks. Because NeMo AutoModel taps into v5's reversible weight conversion to load models, it focuses engineering effort on reusable core operations rather than per-model checkpoint plumbing.

The library stays fully compatible with the Hugging Face Transformers API. Loading a model requires only a swap in the import statement, as NVIDIA's code examples demonstrate.

Performance Benchmarks

Multi-Node: Nemotron 3 Ultra 550B A55B

For full fine-tuning of the 550B-parameter Nemotron 3 Ultra A55B model across 16 H100 nodes (128 GPUs), NeMo AutoModel with Expert Parallelism set to 64 hit 815 tokens per second per GPU, 293 TFLOPs per second per GPU, and peak memory of 58.2 GiB. Critically, Transformers v5 runs out of memory at this scale, making a direct comparison impossible without EP.

Single-Node: Qwen3-30B-A3B

On a single node with 8x H100 80GB GPUs, NeMo AutoModel (EP=8) logged 11,340 TPS/GPU, a 3.69x jump over Transformers v5's 3,075 TPS/GPU. Peak memory fell from 68.2 GiB to 48.1 GiB, a 29 percent reduction. Notably, Transformers v4 deadlocked entirely on this model due to incompatible FSDP wrapping of individual expert modules.

Single-Node: Nemotron 3 Nano 30B A3B

For the 30B Nemotron Nano model, NeMo AutoModel reached 15,421 TPS/GPU, a 3.36x improvement over v5's 4,583 TPS/GPU. Peak memory consumption dropped from 62.1 GiB to 42.5 GiB, a 32 percent reduction.

Sources of Speedup

The performance gains come from three coordinated optimizations:

• Expert Parallelism shards expert weights across GPUs. With EP=8 on 8 GPUs, each GPU holds only one-eighth of the expert parameters. For a model with roughly 55 GiB of expert weights, EP shrinks the per-GPU footprint to about 6.8 GiB, enabling training where FSDP-only approaches run out of memory.
• DeepEP fuses token dispatch and combining into optimized GPU kernels, overlapping communication with expert computation rather than using separate AllGather and ReduceScatter collectives.
• TransformerEngine kernels accelerate core operations like fused attention, linear layers, and RMSNorm, delivering consistent speedups across all layer types.

Leveraging Transformers v5 Features

NeMo AutoModel builds on several key v5 features, including the expert backends system. v5 introduced the experts_implementation parameter with three options: eager (for-loop debugging), batched_mm (single batched GEMM), and grouped_mm (sorted token execution). NeMo AutoModel extends that by using DeepEP fused all-to-all dispatch combined with grouped GEMM kernels and TransformerEngine linear layers.

Dynamic weight loading through v5's WeightConverter and WeightRenaming lets MoE checkpoints be stored in fused 3D tensors. Over 20 model types use this mechanism, including Mixtral, Qwen2 MoE, Qwen3 MoE, DeepSeek V2/V3, and OLMoE. The conversions are fully reversible: save_pretrained() produces standard Hugging Face format safetensors that downstream tools like vLLM and SGLang can load.

Conclusion and Availability

NVIDIA NeMo AutoModel offers a zero-friction upgrade path for Hugging Face users scaling up model training. By changing a single import line, users get a model instance that's more than three times as fast with a much smaller memory footprint. The code, configs, and benchmark scripts are available in the NeMo AutoModel repository. Core contributors to the work include Adil Asif, Hemil Desai, Alexandros Koumparoulis, and Huiying Li.