Nemotron Elastic: Rumo a LLMs de Raciocínio Eficientes com Múltiplas Capacidades

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

November 20, 2025

Autores: Ali Taghibakhshi, Sharath Turuvekere Sreenivas, Saurav Muralidharan, Ruisi Cai, Marcin Chochowski, Ameya Sunil Mahabaleshwarkar, Yoshi Suhara, Oluwatobi Olabiyi, Daniel Korzekwa, Mostofa Patwary, Mohammad Shoeybi, Jan Kautz, Bryan Catanzaro, Ashwath Aithal, Nima Tajbakhsh, Pavlo Molchanov

cs.AI

Resumo

Treinar uma família de modelos de linguagem de grande escala direcionada a múltiplas escalas e objetivos de implantação é proibitivamente caro, exigindo execuções de treinamento separadas para cada tamanho diferente. Trabalhos recentes sobre compressão de modelos por meio de poda e destilação de conhecimento reduziram esse custo; no entanto, esse processo ainda incorre em custos de treinamento equivalentes a centenas de bilhões de tokens por modelo comprimido. Neste artigo, apresentamos o Nemotron Elastic, uma estrutura para construir LLMs orientados para raciocínio, incluindo arquiteturas híbridas Mamba-Atenção, que incorporam múltiplos submodelos aninhados dentro de um único modelo pai, cada um otimizado para diferentes configurações de implantação e orçamentos. Cada um desses submodelos compartilha pesos com o modelo pai e pode ser extraído *zero-shot* durante a implantação sem treinamento ou ajuste fino adicionais. Habilitamos essa funcionalidade por meio de um roteador treinado de ponta a ponta, fortemente acoplado a um currículo de treinamento em dois estágios projetado especificamente para modelos de raciocínio. Adicionalmente, introduzimos a elastificação SSM com consciência de grupo, que preserva as restrições estruturais do Mamba, a elastificação heterogênea de MLP, a importância de camada baseada em MSE normalizado para uma seleção de profundidade aprimorada, e a destilação de conhecimento que permite a otimização simultânea de múltiplos orçamentos. Aplicamos o Nemotron Elastic ao modelo Nemotron Nano V2 12B, produzindo simultaneamente um modelo de 9B e um de 6B usando apenas 110B de tokens de treinamento; isso resulta em uma redução de custo superior a 360x em comparação com o treinamento de famílias de modelos do zero, e cerca de 7x em comparação com as técnicas de compressão SoTA. Cada um dos modelos aninhados apresenta desempenho igual ou superior ao SoTA em precisão. Além disso, ao contrário de outros métodos de compressão, a capacidade aninhada da nossa abordagem permite ter um modelo de raciocínio múltiplo-em-um que possui memória de implantação constante em relação ao número de modelos na família.

English

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.