Nemotron Elastic: Towards Efficient Many-in-One Reasoning LLMs

TL;DR

Nemotron Elastic introduces a novel framework for creating multi-scale, reasoning-oriented large language models with nested submodels that can be extracted zero-shot, significantly reducing training costs and maintaining high accuracy.

Contribution

The paper presents Nemotron Elastic, a method for embedding multiple nested submodels within a single LLM, enabling multi-budget deployment without additional training.

Findings

Achieved over 360x cost reduction compared to training from scratch.

Produced multiple models (9B and 6B) from a 12B model using only 110B tokens.

Nested models perform on par or better than state-of-the-art in accuracy.

Abstract

Training a family of large language models targeting multiple scales and deployment objectives is prohibitively expensive, requiring separate training runs for each different size. Recent work on model compression through pruning and knowledge distillation has reduced this cost; however, this process still incurs hundreds of billions of tokens worth of training cost per compressed model. In this paper, we present Nemotron Elastic, a framework for building reasoning-oriented LLMs, including hybrid Mamba-Attention architectures, that embed multiple nested submodels within a single parent model, each optimized for different deployment configurations and budgets. Each of these submodels shares weights with the parent model and can be extracted zero-shot during deployment without additional training or fine-tuning. We enable this functionality through an end-to-end trained router, tightly coupled to a two-stage training curriculum designed specifically for reasoning models. We additionally introduce group-aware SSM elastification that preserves Mamba's structural constraints, heterogeneous MLP elastification, normalized MSE-based layer importance for improved depth selection, and knowledge distillation enabling simultaneous multi-budget optimization. We apply Nemotron Elastic to the Nemotron Nano V2 12B model, simultaneously producing a 9B and a 6B model using only 110B training tokens; this results in over 360x cost reduction compared to training model families from scratch, and around 7x compared to SoTA compression techniques. Each of the nested models performs on par or better than the SoTA in accuracy. Moreover, unlike other compression methods, the nested capability of our approach allows having a many-in-one reasoning model that has constant deployment memory against the number of models in the family.