Nemotron-H: A Family of Accurate and Efficient Hybrid Mamba-Transformer Models
Nemotron-H: A Family of Accurate and Efficient Hybrid Mamba-Transformer Models
作者/机构: NVIDIA
A1 主要贡献
本文针对大型语言模型(LLM)中自注意力机制(self-attention)带来的推理成本问题,即其计算量和KV缓存随序列长度线性增长,从而限制了模型的推理能力和整体智能,提出了一系列解决方案。随着推理时扩展(inference-time scaling)对增强推理能力变得至关重要,构建推理高效的模型日益重要。
核心问题与研究目标:
传统的自注意力机制在生成每个token时,其计算和内存需求都与序列长度成线性关系,这从根本上限制了模型的性能,尤其是在需要生成更多token以提高响应质量的推理场景下。为了解决这一瓶颈,本文旨在构建在特定准确度水平下能降低推理成本的模型。
创新点与主要贡献:
1. 提出Nemotron-H模型家族:本文引入了一系列混合Mamba-Transformer模型,命名为Nemotron-H,包括8B、56B和47B参数规模。这些模型通过将大部分自注意力层替换为Mamba-2层,实现了每个生成token的恒定计算和内存消耗,从而大幅提升了推理效率。
2. 实现SOTA性能与推理加速:实验证明,Nemotron-H模型在准确性上与同等规模的开源Transformer模型(如Qwen-2.5-7B/72B和Llama-3.1-8B/70B)相当甚至更优,而在长序列(65536输入,1024输出)推理时,其吞吐量在NVIDIA H100 GPU上最高可达3倍。例如,Nemotron-H-56B-Base在评估的17个任务中,有16个优于Llama-3.1-70B。
3. 引入FP8预训练方法:本文提出了一种名为“per-tensor current scaling”的基于FP8的训练方法,并成功用于Nemotron-H-56B-Base的预训练。该方法通过对整个张量使用单一缩放因子进行量化,实现了与BF16训练相当甚至更好的下游任务准确率,同时降低了训练成本。
4. 提出MiniPuzzle模型压缩框架:为了高效地适配不同部署场景,本文引入了一种名为MiniPuzzle的压缩技术,结合了剪枝和蒸馏。通过MiniPuzzle,将56B模型压缩至47B(Nemotron-H-47B-Base),仅使用630亿训练token,推理速度提升20%,且准确率与原始56B模型相当。这使得模型能以FP4精度部署在单张32GiB显存的NVIDIA RTX 5090 GPU上。
5. 展示强大的后训练扩展能力:Nemotron-H基础模型可以被有效地后训练,生成在视觉语言、指令遵循和长上下文(如RULER)基准测试中表现优异的模型。基于Nemotron-H的视觉语言模型(VLM)已作为Cosmos-Reason1项目的一部分,用于开发强大的物理AI推理模型。
本文公开发布了Nemotron-H-56B-Base、Nemotron-H-47B-Base和Nemotron-H-8B-Base的模型检查点,以促进社区的进一步研究。
A2 方法细节
2. 基础模型与预训练
本节讨论了Nemotron-H-8B-Base和Nemotron-H-56B-Base的模型架构,以及生成这些模型的预训练过程中的关键细节。我们还将Nemotron-H-8B/56B-Base与现有的开源SOTA Transformer模型在准确率(在基础模型的通用基准任务上)和推理速度方面进行了比较。
2.1. 模型架构
Nemotron-H模型组成。Nemotron-H模型由Mamba-2(【索引16,Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality,2024】)、自注意力(self-attention)和前馈网络(FFN)层混合组成,具体架构如图2和表1所示。根据先前工作(【索引103,An Empirical Study of Mamba-based Language Models,2024】)的建议,我们将注意力层的数量设置为总层数的约8%,并将其均匀分布在整个模型中。具体来说,Nemotron-H-8B有4个自注意力层(共52层),Nemotron-H-56B有10个自注意力层(共118层)。其余层由FFN和Mamba-2层均等分配。我们还确保:a) 模型的第一个层是Mamba-2层,b) 模型的最后一个层是FFN层,c) 自注意力层总是位于FFN层之前(如同在标准Transformer块中一样,例如【索引101,Attention Is All You Need,2023】)。
模型超参数设置。对于Nemotron-H-8B,我们使用的隐藏层维度为4096;对于Nemotron-H-56B,则为8192。对于较小的模型,我们使用的FFN隐藏层维度为21504,32个注意力查询头,Mamba-2状态维度为128;对于较大的模型,我们使用的FFN隐藏层维度为32768,64个注意力查询头,Mamba-2状态维度为256。两个模型都使用了分组查询注意力(Grouped-Query Attention, GQA)(【索引1,GQA: Training Generalized Multi-Query Transformer Models from Multi-Head Checkpoints,2023】),其中包含8个键值头(key-value heads)、8个Mamba-2组,以及用于FFN层的平方ReLU(squared ReLU)激活函数(【索引95,Primer: Searching for Efficient Transformers for Language Modeling,2022】)。我们没有使用任何位置嵌入。对于Mamba-2层,我们保留了头维度(64)、扩展因子(2)和卷积窗口大小(4)的默认值。我们还使用了RMSNorm(【索引114,Root Mean Square Layer Normalization,2019】)进行归一化,分离了嵌入层和输出层的权重,并且没有使用dropout。我们没有为线性层使用偏置权重。我们在架构中的每个Mamba-2、自注意力和FFN层周围都加入了残差跳跃连接。
| 模型 | 层数 | 模型维度 | FFN维度 | 查询头数 | 键值头数 | 状态维度 | Mamba组数 |
|---|---|---|---|---|---|---|---|
| Nemotron-H-8B | 52 | 4096 | 21504 | 32 | 8 | 128 | 8 |
| Nemotron-H-56B | 118 | 8192 | 32768 | 64 | 8 | 256 | 8 |
Table 1 | Nemotron-H混合Mamba-Transformer架构摘要。
Nemotron-T-8B Transformer基线模型。为了以“苹果对苹果”的方式将Nemotron-H-8B-Base与Transformer模型进行比较,我们还在完全相同的数据上训练了Nemotron-T-8B-Base。Nemotron-T-8B的架构遵循GPT-3的风格(【索引10,Language Models are Few-Shot Learners,2020】)。我们使用了32个Transformer层(每个层包含一个自注意力层和一个FFN层)。与Nemotron-H-8B一样,我们使用了4096的隐藏层维度、32个查询头、带有8个键值头的GQA、21504的FFN隐藏层维度、平方ReLU激活函数、RMSNorm、线性层无偏置权重、无dropout,以及模型嵌入和输出层权重的分离参数。我们使用RoPE(【索引98,RoFormer: Enhanced Transformer with Rotary Position Embedding,2023】)作为位置嵌入。
2.2. 预训练数据
Nemotron-H-8B-Base和Nemotron-H-56B-Base是在一个包含高质量策划数据和合成生成数据的大型语料库上进行预训练的。
2.2.1. 策划数据
数据策划流程。我们为四个主要数据类别建立了独立的数据策划流程:通用网络爬取数据、数学数据、代码数据和“学术”数据。接下来我们将依次讨论每一类。
网络爬取数据。对于Nemotron-H,我们在处理英文Common Crawl数据方面相比Nemotron-4(【索引83,Nemotron4 15B Technical Report. arXiv preprint arXiv:2402.16819,2024】,【索引77,Nemotron-4 340B Technical Report,2024】)做出了几项关键创新;这些创新显著提升了数据质量。关于该数据集的制备细节和各种消融研究,请参阅专门的论文(【索引97,Nemotron-CC: Transforming Common Crawl into a Refined Long-Horizon Pretraining Dataset,2024】)。此处提供一个简要总结。
高质量token提取策略。我们主要致力于提取尽可能多的高质量token,以便能够在长token horizon(例如15万亿或更多token)上训练Nemotron-H。我们从常见的HTML到文本提取、语言过滤、全局模糊去重和精确子串去重开始。然而,在此阶段,我们偏离了常规做法,采用了一个由三个基于模型的分类器组成的集成系统,将每个文档分入五个质量类别。通过使用集成系统,我们获得了比使用单一模型分类器的方法更大、更多样化的高质量token集。为了尽可能保留高质量的token,我们对低、中低和中等质量的类别应用了启发式和困惑度过滤器。我们还对低质量的token进行了改写,以提升其质量;§2.2.2描述了合成数据的生成流程。
数据集构成与效果。最终的数据集包含6.3万亿个token,其中包括4.4万亿个全局去重的“真实”token和1.9万亿个改写后的合成数据token。质量分布如表2所示。如图3所示,与DCLM(【索引39,DataComp-LM: In Search of the Next Generation of Training Sets for Language Models,2024】)和FineWeb-Edu(【索引85,The FineWeb datasets: Decanting the Web for the Finest Text Data at Scale,2024】)等其他领先的英文Common Crawl数据集相比,我们的数据集能产出更高质量的模型:使用高质量子集,MMLU得分比DCLM高出5.6分,而使用未加权的完整数据集,其准确率与DCLM相当,但唯一token数量多4倍。拥有更多的唯一token对于在更长的token horizon上进行预训练至关重要,因为这样可以在整个数据集上进行不超过4到8个epoch的训练。更高的epoch数会导致收益递减(【索引75,Scaling Data-Constrained Language Models,2024】,【索引22,Maximize Your Data’s Potential: Enhancing LLM Accuracy with Two-Phase Pretraining,2024】)。
Table 2 | 6.3万亿英文Common Crawl token的质量分布。
数据集发布。为了社区的利益,我们已根据Common Crawl的使用条款公开发布了该数据集(名为Nemotron-CC1)。
数学数据处理。为确保Common Crawl中的技术页面保留其数学内容,我们利用了Paster等人(【索引84,OpenWebMath: An Open Dataset of High-Quality Mathematical Web Text,2023】)的方法和代码库。我们还发现将此方法应用于维基百科等其他高质量数据源至关重要。我们发现,一个规模更小、质量更高的数据集能在相关基准测试上带来更大的提升,这与FineMath(【索引4,SmolLM2: When Smol Goes Big – Data-Centric Training of a Small Language Model,2025】)的发现类似。
代码数据处理。我们使用大量代码token对Nemotron-H模型进行预训练。我们从先前在Nemotron-4上的工作(【索引83,Nemotron4 15B Technical Report. arXiv preprint arXiv:2402.16819,2024】,【索引77,Nemotron-4 340B Technical Report,2024】)出发,并选择减少HTML、Markdown、YAML、CSS、JSON和Makefile等标记和配置语言中的token数量。原本来自这些语言的token被替换为来自Python、C++、C和Java等流行语言的token。
学术数据处理。Nemotron-H预训练数据集包含了从“高信息量”英文文本中收集的额外token,包括 permissively-licensed 的书籍和文章。我们利用了来自广泛领域(如科学、数学、人文学科、法律和工程学)和大量文档类型(包括同行评审的期刊出版物、教科书、已出版的小说和专利文件)的样本。
学术数据解析与过滤。由于这些文档的原始数据格式差异很大——包括EPuB、HTML、XML、纯文本、PDF、LaTeX和markdown——我们编写了自定义函数将文本从输入格式解析为标准化的输出格式。对于表格、列表和方程式等复杂部分,我们以markdown或LaTeX格式保持其适当的格式。我们利用Éclair(【索引32,Éclair – Extracting Content and Layout with Integrated Reading Order for Documents,2025】)进行PDF到文本的提取。我们还开发了专门的启发式过滤器,以去除页面页眉或页脚中包含的无关信息。然后,我们应用Nemotron-4的启发式和困惑度过滤器,从用于预训练的文档集中移除低质量文档。
学术数据加权。为了更好地在最终的数据混合中对来自这些来源的样本进行加权,我们开发了分类器来评估文档的教育内容和难度。我们还尝试确定每个文档的领域(生物学、商业、化学、计算机科学、经济学、工程学、历史、法律、数学、医学、哲学、物理学、心理学或其他之一)。对于教育内容,我们将所有文档按0(无教育信息)到3(高度相关的教育信息)的数字等级进行评分。对于教育难度,我们将具有高教育内容的文档分为以下级别:小学、初中、高中、本科和研究生。我们能够使用这些类别(例如,“本科水平、高教育内容的生物学”)来确定文档权重(§2.2.3)。正如预期的那样,我们发现在重要领域增加高中/本科水平的高教育内容的权重最为有益。
2.2.2. 合成生成数据
合成数据用途。我们还合成了数据,以去噪低质量数据并增加模型接触的token多样性。我们使用不同的处理流程来合成重写网络爬取、数学和代码数据。
网络爬取数据的合成处理。我们发现,通过LLM对文本进行改写是减少低质量爬取数据中噪音和错误的有效方法,并能生成具有新唯一token的高质量数据变体。这有助于在下游任务中取得更好的结果。
Table 3 | 使用Mistral NeMo 12B指令版生成的网络爬取合成数据token数(十亿)。
改写提示词。我们使用Maini等人(【索引59,Rephrasing the Web: A Recipe for Compute and Data-Efficient Language Modeling,2024】)的“medium Wikipedia”提示词对低质量数据进行了改写。对于高质量文档,我们使用了四个额外的提示词来生成合成数据:
1. 多样化问答(QA)对。以多种形式(例如,是/否问题、开放式问题、多项选择题)就文本中的事实信息提问。
2. 提炼。将文本重写为简洁明了的段落。
3. 提取知识。从文本中重写知识,并忽略无信息的内容。
4. 知识列表。将文本中的关键信息提取为有组织的列表。
合成数据生成细节。这些提示词要求模型提供清晰简洁的回答,同时保留事实信息和具体细节(如数字)。我们总共使用Mistral NeMo 12B的指令版本生成了超过1.8万亿个合成token。按质量和提示词划分的源token和生成token的明细如表3所示。关于此合成数据生成的完整细节,包括使用的提示词,已在Su等人(【索引97,Nemotron-CC: Transforming Common Crawl into a Refined Long-Horizon Pretraining Dataset,2024】)的论文中单独发表。
数学数据的合成处理。我们使用合成数据来增强数学推理基准测试,如GSM8K和MATH(【索引14,Training Verifiers to Solve Math Word Problems,2021】,【索引27,Measuring Mathematical Problem Solving With the MATH Dataset,2021b】),具体细节在Akter等人(【索引2,MIND: Math Informed syNthetic Dialogues for Pretraining LLMs,2024】)的论文中有详细描述。我们使用Nemotron-4-340B(【索引77,Nemotron-4 340B Technical Report,2024】)将OpenWebMath(【索引84,OpenWebMath: An Open Dataset of High-Quality Mathematical Web Text,2023】)从140亿token扩展到超过1000亿token,在受控实验中使GSM8K的性能提升了18个百分点。
数学合成数据生成方法。为了实现这一目标,我们从Common Crawl的技术预训练文档开始,并利用Nemotron-4-340B生成对话,其中知识渊博的角色指导经验较少的角色(例如,学生和老师之间的互动)。这种方法与《语言模型物理学》系列(【索引5,Physics of Language Models: Part 3.2, Knowledge Manipulation,2024】)中的见解相一致,并融入了诸如展示错误答案及修正等策略。通过将内容构建为学习互动,我们的方法从公共语言模型中提炼出广泛的知识,而不会对基准测试产生过拟合。
代码数据的合成处理。为了增加以解决问题为中心的多样化代码数据,我们选择在11种编程语言(如Python、C++、C、C#和Java)中合成生成混合自然语言和代码的数据集。为此,我们提示Mixtral 8x22B根据我们策划的代码数据集中抽取的代码片段生成编程问题。这些样本的代码行数(LOC)从1到15行不等,通常只是一个小函数。我们还提示Mixtral 8x22B解决生成的问题,并以最小的努力移除明显无效的解决方案。例如,我们提取模型生成的Python代码并尝试将其解析为抽象语法树(AST),丢弃解析失败的样本。最后,为了形成训练样本,我们将生成的问题和答案组合成一个单一的样本,该样本混合了自然语言指令、生成的代码(通常包含在markdown代码块中),以及通常对方法的解释。
SFT风格数据的合成。我们进一步在预训练语料库中添加了合成的SFT风格数据;这提高了基础模型遵循指令的能力。我们使用Qwen2.5系列、Mixtral 8x22B、Nemotron-4-340B和DeepSeek-R1(仅用于56B模型)来生成这些数据集。我们遵循OpenMathInstruct-2(【索引100,OpenMathInstruct-2: Accelerating AI for Math with Massive Open-Source Instruction Data,2024】)中提出的流程,使用精心策划的种子数据(如AoPS2)来提高数学能力。我们还采用Genetic Instruct(【索引60,Genetic Instruct: Scaling up Synthetic Generation of Coding Instructions for Large Language Models,2024】)中提出的方法,以tigerbot-leetcode3和the-stack-v24为种子数据,合成了代码数据。为了让基础模型具备高质量的通用知识,我们针对选定的主题和特定领域的文档生成了问题-解决方案对。总共,我们向训练语料库中添加了2300亿个合成SFT风格token(1740亿数学token,350亿代码token,以及210亿通用知识token)。
2.2.3. 数据混合与排序
数据混合构成。我们的数据混合包含七个高级数据类别:网络爬取数据(web crawl)、数学(math)、维基百科(wikipedia)、代码(code)、学术(academic)、crawl++和多语言(multilingual)以及合成SFT风格数据。Crawl++包含OpenWebText、BigScience和Reddit等网络爬取衍生的数据集。我们的多语言数据涵盖九种语言:德语、西班牙语、法语、意大利语、葡萄牙语、中文、日语、韩语和俄语。我们设计数据混合的方式是,使得所有同等质量的数据源权重相似,而高质量数据源的权重要高于低质量数据源。我们在Feng等人(【索引22,Maximize Your Data’s Potential: Enhancing LLM Accuracy with Two-Phase Pretraining,2024】)的论文中提供了更多关于估算数据集质量的细节。
(a) 阶段1的数据混合。
(b) 阶段2的数据混合。
(c) 阶段3的数据混合。
(d) 阶段4的数据混合。
分阶段数据混合方法。我们采用分阶段的数据混合方法(【索引22,Maximize Your Data’s Potential: Enhancing LLM Accuracy with Two-Phase Pretraining,2024】)来预训练Nemotron-H的两个基础模型。在第一阶段,我们使用促进数据多样性的数据混合;在第二和第三阶段,我们主要使用高质量数据集(如维基百科)。我们在训练的60%节点切换到第二阶段,在80%节点切换到第三阶段。第四阶段在最后3800亿训练token时进行。每个阶段使用的数据混合如图4所示。我们对一个在1万亿token上训练的80亿参数模型的初步实验表明,分阶段的预训练方法比随机数据排序的性能高出3.4%。
2.3. FP8 训练方法
FP8训练策略。Nemotron-H-56B-Base采用层级混合精度进行预训练:模型中所有的线性层(包括FFN块中的两个线性层以及注意力块中的QKV/输出投影)都以FP8精度计算,除了前4个和后4个层,它们保持在BF16精度。我们对线性层的前向和后向传播都进行了量化。与Micikevicius等人(【索引70,FP8 Formats for Deep Learning,2022】)的研究一样,我们使用一种混合FP8方法,即对权重和激活张量使用E4M3格式,对梯度张量使用E5M2格式。
FP8量化细节。我们采用了FP8逐张量动态量化,该过程包括几个步骤:首先,我们计算一个量化尺度(并检查除零错误);然后,我们将输入张量与该量化尺度相乘,并将其转换为所需的FP8格式。量化尺度的计算方法是,将该FP8数据格式可表示的最大值除以输入张量的最大绝对值。
FP8训练中的损失变化。我们观察到,FP8预训练实验的对数似然损失曲线几乎总是高于其BF16对应实验,但差距出奇地小,训练和验证集上的典型相对差距小于0.1%。这个损失差距随着训练的进行趋于减小,但在训练的最后四分之一阶段又会扩大。我们推测这是因为在训练后期占主导的极小梯度更新,尽管在E5M2格式下操作,仍被刷新为零。图5展示了预训练期间一个BF16和一个FP8实验之间的典型相对损失差距。
FP8训练的下游任务性能。尽管损失曲线略差,我们观察到使用FP8训练的模型的下游评估结果与BF16实验一样好,甚至更好。传统上作为下游任务代理指标的对数似然,在我们的实验中并非一个可靠的预测器;事实上,我们经常观察到损失曲线更好的模型在下游评估任务上表现更差。
FP8在特定任务上的优势。对于编码和数学任务,使用逐张量动态缩放的实验结果通常远优于BF16实验。图6展示了两个8B参数Transformer在15万亿token上训练的准确率差距,一个用FP8训练,另一个用BF16训练。由于在相同的token horizon上取得了相同或更好的下游结果,我们从不需要对FP8模型进行相对于BF16模型的过度训练。我们对用FP8训练的模型进行了BF16和FP8两种精度的评估;它们的结果通常相同,BF16的评估通常略好。然而,对于56B模型,FP8的评估最终超过了BF16。
FP8方法验证。我们发现在构建我们的FP8方法时,在至少8B参数的模型上进行验证非常重要,因为在较小模型上的结果不具有普适性。所有实验的token horizon至少为1万亿。
2.4. 预训练
预训练参数。我们分别在15万亿token和20万亿token的训练尺度上训练了Nemotron-H-8B-Base和Nemotron-H-56B-Base。我们使用了8192的序列长度和768的全局批量大小(每个批次6,291,456个token)。我们没有使用任何批量大小的 ramp-up 策略。对于Nemotron-H-8B-Base,我们使用的峰值学习率为 $8 \cdot 10^{-4}$,并在83亿token上进行预热;对于Nemotron-H-56B-Base,我们使用的峰值学习率为 $4 \cdot 10^{-4}$,并在250亿token上进行预热。在这两种情况下,我们都使用了余弦学习率衰减,最小值为峰值的1%,权重衰减为0.1,Adam的 $\beta_1$ 和 $\beta_2$ 分别为0.9和0.95。
预训练软件栈与并行策略。我们使用Megatron-LM进行Nemotron-H的预训练;我们依赖Transformer Engine提供FP8支持。我们使用了8路张量模型并行(【索引93,Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism,2020】),并结合序列并行(【索引35,Reducing Activation Recomputation in Large Transformer Models,2022】)以进一步节省内存,同时使用了768路数据并行,并将优化器状态分布在数据并行副本上(【索引87,ZeRO: Memory Optimizations Toward Training Trillion Parameter Models,2020】)。
2.4.1. 弹性训练
大规模训练的挑战与对策。我们在大规模训练模型时也高度关注最小化停机时间:在大量GPU上进行长时间的训练任务时,作业失败和减速是不可避免的,这会导致训练效率的损失。失败可能由GPU相关错误(如无法纠正的ECC错误或GPU掉线)、网络结构或Lustre等共享基础设施的错误,或应用程序引起的失败(如某些操作的数值不稳定性或“坏”数据)引起。我们还可能因多种原因(如单GPU节流或通信集合操作变慢)经历作业减速。
提升训练弹性。为了最小化上述问题的影响,我们专注于 a) 将故障准确归因于正确的基础设施和应用组件,以防止重复中断,以及 b) 减少每次故障的恢复开销。为此,我们利用DGX Cloud Resilience服务来确保故障被正确归因,并且相应的硬件组件(例如,带有故障GPU的特定HGX节点)不会被包含在下一次训练作业实例中,从而防止重复失败。启用DGX Cloud Resilience后,我们集群的平均无故障时间(MTBF)提高了3.3倍。
低开销检查点与快速重启。我们还主动频繁地保存检查点,以最小化故障发生时丢失的工作量(从上一个检查点保存到故障发生时的权重更新会丢失)。我们使用NVRx来促进低开销的检查点操作,通过异步保存检查点而不阻塞训练。我们还优化了作业重启时间(即,Megatron-LM在开始训练迭代前所需的一系列操作,如分布式运行时初始化和检查点加载)。
训练效率。总而言之,上述组件使我们能够在6144个NVIDIA H100 GPU上高效地训练Nemotron-H-56B-Base。
2.5. 基础模型评估
评估框架。除非另有说明,所有模型的评估均由我们自己进行。我们的评估设置建立在lm-evaluation-harness之上,以进行公平比较,并做了以下修改:
1. 数学推理:我们在GSM8K(【索引14,Training Verifiers to Solve Math Word Problems,2021】)和MATH(【索引27,Measuring Mathematical Problem Solving With the MATH Dataset,2021b】)基准上进行评估。我们还特别关注MATH基准中的竞赛级别部分,标记为“MATH Level 5”。在仔细检查不一致的案例后,我们结合使用Math-Verify和NeMo-Skills的数学评分工具来为解决方案评分。
2. MATH-R1评估:我们还引入了MATH任务的MATH-R1版本,其中我们提供了由DeepSeek-R1(【索引19,DeepSeek-R1: Incentivizing Reasoning Capability in LLMs via Reinforcement Learning,2025b】)生成的少样本示例。Nemotron-H和Qwen 2.5模型都通过简单地更改提示词在该基准上显示出显著的改进。
3. 代码任务:对于代码任务(HumanEval(【索引11,Evaluating Large Language Models Trained on Code,2021】)、MBPP(【索引6,Program Synthesis with Large Language Models,2021】)及其EvalPlus变体),我们始终使用EvalPlus(【索引46,Is Your Code Generated by ChatGPT Really Correct? Rigorous Evaluation of Large Language Models for Code Generation,2023b】)对生成结果进行净化。
4. 通用推理基准:通用推理基准(OpenBookQA(【索引71,Can a Suit of Armor Conduct Electricity? A New Dataset for Open Book Question Answering,2018】)、PIQA(【索引9,PIQA: Reasoning about Physical Commonsense in Natural Language,2019】)、Hellaswag(【索引113,HellaSwag: Can a Machine Really Finish Your Sentence?,2019】)、Winogrande(【索引88,WinoGrande: An Adversarial Winograd Schema Challenge at Scale,2019】))保持不变,除了ARC-Challenge(【索引13,Think you have Solved Question Answering? Try ARC, the AI2 Reasoning Challenge,2018】),我们像MMLU(【索引26,Measuring Massive Multitask Language Understanding,2021a】)一样,同时呈现所有选项。
56B模型评估结果。Nemotron-H-56B-Base和Nemotron-H-8B-Base在常见基准任务上的准确率结果分别显示在表4和表5中。Nemotron-H-56B-Base和Nemotron-H-47B-Base均在FP8精度下进行评估,所有其他模型则在BF16精度下评估。在这两种模型规模下,Nemotron-H基础模型相对于同等规模的SOTA Transformer模型,其准确率相当或更优。
Table 4 | Nemotron-H-56B-Base(及其蒸馏版本Nemotron-H-47B-Base,详见§3)与现有SOTA模型的准确率对比。N-H为Nemotron-H的缩写,DS为DeepSeek的缩写。斜体数字来自DeepSeek-V3报告(【索引20,DeepSeek-V3 Technical Report,2025c】),所有其他数字由我们运行。我们加粗了每行的最高准确率,不包括斜体数字。
与70B级别模型的比较。我们将Nemotron-H-56B-Base直接与Qwen-2.5-72B-Base(【索引86,Qwen2.5 Technical Report,2025】)和Llama-3.1-70B-Base(【索引69,The Llama 3 Herd of Models,2024】)进行比较。在我们评估的17个任务中,Nemotron-H-56B-Base在9个任务上取得了三个模型中的最高准确率,Qwen-2.5-72B-Base在7个任务上取得了最高准确率,Llama-3.1-70B-Base在1个任务上取得了最高准确率;Nemotron-H-56B-Base在除Winogrande外的所有任务上都优于Llama-3.1-70B-Base。我们进一步将Nemotron-H-56B-Base与两个规模大得多的SOTA模型——DeepSeek-V3-671B-Base(【索引20,DeepSeek-V3 Technical Report,2025c】)和Llama-3.1-405B-Base(【索引69,The Llama 3 Herd of Models,2024】)进行比较;我们使用了这两个模型在已报告的共同任务子集上的公开报告数字。令人惊讶的是,Nemotron-H-56B-Base在这些模型中仍然具有竞争力,在10个重叠任务中分别有4个和5个任务上超过了DeepSeek-V3-671B-Base和Llama-3.1-405B-Base。
8B模型评估结果。我们观察到Nemotron-H-8B-Base与Qwen-2.5-7B-Base和Llama-3.1-8B-Base相比也取得了类似的准确率结果(表5)。在用于评估的15个任务中,Nemotron-H-8B-Base、Qwen-2.5-7B-Base和Llama-3.1-8B-Base分别在7个、8个和0个任务上取得了三个模型中的最高准确率。Nemotron-H-8B-Base在常识理解任务上尤其强大,在该类别的5个任务中有4个取得了最高准确率。如上所述,我们还与一个更大的模型Gemma-3-12B-Base(【索引17,Gemma 3 Technical Report,2025】)进行了比较;Nemotron-H-8B-Base与这个更大的模型相比具有竞争力。
Table 5 | Nemotron-H-8B-Base与现有SOTA模型以及一个在完全相同数据上训练的Transformer模型(Nemotron-T-8B-Base)的准确率对比。Gemma 3模型的基于似然的评估结果来自Gemma 3报告(【索引17,Gemma 3 Technical Report,2025】),所有其他数字由我们运行。我们加粗了每行的最高准确率,不包括斜体数字。
混合模型与纯Transformer模型的直接比较。表5还包含了迄今为止最大规模的混合Mamba-Transformer与纯Transformer模型的“苹果对苹果”比较(Nemotron-H-8B-Base 对比 Nemotron-T-8B-Base)。Nemotron-H-8B-Base在15个任务中的7个任务上取得了比在完全相同数据上训练的Transformer模型更高的准确率,并且在另外4个任务上差距在一个百分点以内。总的来说,这个实验表明,当在SOTA规模上进行训练时,混合模型可以达到与Transformer模型相同或更好的准确率。
2.6. 推理速度
推理加速。由于自注意力层的减少,Nemotron-H-8B/56B-Base相比于表4和表5中的Transformer模型提供了推理时间的加速。为了量化这些加速效果,我们在图1中绘制了Nemotron-H-56B-Base及同等规模Transformer模型的MMLU-Pro准确率与推理吞吐量的关系图,并在图7中绘制了Nemotron-H-8B-Base及相关基线模型的相应关系图。我们使用的输入序列长度为65536,并要求模型生成1024个输出token。我们使用了Nemotron-H推理的初始Megatron-LM实现和vLLM v0.7.3用于基线模型。在这些实验中,我们通过使用尽可能大的批量大小来最大化单GPU的推理吞吐量,并且所有实验都在NVIDIA H100 GPU上运行。我们报告的结果是按使用的GPU数量归一化的(即,每秒每GPU输出的token数)。
量化加速效果。在上述设置中,Nemotron-H-56B-Base每秒每GPU生成的输出token数量比Qwen-2.5-72B和Llama-3.1-70B多2.4倍(经过蒸馏到Nemotron-H-47B-Base后,则多2.9倍,详见§3)。与Llama-3.1-405B相比,Nemotron-H-56B-Base的吞吐量高出19.6倍。与较大模型类似,Nemotron-H-8B-Base的推理速度也比相应的Transformer模型快。如图7所示,在更长的上下文中,我们测得与Qwen-2.5-7B和Llama-3.1-8B相比,速度分别提升了1.8倍和3倍。推理速度的提升源于两个因素:a) Mamba-2层中的恒定计算量,而非自注意力层中的线性计算量;b) Mamba-2层较低的内存占用使得可以使用更大的批量大小,从而提高效率。我们预计,进一步优化Nemotron-H的推理将带来相对于Transformer模型更多的速度提升,后者在过去几年中已经得到了大量优化。
3. 压缩与蒸馏
模型压缩动机。LLM的高效部署通常需要根据硬件施加的特定约束来调整模型架构。本节中,我们描述了用于将56B模型压缩至47B参数的剪枝和蒸馏过程,目标是在NVIDIA GeForce RTX 5090 GPU上运行更长上下文的推理(在具有32GiB内存的RTX 5090 GPU上以FP4精度存储56B参数模型的权重将需要29.5GiB,仅留下2.5GiB用于KV缓存和激活缓冲区)。我们引入了一种名为MiniPuzzle的新型模型压缩框架,它结合了Minitron(【索引76,Compact Language Models via Pruning and Knowledge Distillation,2024】,【索引96,LLM Pruning and Distillation in Practice: The Minitron Approach,2024】)的简洁性和Puzzle(【索引8,Puzzle: Distillation-Based NAS for Inference-Optimized LLMs,2024】)的多功能性。与从头开始训练相比,我们的方法使用的token数量减少了约300倍,才获得了47B模型。我们压缩后的47B模型实现了1.2倍的推理加速(图1),同时其准确率与原始56B模型相当。
3.1. MiniPuzzle 概述
MiniPuzzle框架。MiniPuzzle将轻量级剪枝与神经架构搜索(NAS)相结合,以生成满足特定目标部署约束的高效压缩模型。图8提供了该框架的高层概述。MiniPuzzle的优化过程包括三个阶段:a) 重要性估计(§3.2),b) 条件NAS(§3.3),以及 c) 知识蒸馏(§3.4)。
3.2. 重要性估计
重要性估计阶段。MiniPuzzle首先为每个模型组件(例如,层、FFN神经元)收集重要性或敏感性分数,以帮助决定移除哪些组件;这就是重要性估计阶段。此阶段计算的分数用于决定哪些模型组件可以被剪枝。我们注意到,基于梯度信息的敏感性分析在现代LLM规模下通常不切实际(【索引76,Compact Language Models via Pruning and Knowledge Distillation,2024】);因此,我们依赖于一种仅使用前向传播的轻量级策略。在这项工作中,我们使用一种在我们的消融研究中表现良好的简化方法:a) 剪枝层,以及 b) 剪枝FFN隐藏维度(实际上是神经元)。我们也尝试了剪枝Mamba头和Mamba头维度;不幸的是,这两个维度都导致了严重的准确率下降(特别是头维度剪枝)。
计算重要性分数的方法。现在我们描述如何计算每个层和FFN神经元的重要性。
层重要性计算。我们使用Puzzle中的评分方法来估计模型中每个层的重要性。给定一个输入张量,层重要性被计算为完整模型的LM头之前的中间激活张量与移除了特定层的模型的相同激活张量之间的均方误差(MSE)。我们在训练数据的一个小的随机子集(在我们的案例中是128个样本)上对这些排名进行平均,以获得一个考虑到样本变异性的可靠重要性估计。
层重要性分析。图9绘制了Nemotron-H-56B-Base中每个层的平均重要性分数。绿色、蓝色和红色的虚线分别对应自注意力层、FFN层和Mamba层。我们从图中注意到,56B模型遵循了Minitron中观察到的典型模式:最重要的层集中在模型的开头和结尾。此外,MSE重要性显示,尽管56B模型只有10个自注意力层,但一些自注意力层被排在最不重要的行列中,特别是第84层(第7个自注意力层)。然而,第40、51、62和73层似乎比其邻近的其他层更重要。
FFN重要性计算。FFN层内部由两个线性算子和一个非线性激活函数组成:
这里,$X$ 表示输入,$W_1$ 和 $W_2$ 是FFN层中相关的两个权重矩阵。$W_1, W_2 \in R^{d_{model} \times d_{ffn}}$,其中 $d_{model}$ 和 $d_{ffn}$ 分别是模型的隐藏维度和FFN的隐藏维度(见表1)。$\delta(\cdot)$ 指的是非线性激活函数(在这项工作中是平方ReLU)。
FFN神经元重要性。我们遵循与Minitron相同的程序,通过检查每个FFN层第一个线性算子中每个神经元产生的输出集来计算其重要性。为此,我们使用了一个包含1024个样本的小型校准数据集。形式上,我们通过聚合给定输入批次B的输出来计算每个神经元的重要性得分:
这里,$w_1^i$ 指的是权重矩阵 $W_1$ 的第 $i$ 行。$\sum_{B,S}$ 指的是沿批次和序列维度的聚合。我们遵循Minitron论文(【索引76,Compact Language Models via Pruning and Knowledge Distillation,2024】)中的观察,沿批次和序列维度使用均值和l2范数聚合函数。对于一个得分序列 $S$,均值聚合定义为 $\frac{1}{n} \sum_{i=1}^{n} |S_i|$,l2范数是 $\sqrt{\sum_{i=1}^{n} S_i^2}$。
3.3. 条件神经架构搜索(Conditional NAS)
条件NAS流程。MiniPuzzle的条件NAS利用每个层和FFN神经元的重要性分数(如§3.2所述计算)来识别既满足内存限制又保留最关键组件的架构。这个过程使我们能够有效地探索可能的层配置的巨大搜索空间,并包括三个步骤:
1. 剪枝候选枚举。我们遍历一个由每种层类型(即自注意力层、FFN层和Mamba-2层的数量)的可能层数和FFN隐藏维度(24576, 25600, ..., 32768)组成的网格。对于每个目标层数和FFN隐藏维度大小,我们根据它们各自的重要性分数(有关重要性估计的细节见§3.2)选择顶级的层和神经元,并实现相应的架构;即,我们丢弃不相关的层,并将所有FFN层剪枝到相同的目标宽度。只有满足在100万token上下文上FP4推理所需的总内存小于31.7 GiB的目标内存约束的架构才被保留,最终得到大约400个候选架构。
2. 架构排名。我们使用一个轻量级指标对候选架构进行评分,该指标将它们与父模型(Nemotron-H-56B-Base)在一个包含100万token的小型验证语料库上进行比较。具体来说,我们使用两种不同的分数来估计每个候选模型的质量:a) 下一词元准确率,计算子模型正确预测下一个真实词元的频率;b) 下一词元父模型一致性,计算子模型与父模型在贪婪预测下一个词元上的一致性频率。我们使用这两种评分指标对所有候选架构进行排名,并保留表现最好的130个。
3. 候选选择。现在,我们通过在一部分评估任务上对这130个排名靠前的架构进行基准测试,将其数量精简到可管理的3个。表6的第一列列出了我们用来评估最后3个候选的任务。我们使用这些任务的平均分进行评估。图10绘制了所有130个候选的平均基准分数,图11展示了所有三个候选具体丢弃的层。候选1和2取得了最高的基准分数,并且丢弃了相似的层。候选3丢弃的层略有不同,在剪枝后表现明显更差。然而,我们在此处包含它,是为了研究蒸馏能在多大程度上弥补损失。
3.4. 使用蒸馏进行再训练
蒸馏恢复精度。为了恢复因剪枝而损失的准确率,模型会进行持续训练。最近的研究表明,将原始模型的知识蒸馏到剪枝后的模型中,其效果优于传统的微调(如Minitron和Puzzle所示);因此,我们采用基于logit的蒸馏进行持续训练,在准确率恢复阶段仅使用前向KL散度损失(关于蒸馏损失公式的更多细节,请参见Minitron论文的第3节)。
两阶段蒸馏。在§3.3中描述的初始候选选择过程之后,我们以56B模型为教师,对前3个候选模型进行了约70亿token的短时蒸馏。我们的实验表明,这一步对于挑选出最终最佳候选至关重要。表6详细列出了3个候选模型在短时蒸馏前后的基准分数。在这三个模型中,平均(在所有基准上)最准确的候选模型将被选中,进行一次使用630亿token的扩展蒸馏,以产生最终的47B模型;在我们的案例中,候选2表现最佳,如表6所示。我们观察到,最终候选的选择取决于最终目标;例如,优先考虑编码任务的准确率还是常识推理任务的准确率。
| 基准 | Nemotron-H 56B-Base | 候选1 47B-pruned | 候选2 47B-pruned | 候选3 47B-pruned |
|---|---|---|---|---|
| MMLU | 84.2 | 81.5 → 83.7 | 80.8 → 83.2 | 81.3 → 83.6 |
| 常识理解 (平均) | 67.6 | 64.2 → 66.6 | 63.8 → 67.2 | 64.5 → 66.8 |
| GSM8k | 91.4 | 67.9 → 89.2 | 60.7 → 91.1 | 19.5 → 90.5 |
| 代码 (平均) | 67.0 | 52.0 → 64.9 | 52.7 → 64.6 | 23.0 → 64.9 |
| 平均 | 70.6 | 63.3 → 69.3 | 62.1 → 69.6 | 53.7 → 69.5 |
Table 6 | 剪枝候选在轻量级蒸馏前后的基准分数对比。分数可能与表4中呈现的略有不同,因为评估设置和不同任务集的平均值存在微小差异。
蒸馏参数。所有的蒸馏运行都使用FP8精度,softmax温度为1.0,该温度控制知识蒸馏过程中概率分布的平滑度(【索引28,Distilling the Knowledge in a Neural Network,2015】),序列长度为8192,批量大小为768个样本。
3.5. 结果
压缩成果。通过使用MiniPuzzle方法,我们高效地将56B模型压缩至47B参数,具体方法是剪枝整个层(深度)和FFN隐藏维度大小,从而提高了推理速度,并使得在32GiB的NVIDIA GeForce RTX 5090 GPU上能够进行更长上下文的推理。47B模型保留了56B模型一半的自注意力层(5层而非10层),以及44个Mamba-2层(从54层减少)和49个FFN层(从54层减少)。此外,FFN的中间大小从32768剪枝至30720。如图1和表4所示,最终得到的47B模型在长上下文推理上速度提升了1.2倍,在基准测试上实现了近乎无损的准确率,而所需的token数量比从头训练减少了约300倍。
4. 视觉语言模型
本节介绍Nemotron-H-8B-VLM和Nemotron-H-56B-VLM,它们分别基于Nemotron-H-8B-Instruct和Nemotron-H-56B-Base构建。通常建议在对齐(aligned)模型的基础上构建视觉语言模型(VLM),以增强其指令遵循能力(【索引15,NVLM: Open Frontier-class Multimodal LLMs,2024】)。在这项工作中,我们选择在Nemotron-H-8B-Instruct和Nemotron-H-56B-Base(因为Nemotron-H-56B-Instruct尚不可用)上构建Nemotron-H-8B-VLM和Nemotron-H-56B-VLM。Nemotron-H-8B-VLM和Nemotron-H-56B-VLM都已被用作构建Cosmos-Reason1(【索引78,Cosmos-Reason1: From Physical Common Sense To Embodied Reasoning,2025】)的基础模型,这是一个用于物理AI的前沿推理模型家族。
4.1. 模型架构
VLM架构选择。在以往的研究中,VLM是通过将一个纯文本LLM与一个视觉编码器集成来构建的。用于这种集成的两种流行架构是仅解码器架构(【索引45,Visual Instruction Tuning,2023a】)和基于交叉注意力的架构(【索引3,Flamingo: A Visual Language Model for Few-shot Learning,2022】)。我们使用仅解码器的NVLM-D架构(【索引15,NVLM: Open Frontier-class Multimodal LLMs,2024】),因为它具有简洁性、参数效率、处理高分辨率图像输入的能力,以及通过将非文本token(如图像、音频、视频)映射到与文本token相同的嵌入空间来统一处理多种模态。Nemotron-H-VLM由一个视觉编码器、一个投影器(一个两层FFN)和Nemotron-H LLM主干组成。
视觉编码器与图像处理。对于8B和56B的VLM,我们都使用InternViT-300M-V2.5(【索引12,Expanding Performance Boundaries of Open-Source Multimodal Models with Model, Data, and Test-Time Scaling,2024】)作为视觉编码器。它处理448×448像素的静态图像作为输入,patch大小为14×14,总共生成1024个视觉token。每个视觉token由一个1024维的向量表示。遵循NVLM设计(【索引15,NVLM: Open Frontier-class Multimodal LLMs,2024】),我们根据输入图像的分辨率将其动态调整到最接近的预定义宽高比 $\{w : h\}$($w$ 和 $h$ 是整数),并将其分割成 $w \times h \le 12$ 个图块,每个图块对应一个448×448像素的图像块。为了保留全局上下文,我们还生成一个448×448的缩略图块,它是整个图像的缩小版本。例如,给定一张1920×1080的图像,最接近的预定义宽高比是3 : 2。该图像被调整为1344×896,并被划分为 $3 \times 2 + 1 = 7$ 个图块,包括一个缩略图块。
视觉特征处理与输入LLM。为了减少LLM的处理开销,我们通过将四个相邻的图像token分组并沿通道维度拼接,将1024个token下采样到256个。来自多个图块的图像token与一个交错的图块ID标签(以原始文本格式)拼接在一起,这为下游的LLM提供了关于动态分块结构的信息;我们发现这对于提高各种视觉-语言基准的准确率至关重要。更多细节,请参见NVLM技术报告(【索引15,NVLM: Open Frontier-class Multimodal LLMs,2024】)。拼接后的视觉token由一个两层的FFN块处理,该块将每个视觉token映射到文本token的嵌入空间。然后,这些嵌入被送入LLM主干。
4.2. 训练方法与数据
两阶段训练法。遵循NVLM(【索引15,NVLM: Open Frontier-class Multimodal LLMs,2024】)的方法,Nemotron-H-VLM的训练分为两个阶段:
1. VL预训练。我们只训练用于模态对齐的两层FFN,同时保持Nemotron-H主干和视觉编码器冻结。
2. VL SFT。我们在各种面向任务的SFT数据上对视觉编码器、FFN投影器和Nemotron-H主干进行端到端的微调。
VL预训练数据。对于VL预训练,我们利用了来自NVLM(【索引15,NVLM: Open Frontier-class Multimodal LLMs,2024】)的一个大规模且多样化的图文预训练数据集,包括图像描述(【索引43,Microsoft COCO: Common Objects in Context,2015】,【索引92,Conceptual Captions: A Cleaned, Hypernymed, Image Alt-text Dataset for Automatic Image Captioning,2018】,【索引82,Im2text: Describing Images using 1 Million Captioned Photographs,2011】,【索引40,BLIP: Bootstrapping Language-Image Pretraining for Unified Vision-Language Understanding and Generation,2022】)、自然图像上的视觉问答(VQA)(【索引24,Making the V in VQA Matter: Elevating the Role of Image Understanding in Visual Question Answering,2017】,【索引36,Visual Genome: Connecting Language and Vision Using Crowdsourced Dense Image Annotations,2017】)、视觉图表(【索引31,DVQA: Understanding Data Visualizations via Question Answering,2018】)和文档理解(【索引61,Docmatix - A Huge Dataset for Document Visual Question Answering,2024】)、光学字符识别(OCR)(【索引72,OCR-VQA: Visual Question Answering by Reading Text in Images,2019】)和场景文本识别(【索引102,COCO-Text: Dataset and Benchmark for Text Detection and Recognition in Natural Images,2016】),以及视觉数学推理(【索引44,CLEVR-Math: A Dataset for Compositional Language, Visual and Mathematical Reasoning,2022】)数据。总体而言,我们的VL预训练数据集包含1.3亿个样本。我们进行了仔细的净化,以移除任何潜在的有害内容。
VL SFT数据。对于VL SFT,我们利用了来自NVLM(【索引15,NVLM: Open Frontier-class Multimodal LLMs,2024】)和Eagle2(【索引41,Eagle 2: Building Post-Training Data Strategies from Scratch for Frontier Vision-Language Models,2025】)的多样化和高质量数据集。除了前面提到的类别外,该数据集还包括基于知识的VQA(【索引62,OK-VQA: A Visual Question Answering Benchmark Requiring External Knowledge,2019】)、视觉推理(【索引30,GQA: A New Dataset for Real-World Visual Reasoning and Compositional Question Answering,2019】)、科学相关的VQA(【索引55,Learn to Explain: Multimodal Reasoning via Thought Chains for Science Question Answering,2022】)和视觉指令遵循数据。总体而言,我们的VL SFT数据集包含600万个图文样本。Dai等人(【索引15,NVLM: Open Frontier-class Multimodal LLMs,2024】)和Li等人(【索引41,Eagle 2: Building Post-Training Data Strategies from Scratch for Frontier Vision-Language Models,2025】)有更多细节。
4.3. 视觉-语言基准测试结果
评估基准。我们在一系列全面的视觉-语言基准上评估了Nemotron-H-8B-VLM和Nemotron-H-56B-VLM,包括MMMU(【索引111,MMMU: A Massive Multi-discipline Multimodal Understanding and Reasoning Benchmark for Expert AGI,2024】)、MathVista(【索引56,MathVista: Evaluating Mathematical Reasoning of Foundation Models in Visual Contexts,2024】)、ChartQA(【索引63,ChartQA: A Benchmark for Question Answering about Charts with Visual and Logical Reasoning,2022】)、AI2D(【索引34,A Diagram Is Worth A Dozen Images,2016】)、OCRBench(【索引47,OCRBench: On the Hidden Mystery of OCR in Large Multimodal Models,2024】)、TextVQA(【索引94,Towards VQA Models That Can Read,2019】)、RealWorldQA(【索引108,Grok-1.5 Vision Preview: Connecting the Digital and Physical Worlds with our First Multimodal Model,2024】)和DocVQA(【索引64,DocVQA: A Dataset for VQA on Document Images,2021】)。这些基准评估了广泛的能力,如多模态推理、视觉环境下的数学推理、自然图像理解、场景文本识别、图表与文档理解、现实世界感知和OCR。
8B VLM性能。我们将Nemotron-H-8B-VLM与基于LLaMA-3.1-8B-Instruct和Qwen2.5-7B-Instruct构建的VLM进行了比较,使用了相同的训练方法和数据集。如表7所示,Nemotron-H-8B被证明是开发同类最佳VLM的强大LLM主干。
Table 7 | Nemotron-H-8B-VLM在视觉-语言基准上的评估。我们与使用相同训练方法构建的、基于Llama-3.1-8B-Instruct和Qwen2.5-7B-Instruct的VLM进行比较。
56B VLM性能。我们进一步将Nemotron-H-56B-VLM与NVLM-1.0-D(【索引15,NVLM: Open Frontier-class Multimodal LLMs,2024】)以及使用相同训练配方构建的基于Qwen2.5-72B-Instruct的VLM在表8中进行了比较。Nemotron-H-56B-VLM取得了SOTA的结果,展示了优于以往模型的质量。
Table 8 | 我们在一系列视觉-语言基准上评估了Nemotron-H-56B-VLM,并将其与使用相同训练配方构建的基于Qwen2.5-72B-Instruct的VLM进行了比较。我们还与NVIDIA之前开发的SOTA VLM NVLM-1.0-D 72B进行了比较。†我们使用VLMEvalKit(【索引21,VLMEvalKit: An Open-Source Toolkit for Evaluating Large Multi-Modality Models,2024】)在MMMU和MathVista上评估了NVLM-1.0-D 72B;这与其他模型使用的评估设置相同,并且取得了比官方数字更好的结果。
5. Nemotron-H 推理模型
5.1. Nemotron-H 与 Transformer 模型的对齐研究
对齐能力研究。我们首先研究了我们的混合Mamba-Transformer模型在多大程度上可以被有效地后训练成指令调优和长上下文变体。我们在预训练于相同数据上的Transformer(Nemotron-T-8B-Exp-Base)和混合(Nemotron-H-8B-Exp-Base)模型上进行了实验,这两个模型在预训练性能上非常相似,如表9所示。
Table 9 | Transformer和混合基础模型在选定任务上的比较。
多阶段后训练流程。然后,我们使用一个多阶段训练程序将这两个基础模型后训练成指令调优变体,即Nemotron-T-8B-Exp-Instruct和Nemotron-H-8B-Exp-Instruct。在第一阶段,我们对一个包含600万样本的混合数据集进行监督微调(SFT),该数据集包括代码、数学和通用指令遵循任务。为了提高长上下文性能,我们加入了对长达512k token的扩展对话的训练,方法是连接较短的样本。我们通过在对话中插入对早期轮次的引用和引入语义相关的片段来进一步增强上下文。这即使在较短的上下文长度(例如128k)下也提高了RULER(【索引29,RULER: What’s the Real Context Size of Your Long-Context Language Models?,2024】)分数。
偏好调优阶段。在第二阶段,我们转向使用离线RPO(【索引99,Reward-aware Preference Optimization: A Unified Mathematical Framework for Model Alignment,2025】)在通用领域提示上进行偏好调优。然后,我们进行一个额外的RPO轮次,专注于狭窄的指令遵循,遵循IFEval(【索引115,Instruction-following Evaluation for Large Language Models,2023】)风格的提示。样本被随机扩展到32k token,使用第一阶段的对话,以保留长上下文能力。
在线策略RPO阶段。在第三阶段,我们应用在线策略RPO,使用第二阶段的检查点生成偏好和不偏好的补全。我们整合了来自AEGIS2.0(【索引23,AEGIS2. 0: A Diverse AI Safety Dataset and Risks Taxonomy for Alignment of LLM Guardrails,2025】)的安全数据,并增加了奖励缩放,这提高了下游基准的性能。
实验方法与结果。我们试验了多轮迭代的离线RPO和DPO。两者都产生了相似的结果;我们报告了来自任一方法的最佳模型的结果。最佳的Transformer模型使用了三轮RPO;最佳的混合模型使用了两轮DPO。
对齐研究结论。表10展示了结果。这些发现支持了混合Mamba-Transformer模型可以被后训练以达到或超过Transformer性能的结论。
Table 10 | 指令调优的Transformer和混合模型在对齐基准上的比较。
5.2. 训练Nemotron-H推理模型
训练流程启动。在初步研究中,我们发现mamba2-混合模型和纯transformer模型具有相似的后训练特性,因此我们着手训练我们的Nemotron-H 8B推理模型和Nemotron-H 47B推理模型。
5.2.1. 监督微调阶段
训练流程概述。训练流程始于监督微调(SFT),使用精心策划的示例,这些示例包含显式的推理轨迹——用 <think>...</think> 标签包裹——以引导模型在达到最终答案之前进行逐步解决问题。这些轨迹通常代表多种可能的解决方案路径,并鼓励模型探索替代方案和迭代,从而提高准确性。然而,增加的详细性也增加了推理成本,特别是对于较长的轨迹。为了平衡这一点,我们引入了去除了推理过程的配对示例,让模型学习何时以及如何直接回答。这种双重格式的训练有助于模型灵活地适应不同的推理要求。
第一阶段:掌握数学、代码和科学推理。微调的第一阶段专注于数学、科学和编码——这些领域中显式推理尤其有价值。此处的训练数据使用了5:1的推理与非推理样本比例,其中一些示例在Llama-Nemotron-Post-Training-Dataset中公开发布。
第二阶段:扩展指令覆盖、对话和安全性。第二阶段转向指令遵循、安全对齐和多轮对话——同时继续从第一阶段抽样以保持强大的STEM性能。这个数据集更紧凑——大约小10倍——并提供了推理和非推理样本的均衡混合。这有助于模型在更广泛的任务上进行泛化,同时提高对推理模式切换的控制。
5.2.2. 长上下文训练
长上下文训练方法。为了提升长上下文能力,我们对第二阶段混合数据的增强版本进行了额外训练,该版本包含每个长达512k token的样本。这个扩展混合数据中的样本是通过连接原始的第二阶段SFT对话,并增加几种额外类型的轮次来创建的,目的是提高模型的长程记忆。增加的轮次构建如下:
1. 对现有对话的长程引用。随机选择一个出现在拼接样本中较早位置的短对话;然后我们生成后续的用户查询,这些查询明确地按主题引用该轮次序列(例如,“在我们讨论现代作家的那部分对话中……”)。引用问题和答案都由DeepSeek-R1(【索引18,Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning,2025a】)生成。这些引用对被插入到数据混合的推理和非推理子集的对话中(见图12)。这种方法不需要依赖生成器模型的长上下文理解能力,因为我们只需要被引用的对话作为上下文。
2. 多轮引用。除了引用来自短混合数据的对话轮次外,我们还生成了多个具有统一主题的轮次,这些轮次可能相互引用。这些序列的主题被选为3个随机关键词的组合。然后,我们生成多个后续轮次,每个轮次引用这些片段中的一个或多个,要求模型确定正在引用哪个相关片段。所有这些轮次随后被随机插入到扩展对话中(见图13)。与前一类型类似,这也不需要长上下文生成器。
3. 文档聚合问答。用户的一轮对话中嵌入一个长文档(最多32k token),该文档来自预训练数据。该轮对话之后是几个问题,其答案需要整合散布在长文档中的信息。
4. 带干扰项的常用词检索。为了针对RULER的“常用词”子任务,我们插入了用户轮次,其中列出了30k-80k个token,包含一个关键词列表,并与随机干扰项交错。助理被要求返回频率最高的 $N$ 个词。这个专门的任务仅在非CoT子集(non-CoT subset)中使用。它是在没有使用语言模型的情况下创建的。
长上下文学习效果。我们发现,混合模型在这些512k-token的扩展混合数据上仅需训练数百步就能学会在更长的上下文长度上操作。我们使用RULER(【索引29,RULER: What’s the Real Context Size of Your Long-Context Language Models?,2024】)基准在非推理模式下评估了这一能力,模型在此阶段后在128k序列长度上达到了83%的RULER分数。
5.2.3. 使用GRPO进行强化学习
GRPO应用。在SFT之后,我们在多个阶段应用了组相对策略优化(Group Relative Policy Optimization, GRPO)(【索引89,Deepseekmath: Pushing the limits of mathematical reasoning in open language models,2024】)。每个阶段都针对特定技能——如指令遵循或工具使用——通过创建带有自动验证器的任务特定数据集,然后用一个通用的奖励模型进行更广泛的微调。
指令遵循调优。为了加强指令遵循能力,我们从LMSYS聊天数据集中抽样了16,000个提示,并将它们与IFEval风格的指令配对。一个基于规则的验证器根据输出满足每条指令的程度进行评分,从而创建了一个优先考虑精确遵循指令的奖励信号。我们使用了TRT-LLM后端进行策略展开。IFEval RL实验显著提升了IFEval能力,而其余基准测试结果略有波动,需要仔细选择检查点。我们使用了1e-6的学习率,这在指令遵循能力上显示出持续的改进,而没有在其他模型能力上出现退步。
通用帮助性调优。在最后的RL阶段,我们引入了一个基于Qwen-32B的奖励模型(在RewardBench上得分92.8)以提高整体响应的帮助性。我们使用GRPO和来自HelpSteer2(【索引104,Helpsteer2: Open-source dataset for training top-performing reward models,2024b】)的提示进行训练。这最后一个阶段导致了输出质量的显著提升。
5.2.4. 可控推理
推理时行为控制。推理时的行为可以通过在系统提示中使用简单的控制标签进行定制:
- {'reasoning': True} 触发推理模式
- {'reasoning': False} 触发直接回答模式
- 省略该标签让模型自行选择
实现机制。我们的Jinja聊天模板会检测这些控制字符串并相应地修改助手的响应。当存在 {'reasoning': True} 时,响应会以 Assistant:<think>\n 为前缀,表示推理轨迹的开始。当找到 {'reasoning': False} 时,响应会以 Assistant:<think></think> 为前缀,表示一个非推理响应。这种机制实现了对推理或非推理模式近乎100%的控制。
5.3. 结果
5.3.1. Nemotron-H-47B-Reasoning-128K
性能对比。在数学、编码、科学、工具使用和对话等基准测试中,Nemotron-H-47B-Reasoning-128K的准确率与Llama-Nemotron Super 49B V1.0相当或更优,并在所有非编码基准上优于Qwen3 32B。该模型支持所有线性层的后训练量化,能够在准确率损失最小的情况下实现高效部署。我们还提供了量化后的检查点及其相应结果,以展示其在实际应用中的有效性。
5.3.2. Nemotron-H-8B-Reasoning-128K
检查点合并。我们观察到,较小的8B模型在内部安全基准测试中表现不佳。为了解决这个问题,我们选择了检查点插值(【索引107,Model soups: averaging weights of multiple fine-tuned models improves accuracy without increasing inference time,2022】),即混入一个从更安全的、长上下文扩展检查点对齐的早期RL检查点。检查点插值通过线性插值模型权重来执行:$(1 - \alpha) \cdot \text{weights}_1 + \alpha \cdot \text{weights}_2$。我们对 $\alpha$ 值从0.1到0.9,增量为0.1进行了参数扫描,发现0.5左右的值在推理能力和安全行为之间提供了一个很好的权衡。
性能对比。在AIME25、科学、工具使用和指令遵循等基准测试中,Nemotron-H-8B-Reasoning-128K的准确率与Llama-Nemotron Nano 8B V1.0相当或更优。该模型支持所有线性层的后训练量化,能够在准确率损失最小的情况下实现高效部署。我们还提供了FP8量化检查点及其相应结果,以展示其在实际应用中的有效性。
Table 11 | 多个模型变体的基准性能比较。每行最佳分数以粗体显示。
Table 12 | 多个8B模型变体的基准性能比较。每行最佳分数以粗体显示。
A3 实验环境
-
模型架构:
- Nemotron-H-8B: 52层,模型维度4096,FFN维度21504,32个查询头,8个键值头,Mamba-2状态维度128。
- Nemotron-H-56B: 118层,模型维度8192,FFN维度32768,64个查询头,8个键值头,Mamba-2状态维度256。
- Nemotron-H-47B: 由56B模型通过MiniPuzzle压缩而来,保留了5个自注意力层,44个Mamba-2层和49个FFN层,FFN中间大小从32768降至30720。
- Nemotron-T-8B (Baseline): 32层纯Transformer架构,参数与Nemotron-H-8B对齐。
- VLM: 采用NVLM-D架构,视觉编码器为InternViT-300M-V2.5。
-
硬件配置:
- 预训练: 在6144个NVIDIA H100 GPU上进行,使用了DGX Cloud Resilience服务。
- 推理: 在NVIDIA H100 GPU上进行性能评测。
- 部署目标: NVIDIA GeForce RTX 5090 GPU(32GiB显存)是47B模型的目标部署平台。
-
软件配置:
- 代码实现: Hugging Face, NeMo, Megatron-LM。
- 依赖库: Transformer Engine (FP8支持), vLLM (推理基线), lm-evaluation-harness (评估框架), Math-Verify, NeMo-Skills, EvalPlus (评估工具), TRT-LLM (推理后端)。
-
数据集:
-
预训练: 包含高质量策划数据和合成数据,总计高达20万亿token。主要来源包括:
- Nemotron-CC1: 自行处理的6.3万亿token的英文Common Crawl数据。
- 数学数据: 基于OpenWebMath扩展,超过1000亿token。
- 代码数据: 专注于Python, C++, C, Java等流行语言。
- 学术数据: 来自书籍、期刊、专利等。
- 合成数据: 包括网络文本改写、数学对话、代码问题生成、SFT风格指令等,总计超过2万亿token。
-
压缩与蒸馏: 使用了70亿token进行短时蒸馏,630亿token进行长时蒸馏。
- VLM训练: 预训练使用1.3亿图文对,SFT使用600万图文对。
- 对齐与长上下文: SFT阶段使用600万样本,长上下文训练样本长度达512k token,RL阶段使用LMSYS Chat和HelpSteer2等数据集。
-
A4 实验结果
-
基础模型准确率 (表4, 表5):
- Nemotron-H-56B-Base: 在17个评估任务中,有9个任务上准确率超过Qwen-2.5-72B和Llama-3.1-70B,并在16个任务上优于Llama-3.1-70B。在MMLU-Pro上达到60.5%,GSM8k达到93.7%。
- Nemotron-H-8B-Base: 在15个评估任务中,有7个任务上准确率最高,与Qwen-2.5-7B性能相当。在常识理解任务上表现尤为突出。
- 混合 vs. Transformer: Nemotron-H-8B-Base与在相同数据上训练的纯Transformer模型Nemotron-T-8B-Base相比,准确率相当或更好,证明了混合架构在SOTA规模下的有效性。
-
推理速度 (图1, 图7):
- 长上下文场景: 在65536输入和1024输出的设置下,Nemotron-H-56B-Base的推理吞吐量是Qwen-2.5-72B和Llama-3.1-70B的2.4倍。Nemotron-H-8B-Base相比Qwen-2.5-7B和Llama-3.1-8B有1.8倍到3倍的加速。
- 加速来源: Mamba-2层的恒定计算复杂度和较低的内存占用是实现加速的关键。
-
FP8预训练效果 (图5, 图6):
- 尽管FP8训练的损失曲线略高于BF16,但在MMLU、代码、数学等下游任务上,FP8训练的模型准确率与BF16相当甚至更优。这证明了FP8作为一种节约成本的预训练方法的可行性。
-
模型压缩 (MiniPuzzle) 效果 (图1, 表4, 表6):
- Nemotron-H-47B-Base: 从56B模型压缩而来,在MMLU-Pro上准确率达到61.8%,略高于56B版本,同时推理速度提升了1.2倍。在HumanEval和MBPP等代码任务上甚至超过了原始56B模型。
- 高效压缩: MiniPuzzle仅用约300分之一的训练token就实现了近乎无损的压缩,并提升了效率。
-
视觉语言模型 (VLM) 性能 (表7, 表8):
- Nemotron-H-8B-VLM: 在MMMU (val)上达到51.3%,与基于Qwen2.5-7B的模型性能相当。
- Nemotron-H-56B-VLM: 在多个基准上达到SOTA水平,例如在MMMU (val)上达到63.6%,MathVista上达到70.7%,优于基于Qwen2.5-72B的VLM和先前的NVLM-1.0-D 72B。
-
推理模型 (对齐与长上下文) 性能 (表10, 表11, 表12):
- 对齐能力: 实验证明混合Mamba-Transformer模型与纯Transformer模型具有相似的后训练潜力,在MT-Bench、IFEval等对齐基准上表现相当或更优。
- Nemotron-H-47B-Reasoning-128K: 在GPQA-D、IFEVAL等基准上优于Llama-Nemotron Super V1,展现了强大的指令遵循和推理能力。
- 长上下文: 经过特定训练后,模型在RULER基准(128k序列长度)上取得了83%的高分。
A5 结论
本文介绍了Nemotron-H模型家族,证明了混合模型架构能够兼具SOTA Transformer模型的能力和更高的推理效率。实验结果表明,Nemotron-H基础模型在准确性上与Qwen-2.5-7B/72B和Llama-3.1-8B/70B等同等规模的开源Transformer模型相当或更优,而推理速度最高可提升3倍。
此外,本文还验证了这些基础模型可以被有效地适配,生成性能强大的视觉语言模型(VLM)和指令调优模型。本文提出的FP8训练方法和MiniPuzzle压缩框架能够显著降低创建这些高性能模型的成本。
为了促进该领域的进一步研究,我们发布了Nemotron-H-56B-Base、Nemotron-H-47B-Base、Nemotron-H-8B-Base以及Nemotron-H-47B-Reasoning和Nemotron-H-8B-Reasoning的模型检查点,并提供了在Hugging Face和NeMo中的支持。
💬 评论讨论
欢迎在这里分享您的想法和见解!