Accelerating Transformers Fine-Tuning with NVIDIA NeMo AutoModel

NVIDIA released NeMo AutoModel, an open-source library that accelerates fine-tuning of Mixture-of-Experts (MoE) transformer models by 3.4-3.7x in training throughput and reduces GPU memory usage by 29-32% compared to native HuggingFace Transformers v5, while maintaining API compatibility. The library builds on Transformers v5's MoE support, adding Expert Parallelism, DeepEP fused all-to-all dispatch, and TransformerEngine kernels to optimize distributed training for models like NVIDIA Nemotron and Qwen3.

Text Generation • 32B • Updated • 1.2M • 773 Accelerating Transformers Fine-Tuning with NVIDIA NeMo AutoModel Enterprise + Article /blog Transformers v5 https://github.com/huggingface/transformers/releases/tag/v5.0.0 release strengthened it with first-class support for Mixture-of-Experts MoE models, now the dominant architecture for frontier models https://www.nvidia.com/en-us/glossary/frontier-models/ . v5 ships the MoE foundations: expert backends, dynamic weight loading, and distributed execution that make MoE extensible and easy to build on. NVIDIA NeMo AutoModel https://github.com/NVIDIA-NeMo/Automodel is an open library part of the NVIDIA NeMo framework https://github.com/NVIDIA-NeMo for building custom generative AI models at scale. NeMo AutoModel builds cleanly on top of v5, adding Expert Parallelism, DeepEP fused all-to-all dispatch, and TransformerEngine kernels, and it leans on v5's dynamic weight loading to bring those optimizations to a broad and growing set of model families. The payoff is 3.4-3.7x higher training throughput and 29-32% less GPU memory on fine-tuning MoE models than native Transformers v5, using the same from pretrained API: a single import line, with no other code changes. This blog details how this combination works and how users can fine-tune MoE models faster without changing their APIs. Background The rise of MoE models has introduced new challenges to efficient training: Routing tokens across hundreds of experts, fusing expert matmuls into a single kernel, sharding weights across GPUs, and overlapping communication with computation all require infrastructure beyond what a general-purpose library provides out of the box. Transformers v5 https://github.com/huggingface/transformers/releases/tag/v5.0.0 “v5” introduced first-class MoE support such as expert backends https://huggingface.co/docs/transformers/en/experts interface , dynamic weight loading https://huggingface.co/docs/transformers/en/weightconverter , and tensor parallel plans for distributed execution. In addition, v5 made distributed training first-class by integrating PyTorch's DeviceMesh directly into from pretrained . NeMo AutoModel https://github.com/NVIDIA-NeMo/Automodel builds on top of v5 by subclassing AutoModelForCausalLM, and adding Expert Parallelism EP , DeepEP fused all-to-all dispatch, and TransformerEngine kernels. DeepEP is the piece v5 doesn't have yet: it overlaps communication with expert compute. And because NeMo AutoModel rides v5's reversible weight conversion to load each model, it can focus its engineering on these reusable core ops instead of per-model checkpoint plumbing, while save pretrained still emits standard HF checkpoints that tools like vLLM and SGLang can load. The next section walks through how the two work together and the performance gains we measured, from full fine-tuning NVIDIA Nemotron 3 Ultra 550B A55B https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16 across 16 nodes down to single-node models such as Qwen3-30B-A3B and Nemotron 3 Nano 30B A3B https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 . NeMo AutoModel: Same API, More Performance One of NeMo AutoModel's goals is API compatibility with HuggingFace Transformers to enable open-source community. NeMoAutoModelForCausalLM subclasses AutoModelForCausalLM, so any code that works with HF models works with AutoModel too. Here's what loading a model looks like in both. Only the import changes: That single import does a lot of work. For popular MoE architectures like Qwen3, NVIDIA Nemotron https://developer.nvidia.com/nemotron , GPT-OSS, and DeepSeek V3, NeMo AutoModel ships hand-tuned implementations https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo automodel/ transformers/registry.py with TransformerEngine attention, fused linear layers, and custom expert kernels. For everything else, it falls back to vanilla HF while still applying optimizations like Liger kernel https://github.com/linkedin/Liger-Kernel patching, among others. And whichever path it takes, the resulting model is ready to scale: pass a device mesh and you have multi-GPU training without further rewrites. Where NeMo AutoModel really shines is scaling MoE models to multi-GPU training. To train Nemotron 3 Nano 30B A3B https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16 with Expert Parallelism across 8 GPUs, one adds the distributed mesh configuration: 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 This gives speed, scalability and memory-optimizations with FSDP2, Expert Parallelism, TransformerEngine kernels and DeepEP dispatch, all from a from pretrained call. Performance Comparison We evaluated NeMo AutoModel in two regimes: full fine-tuning a frontier-scale 550B model across 16 nodes, and training two 30B MoE models on a single node. The 550B result shows why Expert Parallelism is essential at scale; the 30B results quantify the per-GPU speedup over Transformers v5. Nemotron 3 Ultra 550B A55B full fine-tune, multi-node Nemotron 3 Ultra 550B A55B https://huggingface.co/nvidia/NVIDIA-Nemotron-3-Ultra-550B-A55B-BF16 is a 550B-parameter hybrid model shipping with Mamba2, LatentMoE, and Multi-Token Prediction MTP . We benchmark a full fine-tune : every parameter is updated and the Adam optimizer state is materialized, which at this scale spans 16 H100 nodes 128 GPUs . Methodology: | 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 | Why there is no Transformers v5 column. Transformers v5 runs out of memory at this scale, so there is no v5 number to report here. AutoModel's Expert Parallelism shards the experts across GPUs to bring the footprint within budget, which is what lets the full fine-tune run. The 30B comparisons below show the same advantage where v5 fits. Single-node 30B MoE benchmarks We benchmarked three approaches on a single node with 8x H100 80GB GPUs: HF Transformers v4 hub code , HF Transformers v5 with best available optimizations , and NeMo AutoModel EP=8 + custom kernels . Methodology: | Parameter | Value | |---|---| | Hardware | 8x H100 80GB single node | | Sequence length | 4,096 | | Local batch size | 1 | A note on the routing gate. The NeMo AutoModel numbers below use a balanced routing gate, which forces tokens to be distributed uniformly across experts. This emulates the ideal operating point an MoE is trained toward: a well-trained model's load-balancing loss drives expert utilization to near-uniform, so balanced routing reflects the steady-state a real workload converges to and removes the straggler noise that random dummy tokens otherwise inject into expert parallelism . v4/v5 run their native router on the same dummy tokens. The balanced gate therefore measures NeMo AutoModel at its target MoE operating point, and the v4/v5 columns reflect their out-of-the-box behavior. 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 | Why v4 deadlocks: Transformers v4 stores Qwen3 MoE experts as a ModuleList of 128 individual MLP modules, each separately FSDP-wrapped. The forward pass uses a data-dependent loop that only iterates experts that received tokens. With different data per rank, different ranks skip different experts, causing mismatched FSDP AllGather/ReduceScatter collectives and an indefinite hang. Transformers v5 fixes this by storing experts as fused 3D parameter tensors no per-expert modules, no per-expert FSDP collectives . 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 config: trust remote code=True NVIDIA's hub modeling code . The hub code's expert loop is FSDP-safe iterates all experts regardless of token assignment , so it doesn't deadlock like Qwen3 v4. Where the speedup comes from The 3.4-3.7x speedup from NeMo AutoModel over Transformers v5 comes from three sources: Expert Parallelism reduces memory pressure. EP=8 distributes expert weights across GPUs, cutting the per-GPU MoE footprint by 8x. 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 Expert Backends One of the most impactful features in Transformers v5 is the experts implementation https://huggingface.co/docs/transformers/en/experts interface 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: 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 and DeepEP Transformers v5 also ships an Expert Parallelism path https://huggingface.co/docs/transformers/en/expert parallelism . It shards expert weights across GPUs. The GroupedGemmParallel https://github.com/huggingface/transformers/blob/v5.10.2/src/transformers/integrations/tensor parallel.py L1078 style loads only each device's local experts, and RouterParallel https://github.com/huggingface/transformers/blob/v5.10.2/src/transformers/integrations/tensor parallel.py L1123 routes tokens and combines results with an all reduce. It's neatly built on v5's existing tensor-parallel machinery. Enabling it makes the model's tp plan return its expert plan https://github.com/huggingface/transformers/blob/v5.10.2/src/transformers/modeling utils.py L1448 , so expert parallelism shares the device budget with data parallelism ep × dp = world size . For the single-node 30B benchmarks here, we found plain data-parallel v5 dp=8, ep=1 to be the fastest v5 configuration, so that's the v5 setup we report. 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 https://github.com/deepseek-ai/DeepEP that fuses the token routing into optimized GPU kernels, and delivers significant speedups when combined with grouped GEMM for grouped expert computation. In our large-scale MoE benchmarks https://github.com/NVIDIA-NeMo/Automodel/discussions/916 , DeepEP + grouped GEMM reduced cost per iteration by 47% on the full DeepSeek V3 671B model compared to all-gather + looped expert baselines. Dynamic Weight Loading Transformers v5 also introduced a dynamic weight loading https://huggingface.co/docs/transformers/en/weightconverter system through WeightConverter and WeightRenaming. This enables MoE checkpoint 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 https://github.com/NVIDIA-NeMo/Automodel/blob/main/nemo automodel/components/checkpoint/conversion mapping.py 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 To try NeMo AutoModel, please visit our official documentation page to get started https://docs.nvidia.com/nemo/automodel/latest/get-started/installation . For more details, see: NeMo AutoModel HuggingFace API Compatibility Guide https://docs.nvidia.com/nemo/automodel/latest/get-started/hf-compatibility NeMo AutoModel Model Coverage https://docs.nvidia.com/nemo/automodel/latest/model-coverage/overview NeMo AutoModel Performance Summary https://docs.nvidia.com/nemo/automodel/latest/performance/performance-summary NeMo AutoModel on HuggingFace https://huggingface.co/docs/transformers/en/community integrations/nemo automodel finetuning Conclusion NVIDIA NeMo AutoModel is the natural next step for HuggingFace users scaling up model training. By building directly on Transformers v5, AutoModel provides a zero-friction upgrade path: change one import line and get a model instance that is more than three times as fast. On Qwen3-30B-A3B and Nemotron 3 Nano 30B-A3B, this delivers 3.4-3.7x higher training throughput with 29-32% less GPU memory compared to the best Transformers v5 configuration. And because true Expert Parallelism shards experts across GPUs, the same path scales up to full fine-tuning a 550B model like Nemotron 3 Ultra across 16 nodes, the regime where Expert Parallelism becomes essential to fit the model in memory. Because NeMo AutoModel checkpoints are standard HF-format safetensors, you can deploy them on inference frameworks like vLLM and SGLang. The code, configs, and benchmark scripts are all available in the NeMo AutoModel repository https://github.com/NVIDIA-NeMo/Automodel/tree/blog/transformers-v5-automodel/blog experiments . Acknowledgements Core contributors to this work, listed alphabetically by last name: Adil Asif, Hemil Desai, Alexandros Koumparoulis, and Huiying Li.