英伟达正式开源了其不久前在 AI 数学奥林匹克竞赛（AIMO，AI Mathematical Olympiad）中斩获冠军的核心模型系列。

在本届 AIMO-2 Kaggle 竞赛中，超过 2,200 支参赛队伍提交了 AI 模型，挑战在 5 小时内解决 50 道国家奥林匹克级别的复杂数学问题。英伟达的 7 人团队“NemoSkills”最终正确解答了 34 道题目（相比 2024 年的冠军提高了 5 道），夺得了冠军。

图丨此次比赛的排行榜（来源：Kaggle）

现在，英伟达向全球开放了帮助他们获胜的核心技术，包括小参数的 OpenMath-Nemotron-1.5B、OpenMath-Nemotron-7B 和直接用于竞赛并优化的 OpenMath-Nemotron-14B-Kaggle 模型、性能更为强大的旗舰模型 OpenMath-Nemotron-32B，以及训练它们所依赖的 OpenMathReasoning 数据集。

基准测试的结果显示，这几款模型表现出色，在 AIME 和 HMMT 竞赛中数学问题上的准确率全面超越了 14B 的 DeepSeek-R1。

图丨 AIME 和 HMMT 竞赛中数学问题的准确率（来源：arXiv）

英伟达是如何构建 OpenMath-Nemotron 的？

那么，英伟达是如何构建 OpenMath-Nemotron 模型的？

这首先在于一个大规模且高质量的训练数据集。认识到现有资源的不足，英伟达团队首先投入了大量的工作来创建 OpenMathReasoning 数据集。

他们先从“Art of Problem Solving（AoPS）”等在线数学社区收集了大量的原始数学问题和讨论。

随后，团队利用 Qwen2.5-32B-Instruct 开发了一套自动化流程，对这些原始数据进行细致处理。这包括从帖子中提取完整的数学问题，对问题进行分类（例如，剔除选择题和是非题），并将一些需要证明过程的问题巧妙地转化为需要具体答案的形式，以便于模型训练和自动评估。同时，为了保证模型的泛化能力，他们还进行了基准去污染处理，移除了与现有常见数学测试集（如 MATH、GSM8K）中题目过于相似的问题。

最终完成的 OpenMathReasoning 数据集，包含了 54 万个高质量数学问题，其中涵盖了从中学到奥林匹克竞赛等不同难度级别。为了让模型学会“思考过程”，团队更进一步地利用 DeepSeek-R1 和 QwQ-32B 等强大的现有模型，为这些问题生成了 320 万条包含详细解题步骤的“思维链”（CoT，Chain-of-Thought）解决方案。

图丨数据集组成（来源：arXiv）

第二个核心部分是工具集成推理。

现代 AI 研究的一个重要趋势是让语言模型学会使用外部工具，例如调用计算器或执行代码片段，来辅助解决问题，尤其是在需要精确计算或模拟的场景下。然而，团队在实践中发现，即便是当时最强的开源数学模型，也难以通过简单的提示工程来引导它们生成高质量的、将代码执行与自然语言推理深度融合的解决方案（即 TIR）。这些模型似乎对其自身固有的纯文本推理模式产生了某种“路径依赖”。

为了克服这一障碍，NemoSkills 团队设计并实施了一套迭代式开发流程。他们首先选择了一个指令遵循能力较好的基础模型（LIMO-Qwen-32B），用少量推理数据对其进行初步微调。然后，引导这个模型生成第一批包含 Python 代码的 TIR 解决方案。关键的下一步是进行严格的质量过滤：利用另一个强大的大模型（ Qwen2.5-32B-Instruct），来判断每个代码块的“新颖性”（是产生了新结果还是仅仅验证已知步骤）和“重要性”（是解决问题的关键环节还是可以被几步简单 CoT 取代）。只有那些代码执行提供了显著推理价值（而非冗余计算）的样本才被保留下来，形成了约 1.5 万个样本的初始 TIR 训练集。

接下来，他们用这个高质量的初始集去微调更强大的模型（如 QwQ-32B），使其初步具备生成 TIR 的能力。随后，利用这个微调后的模型生成更多、更高质量的 TIR 数据，并再次运用上述过滤标准进行筛选。这个“生成-过滤-训练”的闭环被重复执行，每一轮都提升了 TIR 数据的规模和质量。最终，团队构建起了一个包含 170 万条高质量 TIR 解决方案的数据集。基于此训练出的 OpenMath-Nemotron 模型，能够熟练地在自然语言推理中嵌入 Python 代码执行，从而攻克那些纯文本推理难以解决的复杂计算问题。此外，他们还设计了一种机制，使得模型在生成答案时能够遵循对代码块使用次数的限制，这对于资源受限的推理场景至关重要。

第三个核心部分则是团队提出的生成式解决方案选择。

在解决困难问题时，让模型生成多个候选答案并从中择优，是提升最终准确率的常用技巧。传统的“多数投票”方法虽然直观，但往往无法充分发掘模型生成的所有答案中的潜在正确信息，其性能通常远低于理论上的“pass@k”（即 k 个答案中至少有一个正确的概率）上限。

为了弥补这一差距，英伟达团队开发了 GenSelect 技术。其核心思想不再是简单地对最终答案进行投票，而是训练一个模型，让它扮演“评审员”的角色，能够“阅读”并“理解”多个候选解决方案的完整摘要，然后基于对解题逻辑、步骤合理性等的判断，选出最可信、最有可能正确的那一个。

图丨 GenSelect 的数据构建流程（来源：arXiv）

具体来说，团队首先利用 Qwen2.5-32B-Instruct 模型为 OpenMathReasoning 数据集中所有已生成的 CoT 和 TIR 解决方案重新生成了结构化的、信息更丰富的摘要。然后，他们构建了 GenSelect 的训练数据：为每个原始问题，随机抽取 2 到 16 个候选方案的摘要（特别设计以确保样本组中至少包含一个正确和一个错误的解），将这些摘要连同原问题一起输入给 QwQ-32B 模型，并要求它生成一段详细的比较分析文本，最终明确指出哪个索引号的解决方案是最佳的。通过筛选掉那些模型判断错误（即选择了错误答案）的案例，他们构建了一个包含 56.6 万个样本的 GenSelect 训练数据集。

实验结果表明，经过 GenSelect 加持的模型，其最终准确率相比简单的多数投票有了显著提升，尤其是在候选方案数量不多时效果更为明显。虽然由于 AIMO 竞赛严格的时间和计算限制，GenSelect 未能被纳入最终的获胜提交方案中，但这项技术已被完全整合到此次发布的 OpenMath-Nemotron-32B 模型中，构成了其支持的三大推理模式之一。

基于上述三大支柱和海量数据，英伟达团队训练了一系列名为 OpenMath-Nemotron 的模型，参数规模涵盖 1.5B、7B、14B 和 32B。这些模型均基于强大的 Qwen2.5 基座模型进行微调。对于 1.5B 和 7B 版本，他们甚至使用了专门为数学任务优化的 Qwen2.5-Math 版本作为起点。

训练过程采用了监督微调，混合使用了 CoT、TIR 和 GenSelect 三种任务的数据，总计达 550 万个样本。这意味着同一个模型可以通过不同的提示（prompt）在 CoT（纯文本推理）、TIR（工具集成推理）和 GenSelect（多方案选择）模式下工作。

图丨训练过程中准确率的提升（来源：arXiv）

为了处理长达数千甚至上万个 token 的长序列推理，团队应用了旋转位置编码（RoPE，Rotary Position Embedding）扩展技术，并将训练过程中的上下文窗口扩展到支持长序列。训练使用了英伟达自家的 NeMo-Aligner 工具包，并结合了序列打包、上下文并行等技术来加速长序列训练。此外，他们还采用了检查点平均（checkpoint averaging）和在更难问题子集上进行第二轮微调等策略，进一步提升模型性能。

多项优化推理措施

赢得 AIMO-2 竞赛不仅需要模型本身强大，还需要在极其苛刻的 5 小时、4x L4 GPU 限制下高效完成推理。这要求团队在模型选择和推理优化上做出极致权衡。

他们的最终提交方案基于 OpenMath-Nemotron-14B 模型的一个早期版本，该版本在一个稍小的 CoT 数据集（仅 DeepSeek-R1 生成）上训练，并进行了轻量级的 TIR 微调。值得注意的是，他们采用了模型合并技术，将纯 CoT 训练的检查点和经过 TIR 微调的检查点进行线性组合。这种简单而有效的方法，让他们能够在保持 TIR 能力的同时，部分恢复 CoT 模型的生成流畅性和速度优势，并减少代码调用次数，从而更好地适应竞赛环境。

为了在有限的时间内最大化解题数量和准确率，团队实施了多项推理优化措施：

首先，他们使用 TensorRT-LLM 将预训练模型转换为 TensorRT 引擎。这一工具的动态批处理功能通过动态分组推理请求提高了吞吐量，在样本完成后即刻释放，减少延迟并优化 GPU 利用率。由于样本是独立处理的，批处理可以无缝混合不同的提示或推理参数。TensorRT-LLM 还包括自定义注意力内核和分页 KV 缓存等多种优化。

在量化方面，团队优先采用 int8 权重量化（W8A16）和 FP8 量化，相比 BF16 格式速度提升了 1.5 倍，同时对准确率的影响最小。减小的权重大小还为更大的键值缓存释放了内存，允许处理更长的序列。团队还使用了苹果开发的 ReDrafter 技术，这是一种循环推测解码方法，使用基于 RNN 的起草器在每个解码步骤提出并验证多个 token。他们训练了一个能够在每一步提出最多三个 token 的起草器，在大约 65% 的步骤中成功接受所有三个 token，显著加速了生成过程。

图丨在 4 个 L4 GPU 上对具有不同优化方法的提交管道进行基准测试（来源：arXiv）

此外，团队通过将 CoT 和 TIR 检查点线性组合创建了最终模型，这种策略允许他们控制每个微调阶段对最终模型行为的影响程度。最佳模型是使用 CoT0.3+TIR0.7 的组合创建的，这不仅提高了准确率，还通过减少解决方案长度和代码执行次数加速了生成。团队实现了一种缓冲策略，为每个问题分配 350 秒的基本时间限制，如果一个问题提前完成，未使用的时间会被添加到共享缓冲区，供后续问题使用。

团队还利用了 NeMo-Skills 的异步生成功能实现批量处理和早停。例如，在 16 个样本的批处理中，如果前 4-5 个完成的样本就已经对最终答案达成一致，则取消剩余的生成并继续下一个问题。这种机制极大地节约了在简单或中等难度问题上可能浪费的时间，为攻克难题争取了宝贵的时间窗口。早停策略增加了响应相关性，因为较短的答案往往质量更高。

图丨异步批处理流程（来源：Kaggle）

实验结果显示，在 Comp-Math-24-25 测试集（包含来自 AIME 和 HMMT 竞赛的问题）上，团队的模型表现出色。1.5B 模型在 CoT 模式下单次通过准确率为 58.2%，多数投票准确率达 80.0%；在 TIR 模式下，这些数字分别提高到 64.5% 和 83.3%；使用 GenSelect 技术后，准确率进一步提升至 83.3%。14B 模型的表现更为出色，在 TIR 模式结合 GenSelect 使用时，准确率高达 90.0%。最大的 32B 模型在相同条件下甚至达到了 93.3% 的准确率。这些结果也表明，无论模型大小如何，TIR 模式始终优于纯 CoT 模式，而 GenSelect 技术能进一步提高准确率。

图｜数学基准测试的评估结果（来源：arXiv）

目前，英伟达团队已将完整的 OpenMathReasoning 数据集、训练好的 OpenMath-Nemotron 模型系列以及所有相关代码以商业许可方式发布到 Hugging Face 和 GitHub 上（项目地址：https://huggingface.co/collections/nvidia/openmathreasoning-68072c0154a5099573d2e730）。

参考资料：

1.https://arxiv.org/abs/2504.16891

2.https://www.kaggle.com/competitions/ai-mathematical-olympiad-progress-prize-2/discussion/574765

运营/排版：何晨龙