Nemotron Elastic:面向高效多任务推理大语言模型
计算与语言
2025-11-21 v1
摘要
面向多个尺度和部署目标训练大语言模型家族的成本极高,为每个不同的规模需要独立的训练运行。最近的模型压缩通过修剪和知识蒸馏已降低了这一成本;然而,这一过程仍需为每个压缩模型消耗数千亿代币的训练成本。本文提出了 Nemotron Elastic,一个用于构建面向推理的大语言模型的框架,包括混合 Mamba-Attention 架构,通过在单个父模型中嵌入多个嵌套子模型来实现多任务,每个子模型针对不同的部署配置和预算进行优化。这些子模型与父模型共享权重,可在部署时无需额外训练或微调即零-shot 提取。我们通过一个与两个阶段训练课程紧密耦合的 end-to-end 训练路由器来实现此功能。我们还引入了保持 Mamba 结构约束的 group-aware SSM 弹性化、异构 MLP 弹性化、归一化 MSE 基于层重要性的深度选择,以及知识蒸馏,实现同时的多预算优化。我们将 Nemotron Elastic 应用于 Nemotron Nano V2 12B 模型,仅用 110B 训练 token 同时生产 9B 和 6B 模型;这相当于相对于从头训练模型家族的成本降低超过 360 倍,相对于最先进的压缩技术约降低 7 倍。每个嵌套模型在准确性方面都与最先进技术持平或更好。而且,与其他压缩方法不同,本方法的嵌套能力允许拥有一个针对模型家族中模型数量常数部署内存的多任务推理模型。
引用
@article{arxiv.2511.16664,
title = {Nemotron Elastic: Towards Efficient Many-in-One Reasoning LLMs},
author = {Ali Taghibakhshi and Sharath Turuvekere Sreenivas and Saurav Muralidharan and Ruisi Cai and Marcin Chochowski and Ameya Sunil Mahabaleshwarkar and Yoshi Suhara and Oluwatobi Olabiyi and Daniel Korzekwa and Mostofa Patwary and Mohammad Shoeybi and Jan Kautz and Bryan Catanzaro and Ashwath Aithal and Nima Tajbakhsh and Pavlo Molchanov},
journal= {arXiv preprint arXiv:2511.16664},
year = {2025}
}