NVIDIA NeMo AutoModel: Revolutionizing Transformer Fine-Tuning

The Rise of Transformers and NVIDIA's Innovation

In the field of artificial intelligence, Transformers from HuggingFace have established themselves as a cornerstone of the open-source ecosystem. The release of Transformers v5 has solidified this position, particularly due to the integration of Mixture-of-Experts (MoE) models, which have become the standard for state-of-the-art architectures. This version introduced essential features such as expert backends, dynamic weight loading, and distributed execution, making MoE more scalable and easier to build.

In this context, NVIDIA NeMo AutoModel stands out as an open library integrated into the NVIDIA NeMo framework, dedicated to creating large-scale custom generative AI models. Building on the innovations of v5, NeMo AutoModel introduces significant improvements, including expert parallelism, merged all-to-all DeepEP dispatch, and TransformerEngine kernels. These optimizations enable a training throughput 3.4 to 3.7 times higher and a 29 to 32% reduction in GPU memory when fine-tuning MoE models compared to native Transformers v5, all while using the same from_pretrained() API.

This blog details how this combination works and how users can fine-tune MoE models more quickly without changing their APIs.

Challenges and Solutions for Training MoE Models

The growing adoption of MoE models has introduced new challenges for the efficient training of AI models. Among these challenges are routing tokens through hundreds of experts, merging expert matrix multiplications into a single kernel, partitioning weights across GPUs, and overlapping communication with computation. These requirements necessitate infrastructure far beyond what a general-purpose library can offer by default.

Transformers v5 brought first-class support for MoE, introducing expert backends, dynamic weight loading, and tensor parallelism plans for distributed execution. Additionally, v5 integrated PyTorch's DeviceMesh directly into the from_pretrained() function, facilitating distributed training.

NeMo AutoModel builds on these foundations by subclassing AutoModelForCausalLM and adding features such as expert parallelism (EP), merged all-to-all DeepEP dispatch, and TransformerEngine kernels. DeepEP, in particular, overlaps communication with expert computation, optimizing training efficiency. Thanks to the reversible weight conversion from v5, NeMo AutoModel can focus on these reusable core operations while allowing save_pretrained() to produce checkpoints compatible with tools like vLLM and SGLang.

The next section explains how the two work together and the performance gains we measured, ranging from the full fine-tuning of the NVIDIA Nemotron 3 Ultra 550B A55B on 16 nodes to single-node models like Qwen3-30B-A3B and Nemotron 3 Nano 30B A3B.

API Compatibility and Enhanced Performance

One of the main goals of NeMo AutoModel is to maintain compatibility with the HuggingFace Transformers API, allowing the open-source community to easily adopt it. By subclassing AutoModelForCausalLM, NeMo AutoModel ensures that any code working with HF models is also compatible with AutoModel.

Loading a model with NeMo AutoModel is straightforward and requires only a minor change to the import statement. This single import activates a series of optimizations for popular MoE architectures such as Qwen3, NVIDIA Nemotron, GPT-OSS, and DeepSeek V3. NeMo AutoModel provides optimized implementations with TransformerEngine attention, fused linear layers, and custom expert kernels. For other models, it reverts to HF standards while applying optimizations like Liger kernel patching.

Where NeMo AutoModel truly excels is in scaling MoE models for multi-GPU training. For example, to train the Nemotron 3 Nano 30B A3B model with expert parallelism on 8 GPUs, it is sufficient to add a specific distributed mesh configuration. This offers speed, scalability, and memory optimizations with FSDP2, expert parallelism, TransformerEngine kernels, and DeepEP dispatch, all from a simple from_pretrained() call.

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",
    model = NeMoAutoModelForCausalLM.from_pretrained(
        "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16",
        dtype=torch.bfloat16,
        distributed_setup=dist_setup,
    )
)

dist.destroy_process_group()

Comparative Performance Analysis

To evaluate the effectiveness of NeMo AutoModel, two scenarios were tested: the full fine-tuning of a 550 billion parameter model on 16 nodes, and the training of two 30 billion parameter MoE models on a single node. The Nemotron 3 Ultra 550B A55B model, a hybrid with 550 billion parameters, was used for the full fine-tuning benchmark. This model, equipped with Mamba2, LatentMoE, and Multi-Token Prediction (MTP), was tested on 16 H100 nodes (128 GPUs).

In this test, expert parallelism allowed the distribution of experts across GPUs, thereby reducing memory footprint and enabling the full fine-tuning to execute efficiently. Tests on the 30 billion parameter models showed similar gains, where NeMo AutoModel outperformed Transformers v5 in terms of throughput and memory efficiency.

Nemotron 3 Ultra 550B A55B (Full Fine-Tuning, Multi-Node)

The Nemotron 3 Ultra 550B A55B model is a hybrid model with 550 billion parameters delivered with Mamba2, LatentMoE, and Multi-Token Prediction (MTP). We conducted a full fine-tuning benchmark: each parameter is updated, and the state of the Adam optimizer is materialized, which at this scale spans 16 H100 nodes (128 GPUs).

• 16x H100 80GB (128 GPUs)
• Expert parallelism
• Local batch size
• MTP, activation checkpointing, fused linear cross-entropy
• DeepEP dispatch + torch_mm experts + TransformerEngine

Why is there no column for Transformers v5? Transformers v5 runs out of memory at this scale, so there is no v5 figure to report here. AutoModel's expert parallelism distributes the experts across GPUs to bring the footprint within limits, allowing the full fine-tuning to execute. The comparisons for 30 billion below show the same advantage where v5 fits.

30 Billion Single-Node MoE Benchmarks

We benchmarked three approaches on a single node with 8x H100 80GB GPUs: HF Transformers v4 (hub code), HF Transformers v5 (with the best available optimizations), and NeMo AutoModel (EP=8 + custom kernels).

• 8x H100 80GB (single node)
• Local batch size

A note on the routing gate. The NeMo AutoModel figures below use a balanced routing gate, which forces tokens to be distributed evenly among the experts. This emulates the ideal operating point toward which an MoE is trained: the load balancing loss of a well-trained model pushes expert utilization toward near-perfect uniformity, so balanced routing reflects the stable state toward which a real workload converges (and eliminates the noise from stragglers that random tokens might inject into expert parallelism). The v4/v5 columns run their native router on the same dummy tokens. Thus, the balanced gate measures NeMo AutoModel at its target MoE operating point, and the v4/v5 columns reflect their default behavior.

Why the Acceleration Comes From Here

The 3.4 to 3.7 times acceleration of NeMo AutoModel compared to Transformers v5 comes from three sources:

• Expert parallelism reduces memory pressure. EP=8 distributes expert weights across GPUs, reducing the MoE footprint per GPU by 8x. For Qwen3, this drops the maximum memory from 68.2 GiB to 48.1 GiB (-29%). For Nemotron Nano, it goes from 62.1 GiB to 42.5 GiB (-32%), freeing up headroom for larger batch sizes or longer sequences.

• DeepEP merges communication and computation. Instead of separate AllGather/ReduceScatter collects for expert routing, DeepEP merges token dispatch and combines into optimized GPU kernels, overlapping communication with expert computation.

• TransformerEngine kernels accelerate core operations. The fused attention, linear layer, and RMSNorm implementations from TE offer consistent speedups over their PyTorch/Flash Attention equivalents across all layer types, not just MoE layers.

Features of Transformers v5 Leveraged by HuggingFace AutoModel

One of the most impactful features of Transformers v5 is the experts_implementation parameter, which includes three expert backends:

• For loop over selected experts
• Debugging, compatibility, and correction. Also available for v4.
• Duplicates expert parameters, single grouped GEMM via torch.bmm
• Small inputs, fast with torch.compile. Added for v5.
• Sorts tokens by expert, single grouped GEMM via torch.nn.functional.grouped_mm
• Training (memory-efficient, no parameter duplication). Added for v5.

The grouped_mm backend is the key optimization for training: 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 goes further. For models with custom implementations, it uses the merged all-to-all dispatch DeepEP combined with grouped GMM kernels and TransformerEngine linear layers. The progression is as follows:

• v4 (impatient for loop) → v5 (grouped_mm) → NeMo AutoModel (DeepEP + GMM + TE)

In NeMo AutoModel, the expert backend is configured via BackendConfig:

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

backend = BackendConfig(
    experts="torch_mm",
    dispatcher="deepep",
)

In summary, NVIDIA NeMo AutoModel represents a major advancement in fine-tuning Transformers models, offering enhanced performance and API compatibility that facilitates its adoption by the AI community.