NVIDIA Llama-Nemotron:高效推理LLM开放家族

在人工智能浪潮席卷全球的今天,大型语言模型的研发突飞猛进。其中,推理能力作为衡量模型智能化程度的核心标尺,已然成为各大AI巨头竞相角逐的技术制高点。

然而,随着模型体量的不断增大,推理效率逐渐显现为制约模型实际部署与性能发挥的关键瓶颈。

洞察及此,NVIDIA基于Meta AI的Llama模型家族,推出了Llama-Nemotron系列模型。这是一个专为高效推理而生的大型模型开放家族,不仅在推理能力和效率上表现卓越,更采用了对企业用户十分友好的开放许可协议。

Llama-Nemotron系列目前包含三个不同规模的模型:

  • Nano (80亿参数)
  • Super (490亿参数)
  • Ultra (2530亿参数)
    此外,还有一个专为处理超长上下文设计的独立变体 UltraLong (80亿参数)。

image.png

这一系列模型的推出可谓意义重大。它们不仅拥有出色的推理硬实力,还为企业应用敞开了大门,提供了开放许可。模型的权重和部分训练数据均已在Hugging Face平台公开,遵循NVIDIA Open Model License以及Llama社区许可,允许商业用途。

值得一提的是,Llama-Nemotron系列模型是首批支持动态推理模式切换的开源模型。这意味着用户在进行推理交互时,可以根据需要在标准的聊天模式与专门的推理模式之间灵活切换,极大地增强了模型应用的灵活性和场景适应性。

研究团队主要通过推理类和非推理类两大基准测试来全面评估Llama-Nemotron系列模型的性能。结果显示,该系列模型在各种参数规模下均展现出优异的性能表现。特别是LN-Ultra模型,相较于DeepSeek-R1等先进模型,在推理吞吐量和部署效率方面实现了显著提升。

image.png

Llama-Nemotron的卓越性能得益于其精心设计的多阶段后训练流程,该流程旨在同时强化模型在推理任务和非推理任务上的表现。

  • 监督微调 (SFT) 阶段:重点关注数学、代码生成、逻辑推理以及工具调用等能力的培养。
  • 强化学习 (RL) 阶段:采用了REINFORCE算法(具体实现为RLOO)以及支持在线奖励感知偏好优化的RPO (Online Reward-aware Preference Optimization) 方法,用以优化模型的对话生成质量和指令遵循的精准度。

在Llama-Nemotron的训练过程中,QwenDeepSeek-R1这两大模型扮演了不可或缺的关键角色。

  • Qwen (例如Qwen2.5-32B-Instruct) 主要负责数学和科学领域相关数据的生成、分类以及去污染处理,为构建高质量的训练数据集奠定了坚实基础。
  • DeepSeek-R1 则作为核心的“教师模型”,负责生成包含多步推理过程和代码解决方案的优质数据,并通过监督微调与强化学习等手段,将其深厚的逻辑推理能力有效地迁移至目标模型(Llama-Nemotron)。

那么,NVIDIA究竟是如何打造出Llama-Nemotron这一系列强大的模型的呢?其背后又蕴藏着哪些独特的训练秘辛?

接下来,让我们一同深入剖析其技术内核。

构建面向推理优化的模型架构

LN-Super和LN-Ultra这两款大规模模型,其高效推理能力的实现,关键在于采用了Puzzle框架。Puzzle是一个先进的神经网络架构搜索 (Neural Architecture Search, NAS) 框架,它能够在真实的部署环境约束下,将现有的大型语言模型巧妙地转化为针对特定硬件优化的、更为高效的变体。具体流程如图3所示。

image.png

以Llama 3 Instruct模型作为优化的起点(具体而言,LN-Super基于Llama 3.3-70B-Instruct,而LN-Ultra则基于Llama 3.1-405B-Instruct),Puzzle框架通过一种逐模块局部蒸馏 (block-wise local distillation) 的方法,构建了一个包含多种可替代Transformer模块的“模块库”。库中的每一个模块都经过独立训练,并且可以并行处理,其目标是在尽可能接近原始模块功能的同时,显著提升计算性能。

这种方法的精妙之处在于,它允许每个替代模块在精度与效率之间进行权衡。也就是说,模块库中的某些变体可能拥有更高的计算效率,但这通常会伴随着一定程度的准确性下降,从而形成一种明确的“精度-效率权衡”空间。这些模块变体主要包括以下几种类型:

  • 移除注意力机制 (Attention removal):部分模块彻底移除了注意力机制,这样做能够显著减少计算开销,并大幅降低KV缓存 (Key-Value cache) 的内存占用。
  • 可变前馈网络维度 (Variable FFN dimensions):通过灵活调整前馈网络 (Feed-Forward Network, FFN) 的中间层维度,可以在不同粒度上实现模型压缩(例如,将隐藏层维度压缩至原始尺寸的87%、75%、50%,甚至极端情况下的10%)。

尽管Puzzle框架同样支持其他类型的结构替换方案(例如,在多组查询注意力机制 (Grouped-Query Attention, GQA) 中调整键值头的数量、采用线性注意力替代方案,或是引入不执行任何操作的“空操作”替换模块),但实际评估结果清晰地表明,在优化LN-Super和LN-Ultra这两款模型的整体吞吐量与内存节省方面,移除注意力机制FFN压缩依然是最为行之有效的技术手段。

在模块库构建完毕之后,Puzzle框架便通过逐层选取模块的方式来组装最终的完整模型。模块的选择过程由一个整数混合规划 (Mixed-Integer Programming, MIP) 求解器精确控制。该求解器会根据预设的约束条件(例如硬件兼容性要求、最大可接受的推理延迟、总内存预算限制,或是特定的推理吞吐量目标),自动确定能够达到最高效率的模块配置方案。

由于模型的每一层都支持多个具有不同“精度-效率”权衡方案的模块变体,Puzzle框架使得用户能够精准地定位到位于“精度-效率帕累托前沿 (Pareto frontier)”上的任意模型配置点。举例来说,Puzzle可以生成完全满足特定智能体系统 (agentic systems) 或部署流程所提出的严格约束(如内存占用不得超出上限,或端到端响应时间必须严格控制在某一阈值内)的模型。

FFN融合:实现模型的纵向压缩

针对参数规模更为庞大的LN-Ultra模型,研究人员还引入了一种额外的压缩技术——FFN融合 (FFN Fusion)。该方法的核心目标是降低模型的有效序列深度,从而进一步缩短推理过程中的延迟。

FFN融合巧妙地利用了Puzzle框架在移除部分注意力层后所形成的特殊结构特性:在这种结构下,模型中经常会出现连续的FFN模块序列。FFN融合技术会精准识别出这类连续的FFN序列,并将其替换为数量更少但“宽度”更大的FFN层。这些“更宽”的FFN层可以并行执行,从而在保持模型表达能力的同时,有效减少了序列化处理的步骤数量。

更进一步地,这种融合方式显著提升了计算资源的利用效率,尤其是在多GPU并行计算的环境中,能够有效地降低因跨层通信所带来的开销。

部署约束与效率目标

LN-Super 的设计目标是在单块NVIDIA H100 GPU上实现高效运行,其配置采用了张量并行系数为1 (Tensor Parallelism 1, TP1)的方案。通过Puzzle框架的深度优化,该模型在批量大小 (batch size) 为256、TP1配置的条件下,相较于其基线模型Llama 3.3-70B-Instruct,实现了高达5倍的推理吞吐量提升。即便Llama 3.3-70B-Instruct采用其最优配置(即张量并行度为4,TP4),LN-Super在TP1条件下依然能够保持至少2.17倍的吞吐量优势。

LN-Super的设计还需满足约30万个缓存Token (cached tokens) 的运行约束(缓存Token数量 = 批量大小 × 序列长度),这一数据是在FP8精度下,基于单张H100 GPU测得的。例如,批量大小为16、序列长度为18,750的配置便能满足此缓存量的要求。

LN-Ultra 的优化目标则更为宏大,瞄准了整个H100节点(包含8张GPU)。在Puzzle框架的结构搜索阶段,模型被施加了推理延迟必须比Llama 3.1-405B-Instruct至少缩短1.5倍的约束。在应用了FFN融合技术之后,最终得到的模型在延迟方面实现了1.71倍的显著提升

与LN-Super类似,LN-Ultra同样受到缓存Token数量的限制:

  • 在FP8精度下,最多支持300万个Token。
  • 在BF16精度下,支持60万个Token。
    这些数据均以整个H100节点作为计算基准。

下图(图4)清晰地展示了在两种不同设置下,GPQA-Diamond基准测试的准确率(以百分比计)与处理吞吐量(以Token/秒计)之间的权衡曲线。值得注意的是,LN-Ultra在准确率和效率两个维度上均优于DeepSeek-R1和Llama 3.1-405B,这充分表明在“精度-吞吐率帕累托曲线”上,LN-Ultra是一个更具竞争力的选择。

image.png

NAS后训练阶段:知识蒸馏与持续预训练的精进

在神经网络架构搜索 (NAS) 阶段圆满结束后,为了进一步提升不同模块间的兼容性,并弥补因模块替换可能带来的性能损失,LN-Super和LN-Ultra两款模型均接受了进一步的精细化训练。

  • LN-Super:使用了Bercovich等人提出的Distillation Mix数据集,并采用知识蒸馏的目标函数,进行了长达400亿个Token的训练。
  • LN-Ultra:首先同样使用Distillation Mix数据集进行了650亿Token的蒸馏训练;随后,在Nemotron-H第四阶段的预训练数据集上,额外进行了880亿Token的持续预训练。

通过这最后阶段的预训练,LN-Ultra不仅成功地在性能上追平了其基准模型Llama 3.1-405B-Instruct,更在多个关键的基准测试中实现了超越。这有力地证明了,即使在进行了较为激进的架构优化之后,通过短周期的知识蒸馏与持续预训练,依然可以有效恢复乃至提升模型的综合性能(详见表1)。

image.png

推理能力强化学习:赋予模型灵活的思考模式

为了使模型能够在不同的任务场景下,灵活地切换其推理的深度与回答的风格,研究人员精心设计了一种**「detailed thinking on/off」(详细思考开启/关闭)指令机制**。通过在合成的训练数据中明确标记是否需要模型展开详细的推理过程,从而引导模型在训练中学习何时应当进行逐步思考、清晰展示推理链条,又何时可以直接给出简洁明了的答案。

具体来说:

  • 当指令为「on」时,模型被期望输出完整的中间推理步骤,并清晰地展示其解题思路。
  • 当指令为「off」时,模型则仅需呈现最终的结论或答案。

这种机制的引入,显著提升了模型对用户指令响应的可控性,同时也增强了其推理行为在不同应用场景中的适应能力,使得模型能够根据实际需求智能地调整输出内容的详略程度。

在此基础上,模型通过监督微调 (SFT) 从“教师模型”那里学习复杂的多步推理路径,并有效地将推理能力与通用的任务处理风格相融合,最终构建出一个兼具高推理精度与出色使用灵活性的智能响应系统。

如表5所示,LN-Ultra在推理类与非推理类的各项基准测试中,均达到甚至超越了当前顶尖的开源权重模型的水平。这充分证明了通过从强大的教师模型中进行知识蒸馏,模型能够通过监督微调获得非常强大的综合能力。

image.png

然而,知识蒸馏在本质上为“学生模型”设定了一个性能上限,尤其是在学生模型自身能力尚未超越教师模型的情况下。

例如,通过监督微调,LN-Ultra虽然可以逼近DeepSeek-R1的性能水平,但要实现超越则颇具挑战。为了让学生模型有机会突破教师模型的限制,大规模强化学习 (RL) 提供了一条可行的路径,因为它能够促使模型持续探索新的策略空间,并从中学习和提升。

研究人员的初步实验表明,在参数规模较小的模型上应用强化学习,其性能提升通常不如直接进行知识蒸馏来得显著。考虑到计算资源的限制,研究团队决定仅对LN-Ultra应用针对推理能力的强化学习训练,目标是打造出一个能够超越其教师模型的最终版本。

精心设计的训练流程

针对LN-Ultra模型,研究人员运用GRPO (Generalized Reward-Parameterized Optimization) 算法,通过大规模强化学习来着力提升其在科学推理任务上的能力。在训练过程中,每个rollout(即一次完整的交互序列)的提示词长度被设定为72个token,并且针对每个提示词,模型会采样生成16个不同的响应,采样参数设置为temperature = 1,top_p = 1。

全局批量大小 (global batch size) 设置为576。在每个rollout之后,会进行两次梯度更新。整个训练过程持续进行,直至模型在推理任务上的性能达到收敛状态。下图(图5)展示了模型在GPQA-Diamond基准测试上的准确率随着训练进展而发生的变化。得益于优化后的高效训练基础设施,整个训练过程总共消耗了约14万个H100 GPU小时。

image.png

在本阶段的训练中,主要使用了以下两类奖励信号:

  1. 准确率奖励 (Accuracy rewards):每个训练样本都提供了标准的参考答案(可能是一个数字、一个句子或一段文字)。研究人员利用Llama-3.3-70B-Instruct模型来判断策略模型生成的响应是否与标准答案一致,并据此给出奖励。
  2. 格式奖励 (Format rewards):遵循DeepSeek-AI等研究团队的做法,当模型处于“详细思考开启 (detailed thinking on)”模式时,其推理过程必须被包含在<thinking></thinking>标签之内;而在“详细思考关闭 (detailed thinking off)”模式下,则要确保输出中不包含这类思考标签。格式奖励旨在确保模型能够按照预设的规范格式来输出其推理过程。

为了进一步增加训练的挑战性,研究人员对训练数据进行了预处理:首先,由LN-Super模型为数据集中的每道题目生成8个独立的回答;然后,计算这些回答的通过率 (pass rate),并筛选掉那些通过率高于或等于0.75的相对简单的样本,从而有效提升了整体训练数据的难度。

除了巧妙的数据筛选策略外,研究人员还发现,课程化学习 (curriculum learning) 策略能够显著帮助模型在复杂的推理问题上实现更好的收敛和泛化能力。他们采用了一种渐进式批处理策略 (progressive batching)。具体做法是:利用预先计算得到的样本通过率作为衡量样本难度的指标,在保持固定批量大小的前提下,动态地调整每个批次中目标难度的分布。

这个难度分布采用高斯函数进行建模,其均值会从训练早期的集中在高通过率(对应简单样本)区域,逐渐平滑过渡到训练后期的集中在低通过率(对应高难度样本)区域。在每个批次的构建过程中,样本会根据当前的目标难度分布被随机分配,同时也会根据不同通过率区间的样本池中剩余的样本数量进行动态的容量填充。

这种精心设计的策略确保了样本的难度在批次 (batch) 的层面上是逐步递进的,同时又能在每个批次内部保持样本的随机性。下图(图6)直观地展示了这种课程化学习策略在降低训练方差、稳定训练过程以及提升最终模型准确率方面的显著效果。

image.png

FP8精度生成阶段的优化

研究人员敏锐地识别出,生成阶段是推理过程中主要的性能瓶颈所在。为了显著提升模型在该阶段的性能表现,他们成功开发了一条支持在vLLM框架下实现在线FP8精度生成模式的技术路径。该模式能够在FP8精度下执行所有的矩阵乘法 (GEMM) 操作,并巧妙地结合了针对每个token激活值的缩放因子以及针对每个张量权重的缩放因子。

为了与训练时输出的BF16精度权重相配合,研究人员还特别开发了一个自定义的vLLM权重加载器。这个加载器能够在模型运行时,将BF16格式的权重动态转换为FP8格式及其对应的缩放参数。由于vLLM框架当前尚不支持直接以FP8模式初始化模型,研究人员为此实现了元权重张量初始化 (meta-weight tensor initialization) 的机制。这一机制避免了因需要载入完整的BF16推理引擎而可能导致的GPU显存溢出问题。

通过上述一系列优化措施,在FP8模式下,单个GPU针对每个prompt的生成吞吐量最高可达32 token/秒,相较于传统的BF16精度模式,实现了高达1.8倍的性能提升。其中,FP8精度本身带来了约1.4倍的加速效果,另外的0.4倍性能增益则源于内存占用的显著减少。内存占用的降低使得研究人员能够启用vLLM框架中的cudagraph特性,从而进一步优化了整个系统的运行效率。

用于偏好优化的强化学习:让模型更懂用户

指令跟随能力的精细打磨

在完成了针对科学推理任务的强化学习训练之后,研究团队紧接着对LN-Super和LN-Ultra两款模型开展了短周期的强化学习训练,其核心目标是进一步优化它们的指令跟随能力。参照Zhou等人提出的验证方案,研究人员生成了一批包含1至10条详细指令的合成提示词,专门用于此阶段的训练。

在此阶段,研究者们采用了RLOO (Reinforcement Learning from Offline data and Online Optimization) 算法,进行了不超过120步的强化学习训练。他们使用了一个自定义的指令跟随验证器作为奖励函数,训练时的批量大小设置为128条提示。实验结果表明,这类针对性的训练不仅显著提升了模型在传统指令跟随评测任务中的表现,同时也对模型在推理类基准测试任务上的性能产生了积极的正面影响。

基于人类反馈的强化学习 (RLHF):提升模型的实用性和交互性

为了进一步增强模型在实际应用中的通用协助能力 (helpfulness) 与多轮对话的流畅性,同时确保其在其他各项任务上的已有性能不被削弱,研究人员采用了基于人类反馈的强化学习 (RLHF) 技术。

如表4所示,LN-Super(490亿参数规模)在极具挑战性的Arena Hard评测中取得了高达88.3的优异分数。这一成绩不仅超越了包括Claude 3.5 Sonnet和GPT-4o-2024-05-13在内的数个知名专有模型,也优于Llama-3.1-405B-Instruct和Mistral-large-2407等更大规模的开源模型。

image.png

为达成这一卓越表现,研究人员采用了一种迭代式的在线RPO (Online Reward-Parameterized Optimization) 训练方法。他们在HelpSteer2数据集提供的提示语上,以最大化由Llama-3.1-Nemotron-70B-Reward模型预测的偏好奖励为目标进行优化。

具体的训练参数设置如下:

  • 学习率 (α): 4e-7
  • KL散度惩罚项系数 (β): 1e-5
  • 奖励缩放因子 (η): 3.0
  • 批量大小 (batch size): 64
  • 训练步数: 500步

经过两轮在线RPO训练后,模型在Arena Hard评测上的分数从最初的69.1分大幅提升至88.1分。

尤为值得关注的是,这一RLHF过程几乎在所有基准测试任务中都带来了性能的提升,仅在IFEval数据集上观察到轻微的性能下降。研究人员推测,由于该数据集以及所使用的奖励模型并未专门针对数学、代码生成、科学推理或函数调用等特定场景进行设计和优化,RLHF过程可能更多地帮助模型更好地整合和调动其已有的知识储备和技能,从而在更广泛的通用能力上表现更佳。

针对参数规模更大的LN-Ultra,研究人员延续了上述的训练流程,但算法上选用了GRPO。对于每条提示词,模型会生成8个样本响应。训练参数设置为:学习率3e-7,批量大小288,KL散度惩罚项系数β为1e-3,共进行30步训练。

对于系列中参数规模较小的LN-Nano模型,研究人员则进行了两轮离线RPO (Offline RPO) 训练,使用的是策略内数据 (on-policy data)。

  • 第一轮训练:混合使用了包含推理内容和非推理内容的数据,并配合相应的系统提示,其主要目的是提升模型对推理任务的控制能力。
  • 第二轮训练:则聚焦于提升模型的指令跟随表现,训练数据主要来源于模型自身生成的策略内响应。
    每轮训练最多进行400步,学习率α设置为7e-7,KL散度惩罚项系数β为3e-2,批量大小为512。

更详尽的技术细节和实验结果,敬请参阅原论文。