PaperlayerView exact revision on arXiv

Source-linked AI summary

Nemotron-H: A Family of Accurate and Efficient Hybrid Mamba-Transformer Models

NVIDIA, :, Aaron Blakeman, Aarti Basant, Abhinav Khattar, Adithya Renduchintala, Akhiad Bercovich, Aleksander Ficek, Alexis Bjorlin, Ali Taghibakhshi, Amala Sanjay Deshmukh, Ameya Sunil Mahabaleshwarkar, Andrew Tao, Anna Shors, Ashwath Aithal, Ashwin Poojary, Ayush Dattagupta, Balaram Buddharaju, Bobby Chen, Boris Ginsburg, Boxin Wang, Brandon Norick, Brian Butterfield, Bryan Catanzaro, Carlo del Mundo, Chengyu Dong, Christine Harvey, Christopher Parisien, Dan Su, Daniel Korzekwa, Danny Yin, Daria Gitman, David Mosallanezhad, Deepak Narayanan, Denys Fridman, Dima Rekesh, Ding Ma, Dmytro Pykhtar, Dong Ahn, Duncan Riach, Dusan Stosic, Eileen Long, Elad Segal, Ellie Evans, Eric Chung, Erick Galinkin, Evelina Bakhturina, Ewa Dobrowolska, Fei Jia, Fuxiao Liu, Gargi Prasad, Gerald Shen, Guilin Liu, Guo Chen, Haifeng Qian, Helen Ngo, Hongbin Liu, Hui Li, Igor Gitman, Ilia Karmanov, Ivan Moshkov, Izik Golan, Jan Kautz, Jane Polak Scowcroft, Jared Casper, Jarno Seppanen, Jason Lu, Jason Sewall, Jiaqi Zeng, Jiaxuan You, Jimmy Zhang, Jing Zhang, Jining Huang, Jinze Xue, Jocelyn Huang, Joey Conway, John Kamalu, Jon Barker, Jonathan Cohen, Joseph Jennings, Jupinder Parmar, Karan Sapra, Kari Briski, Kateryna Chumachenko, Katherine Luna, Keshav Santhanam, Kezhi Kong, Kirthi Sivamani, Krzysztof Pawelec, Kumar Anik, Kunlun Li, Lawrence McAfee, Leon Derczynski, Lindsey Pavao, Luis Vega, Lukas Voegtle, Maciej Bala, Maer Rodrigues de Melo, Makesh Narsimhan Sreedhar, Marcin Chochowski, Markus Kliegl, Marta Stepniewska-Dziubinska, Matthieu Le, Matvei Novikov, Mehrzad Samadi, Michael Andersch, Michael Evans, Miguel Martinez, Mike Chrzanowski, Mike Ranzinger, Mikolaj Blaz, Misha Smelyanskiy, Mohamed Fawzy, Mohammad Shoeybi, Mostofa Patwary, Nayeon Lee, Nima Tajbakhsh, Ning Xu, Oleg Rybakov, Oleksii Kuchaiev, Olivier Delalleau, Osvald Nitski, Parth Chadha, Pasha Shamis, Paulius Micikevicius, Pavlo Molchanov, Peter Dykas, Philipp Fischer, Pierre-Yves Aquilanti, Piotr Bialecki, Prasoon Varshney, Pritam Gundecha, Przemek Tredak, Rabeeh Karimi, Rahul Kandu, Ran El-Yaniv, Raviraj Joshi, Roger Waleffe, Ruoxi Zhang, Sabrina Kavanaugh, Sahil Jain, Samuel Kriman, Sangkug Lym, Sanjeev Satheesh, Saurav Muralidharan, Sean Narenthiran, Selvaraj Anandaraj, Seonmyeong Bak, Sergey Kashirsky, Seungju Han, Shantanu Acharya, Shaona Ghosh, Sharath Turuvekere Sreenivas, Sharon Clay, Shelby Thomas, Shrimai Prabhumoye, Shubham Pachori, Shubham Toshniwal, Shyamala Prayaga, Siddhartha Jain, Sirshak Das, Slawek Kierat, Somshubra Majumdar, Song Han, Soumye Singhal, Sriharsha Niverty, Stefania Alborghetti, Suseella Panguluri, Swetha Bhendigeri, Syeda Nahida Akter, Szymon Migacz, Tal Shiri, Terry Kong, Timo Roman, Tomer Ronen, Trisha Saar, Tugrul Konuk, Tuomas Rintamaki, Tyler Poon, Ushnish De, Vahid Noroozi, Varun Singh, Vijay Korthikanti, Vitaly Kurin, Wasi Uddin Ahmad, Wei Du, Wei Ping, Wenliang Dai, Wonmin Byeon, Xiaowei Ren, Yao Xu, Yejin Choi, Yian Zhang, Ying Lin, Yoshi Suhara, Zhiding Yu, Zhiqi Li, Zhiyu Li, Zhongbo Zhu, Zhuolin Yang, Zijia Chen

arXiv:2504.03624v4cs.CLcs.AIcs.LG

TL;DR

Nemotron-H targets the inference cost of self-attention as reasoning models generate longer outputs. It introduces hybrid Mamba-Transformer models, MiniPuzzle compression, and FP8 training, achieving Transformer-comparable capabilities with improved inference efficiency.

  • Problem

    Self-attention requires computation and KV-cache memory that scale linearly with sequence length, making efficient inference important for inference-time scaling.

  • Method

    Nemotron-H replaces most self-attention layers with Mamba-2 in hybrid models, and applies MiniPuzzle pruning-distillation compression plus an FP8-based pre-training recipe.

  • Results

    Nemotron-H achieves either better or on-par accuracy than similarly sized open-source Transformer models while being up to 3× faster at inference.

  • Takeaways & Limitations

    Hybrid architectures, FP8 pre-training, and MiniPuzzle can support capable models with improved inference speed and lower model-creation cost.

  • Takeaways & Limitations

    FP8 recipe results with smaller models did not generalize, so verification on at least 8B parameters was important.

Abstract

from arXiv · show

As inference-time scaling becomes critical for enhanced reasoning capabilities, it is increasingly becoming important to build models that are efficient to infer. We introduce Nemotron-H, a family of 8B and 56B/47B hybrid Mamba-Transformer models designed to reduce inference cost for a given accuracy level. To achieve this goal, we replace the majority of self-attention layers in the common Transformer model architecture with Mamba layers that perform constant computation and require constant memory per generated token. We show that Nemotron-H models offer either better or on-par accuracy compared to other similarly-sized state-of-the-art open-sourced Transformer models (e.g., Qwen-2.5-7B/72B and Llama-3.1-8B/70B), while being up to 3$\times$ faster at inference. To further increase inference speed and reduce the memory required at inference time, we created Nemotron-H-47B-Base from the 56B model using a new compression via pruning and distillation technique called MiniPuzzle. Nemotron-H-47B-Base achieves similar accuracy to the 56B model, but is 20% faster to infer. In addition, we introduce an FP8-based training recipe and show that it can achieve on par results with BF16-based training. This recipe is used to train the 56B model. We are releasing Nemotron-H base model checkpoints with support in Hugging Face and NeMo.

1. Introduction

Nemotron-H addresses the inference cost of attention during inference-time scaling by using hybrid Mamba-Transformer architectures. The models target Transformer-level capabilities with faster inference, while MiniPuzzle compression and FP8 training further improve deployment and creation efficiency.

  • Motivation: Self-attention computation and KV-cache memory scale linearly with sequence length, which can limit inference-time scaling for reasoning models.Mamba-based alternatives can provide sub-linear or constant compute and memory requirements.
  • Approach: Nemotron-H replaces most Transformer self-attention with a mixture of Mamba-2, self-attention, and MLP layers.The family includes 8B and 47B/56B model series.
  • Results: Nemotron-H offers either better or on-par accuracy than similarly sized Qwen-2.5 and Llama-3.1 models, with up to 3× higher inference throughput.Nemotron-H-56B-Base outperforms Llama-3.1-70B on 16 out of 17 evaluated tasks; the throughput result uses 65536 input and 1024 output tokens on NVIDIA H100 GPUs.
  • Training: FP8 per-tensor scaling reaches equal or better downstream task accuracy than BF16 in 8B ablations, while Nemotron-H-56B-Base is fully pretrained with the FP8 recipe.The recipe preserves tensor maxima and leaves the first and last 4 GEMMs in BF16 for stability and convergence.
  • Compression: MiniPuzzle compresses Nemotron-H-56B-Base to Nemotron-H-47B-Base through pruning and distillation for deployment under memory and latency constraints.The compressed model is intended for FP4 deployment on a single 32GiB NVIDIA RTX 5090 GPU.
  • Extensions and release: The paper reports that Nemotron-H base models can be post-trained into effective vision-language, instruction-following, and long-context models.The released base checkpoints are provided in Hugging Face and NeMo formats.

2. Base Models and Pre-Training

This section compares Nemotron-H-8B/56B-Base with open-source state-of-the-art Transformer models on base-model accuracy and inference speed.

  • Nemotron-H-8B/56B-Base are evaluated against existing open-source state-of-the-art Transformer models for benchmark accuracy and inference speed.

2.1. Model Architecture

Nemotron-H architectures combine mostly Mamba-2 and FFN layers with a small, evenly distributed set of self-attention layers. The design preserves selected Transformer ordering constraints while using grouped-query attention and model-size-specific dimensions.

  • Nemotron-H models combine Mamba-2, self-attention, and FFN layers in a hybrid architecture.
  • Roughly 8% of layers are self-attention layers, evenly dispersed through the model: 4 of 52 in Nemotron-H-8B and 10 of 118 in Nemotron-H-56B.
  • The remaining layers are evenly split between FFN and Mamba-2 layers, with Mamba-2 first, FFN last, and self-attention preceding FFN layers.
  • Nemotron-H-8B uses hidden dimension 4096, FFN dimension 21504, 32 query heads, and Mamba-2 state dimension 128; Nemotron-H-56B uses 8192, 32768, 64, and 256 respectively.
  • Both models use grouped-query attention with 8 key-value heads, 8 Mamba-2 groups, squared ReLU FFN activation, and no position embeddings.
  • The Nemotron-T-8B baseline uses 32 Transformer layers, each consisting of self-attention followed by an FFN layer.

2.2. Pre-Training Data

Nemotron-H pre-training uses curated web, math, code, and academic data alongside synthetic data generation to improve quality, diversity, and long-horizon training capacity.

  • Nemotron-H uses separate curation pipelines for web crawl, math, code, and academic data.
  • Web crawl data: 6.3 trillion English Common Crawl tokens include 4.4 trillion globally deduplicated real tokens and 1.9 trillion rephrased synthetic tokens.
  • Web crawl data: A high-quality subset increases MMLU by 5.6 points over DCLM, while the full dataset provides comparable accuracy with 4× more unique tokens.
  • Academic data: Academic data is weighted using classifiers for educational content, difficulty, and domain, with highly educational high-school and undergraduate material most helpful.
  • Synthetically-Generated Data: Synthetic generation rephrases low-quality web data, creates factual transformations of high-quality documents, and produces over 1.8 trillion synthetic tokens.
  • Math data: Expanding OpenWebMath from 14 billion to over 100 billion tokens yielded an 18-point improvement on GSM8K in controlled experiments.
  • Code data: Synthetic code data combines generated programming problems, solutions, and explanations across 11 programming languages after minimal validity filtering.

2.3. FP8 Recipe

Nemotron-H uses mixed-precision FP8 pre-training with per-tensor dynamic quantization, and reports downstream accuracy that matches or exceeds BF16 despite small loss differences.

  • Nemotron-H-56B-Base computes most linear layers in FP8 while keeping the first and last four layers in BF16.
  • Per-tensor dynamic quantization scales each input tensor from its maximum absolute value before casting it to the target FP8 format.
  • FP8 loss was typically less than 0.1% higher than BF16, decreasing during training before widening in the final quarter.
  • FP8-trained models achieved downstream evaluations as good as or better than BF16 experiments, despite marginally worse loss curves.
  • FP8 experiments using per-tensor current scaling were often substantially better than BF16 on coding and math tasks.

2.4. Pre-Training

Nemotron-H models were trained for 15- and 20-trillion-token horizons using large-scale distributed infrastructure and resilience measures that reduce downtime and recovery costs.

  • Nemotron-H-8B-Base was trained for 15 trillion tokens, while Nemotron-H-56B-Base was trained for 20 trillion tokens.
  • Training used sequence length 8192, global batch size 768, cosine learning-rate decay, and weight decay of 0.1.
  • Pre-training used Megatron-LM, Transformer Engine FP8 support, 8-way tensor parallelism, sequence parallelism, and 768-way data parallelism.
  • Training resilience: The training system addresses failures from GPUs, shared infrastructure, and application-induced errors by attributing causes and excluding faulty hardware from subsequent jobs.
  • Training resilience: Asynchronous checkpointing and optimized restart operations reduce work lost and recovery overhead after failures.
  • Nemotron-H-56B-Base was trained efficiently on 6144 NVIDIA H100 GPUs.

2.5. Base Model Evaluations

Across common benchmarks, Nemotron-H base models are comparable or better than similarly sized Transformers, while the 8B hybrid model also matches a same-data Transformer and the 56B model is competitive with larger models.

  • Across both model sizes, Nemotron-H base models reach comparable or better accuracy than similarly sized state-of-the-art Transformer models.
  • 56B evaluations: Nemotron-H-56B-Base has the highest accuracy among three 56B/70B-class models on 9 of 17 tasks and outperforms Llama-3.1-70B-Base on 16 of 17 tasks.
  • 56B evaluations: Nemotron-H-56B-Base outperforms DeepSeek-V3-671B-Base on 4 of 10 overlapping tasks and Llama-3.1-405B-Base on 5 of 10.
  • 8B evaluations: Nemotron-H-8B-Base achieves the highest accuracy among three 8B-class models on 7 of 15 tasks and leads on 4 of 5 commonsense tasks.
  • 8B evaluations: In an apples-to-apples comparison, Nemotron-H-8B-Base exceeds the same-data Transformer on 7 of 15 tasks and is within one point on 4 more.

2.6. Inference Speed

Nemotron-H reduces inference cost by replacing most self-attention layers with more efficient Mamba layers. Under long-context settings, the models achieve substantially higher throughput than similarly sized Transformer baselines.

  • MMLU-Pro accuracy is plotted against inference throughput to compare Nemotron-H-56B-Base and Nemotron-H-8B-Base with similarly sized Transformer models.Experiments use long input contexts, 1024 generated tokens, and throughput normalized by GPU count.
  • 2.4× more output tokens per second per GPU than Qwen-2.5-72B and Llama-3.1-70B is achieved by Nemotron-H-56B-Base.Against Llama-3.1-405B, it achieves 19.6× higher throughput.
  • 2.9× more output tokens per second per GPU is achieved after distillation to Nemotron-H-47B-Base.

3. Compression and Distillation

MiniPuzzle compresses Nemotron-H-56B-Base through importance-guided pruning, conditional architecture search, and knowledge distillation. The resulting 47B model meets stringent memory constraints while retaining near-original accuracy and improving inference speed.

  • 3. Compression and Distillation: MiniPuzzle combines importance estimation, conditional NAS, and knowledge distillation to produce compressed models satisfying deployment constraints.
  • 3.2. Importance Estimation: Importance estimation scores layers and FFN neurons using lightweight forward-pass analyses to determine which components can be pruned.Layer importance uses activation MSE, while FFN neuron importance aggregates outputs over calibration data.
  • 3.3. Conditional NAS: MiniPuzzle prunes layers and FFN hidden dimensions while retaining architectures below the 31.7 GiB FP4 memory constraint.The search enumerates candidate layer counts and FFN widths, producing around 400 retained architectures.
  • 3.4. Knowledge Distillation: Logit-based distillation with forward KL divergence is used to recover accuracy lost during pruning.
  • 3.5. Results: Nemotron-H-47B-Base achieves 1.2× faster inference and near-lossless benchmark accuracy than the 56B model.It is designed for long-context inference on a 32GiB NVIDIA GeForce RTX 5090 GPU and uses roughly 300× fewer training tokens than training from scratch.
  • 3.5. Results: The selected 47B architecture retains 5 self-attention, 44 Mamba-2, and 49 FFN layers, with FFN width reduced from 32768 to 30720.

4. Vision-Language Models

Nemotron-H backbones are adapted into vision-language models using a vision encoder, projector, and decoder-only language model architecture. The resulting models achieve strong or state-of-the-art results across broad vision-language benchmarks.

  • 4.1. Model Architecture: Nemotron-H-VLM comprises a vision encoder, a two-layer FFN projector, and a Nemotron-H LLM backbone.
  • 4.1. Model Architecture: InternViT-300M-V2.5 processes tiled 448×448 images and produces 1024 visual tokens per tile, with a thumbnail preserving global context.
  • 4.1. Model Architecture: The VLM pipeline downsamples 1024 image tokens to 256 by grouping neighboring tokens and adds interleaved tile ID tags.The paper reports that tile ID tags are crucial for accuracy on various vision-language benchmarks.
  • 4.2. Training: VLM training uses frozen-backbone modality alignment during VL pre-training, followed by end-to-end supervised fine-tuning.
  • 4.3. Vision-Language Benchmark Results: Nemotron-H-56B-VLM achieves state-of-the-art results and superior quality compared with NVLM-1.0-D and a Qwen2.5-72B-Instruct VLM.
  • 4.3. Vision-Language Benchmark Results: Nemotron-H-8B-VLM is reported as a strong backbone for developing best-in-class VLMs against Llama-3.1-8B-Instruct and Qwen2.5-7B-Instruct.The comparison uses the same training methodology and dataset.

5. Nemotron-H Reasoning Models

The reasoning-model study finds that hybrid and Transformer models have similar post-training properties, then develops Nemotron-H reasoning variants through staged supervised and preference or reinforcement optimization. Extended-context training enables strong RULER performance, while evaluated reasoning models match or exceed cited baselines and support quantized deployment.

  • 5.1. Post-training hybrid models: Similar pretraining performance motivated comparing Transformer and hybrid models’ post-training behavior before training Nemotron-H reasoning models.The pilot used Transformer and hybrid models pretrained on identical data, and the study reports similar pretraining performance.
  • 5.1. Post-training hybrid models: A three-stage post-training procedure combines supervised fine-tuning, offline preference tuning, and on-policy preference tuning with safety data.The stages cover code, math, general instruction following, long-context samples, narrow instruction following, and safety alignment.
  • 5.2. Long-context training: Long-context data augments concatenated conversations with generated references to earlier turns and semantically related segments placed later in samples.These additions are designed to improve long-range memory and include references to prior conversation turns and scattered themed segments.
  • 5.2. Long-context training: Hybrid models learn higher context lengths within a few hundred training steps and achieve an 83% RULER score at 128k sequence length.The result follows training on extended blends containing samples up to 512k tokens.
  • 5.3. Reasoning-model results: Across math, coding, science, tool-use, and dialogue benchmarks, Nemotron-H-47B-Reasoning-128K matches or exceeds Llama-Nemotron Super 49B V1.0 and beats Qwen3 32B on all non-coding benchmarks.The model also supports post-training quantization of all linear layers with minimal accuracy loss.
  • 5.3. Reasoning-model results: Nemotron-H-8B-Reasoning-128K matches or exceeds Llama-Nemotron Nano 8B V1.0 across AIME25, science, tool-use, and instruction-following benchmarks.An FP8 quantized checkpoint is provided, with minimal accuracy loss reported for quantized deployment.

6. Conclusions

The conclusions present Nemotron-H as a hybrid architecture that preserves capabilities comparable to state-of-the-art Transformer models while improving inference speed. The authors also report successful adaptation to vision-language and instruction-tuned models, with FP8 training and MiniPuzzle reducing model-creation costs, and release base and reasoning checkpoints.

  • 6. Conclusions: Hybrid Nemotron-H models offer capabilities comparable to state-of-the-art Transformer models while improving inference speed.The conclusion frames this as combining Transformer-like capabilities with more efficient inference.
  • 6. Conclusions: Nemotron-H base models can be adapted into effective vision-language and instruction-tuned models.The released family includes base checkpoints for 8B, 47B, and 56B variants.
  • 6. Conclusions: FP8-based training and MiniPuzzle can reduce the cost of creating Nemotron-H models.The authors describe both techniques as part of the family’s practical development and deployment approach.
  • 6. Conclusions: Nemotron-H base checkpoints are released with Hugging Face and NeMo support, while 47B and 8B reasoning checkpoints are available through Hugging Face.The listed base releases include Nemotron-H-56B-Base, Nemotron-H-47B-Base, and Nemotron-H-8B-Base.
Loading 2504.03624v4…