第一名解决方案 - NemoSkills

我们要感谢 AIMOPrize、XTX Markets 和 Kaggle 的主办方举办了今年的挑战赛，显然他们在问题和 Kaggle 环境上花费了大量的思考和准备。特别感谢 Simon Frieder 在整个挑战过程中如此平易近人。

我们对结果感到非常兴奋——在开发了我们自己的 TIR 模型后，我们很难让公共榜单 (Public LB) 与强大的交叉验证 (CV) 结果相匹配，但最终私有榜单 (Private LB) 与我们的 CV 结果更加一致。

资源

• 论文
• Kaggle 提交脚本
• 训练数据集 - nvidia/OpenMathReasoning - 550 万 OpenMath 解决方案 - CoT & TIR
• 数学推理模型集合 - released 1.5B/7B/14B/32B - 1.5B 在 AIME 上击败 R1 (Kaggle 上传进行中)
• 提交的模型
• 在 Nemo-Skills 中复现结果

摘要 (TL;DR)

数据集是通过收集和过滤 AoPS 论坛上的 54 万个数学问题构建的。对于我们的提交，我们使用 DeepSeek-R1 等模型生成了 220 万个思维链 (CoT) 响应，并通过验证过滤掉错误答案。我们还创建了一个包含 1.5 万个问题和生成的高质量工具集成推理 (TIR) 数据集。
训练分两个阶段进行：首先在一个 CoT 数据上微调 Qwen2.5-14B 模型，其次使用较小的 TIR 数据集进一步细化，增强其结构化推理的问题解决能力。
为了保留 TIR 的优势以及 CoT 的较短响应速度，我们发现 mergekit 能很好地将两个模型合并为一个，既保留了 TIR 的优势，又减少了输出 token 的数量。
对于推理优化，我们实施了一个结合 TensorRT-LLM 转换和 FP8 量化的管道，实现了比 BF16 快 1.5 倍的速度。Apple 开发的投机解码技术 ReDrafter 被用于进一步实现 1.8 倍的加速。我们的服务系统具有动态批处理、早期停止和时间缓冲策略，可有效分配计算资源。该解决方案包括通过沙盒执行环境实现的工具集成推理能力。

训练数据

大规模长推理数据集

我们首先从互联网上收集了大量数学问题。我们开发了一个基于 LLM 的问题提取和细化管道，构建了一个包含 54 万个独特问题的数据集。使用该数据集，我们生成了 320 万个长推理 CoT 解决方案。该过程使用 DeepSeek-R1 和 QwQ-32B 等模型为每个问题生成多个解决方案候选（最多 32 个），温度为 0.7，top-p 为 0.95。更难的问题——由 Qwen2.5-72B-Math-Instruct 的通过率识别——会获得更多的候选方案。 incorrect 解决方案通过验证与 Qwen2.5-32B-Instruct 的答案等价性进行过滤。如果找不到答案，则使用最频繁的候选答案。对于我们的提交，使用了由 DeepSeek-R1 生成的 220 万子集。

工具集成推理 (Tool-Integrated Reasoning)

对于工具集成推理，模型提示代码进行计算，然后代码在沙盒中执行。我们使用特殊 token <tool_call> 和 </tool_call> 识别代码片段。然后将代码附加到 LLM 输出中，位于文本 ``` 和 ```output 之间。以下是输出中的一个示例片段。

Then, k must divide 56, and k > 16. So, k can be 28 or 56. Therefore, b = k - 7 = 21 or 49. So, same result. Therefore, sum is 70.\n\nAlternatively, maybe I can write a small program to check for all bases b > 9, compute 9b + 7 and b + 7, check if the latter divides the former, and collect all such bases. Then sum them. Let's do that to verify.\n\nHere's a Python code to perform the check:\n\n<tool_call>\n# Initialize a list to store valid bases\nvalid_bases = []\n\n# Check bases from 10 upwards\nfor b in range(10, 10000):  # Arbitrary large upper limit\n    num1 = 9 * b + 7\n    num2 = b + 7\n    if num1 % num2 == 0:\n        valid_bases.append(b)\n        print(f"Found base: {b}")\n\n# Sum the valid bases\nsum_bases = sum(valid_bases)\nprint(f"Sum: {sum_bases}")\n\n# If sum is over 1000, take modulo 1000\nif sum_bases > 1000:\n    result = sum_bases % 1000\nelse:\n    result = sum_bases\n\nprint(f"Final Result: {result}")\n</tool_call>
```output
Found base: 21
Found base: 49
Sum: 70
Final Result: 70
```
The code confirms that the valid bases are 21 and 49, summing to 70. Therefore, the final answer is 70. Since 70 is less than 1000, no modulo operation is needed.

我们开发了一种将代码执行集成到长推理生成中的方法。我们最初尝试通过简单提示从 DeepSeek-R1 elicited 工具集成推理 (TIR) 生成，但未成功。我们假设这些模型由于 extensive training on reasoning tasks 和有限的 instruction-following 暴露，难以偏离其标准解决方案格式。为了克服这一挑战，我们构建了一个管道，从指令跟随模型 LIMO 的小规模推理微调开始。通过提示该模型生成长推理 TIR 解决方案， followed by aggressive quality filtering，我们建立了一个适合训练的初始数据集。通过多次训练、生成和过滤迭代，我们构建了一个 170 万的 TIR 解决方案集，这对于提高最终模型的准确性至关重要。

对于最终阶段的 TIR 训练，170 万解决方案根据代码执行次数和答案正确性被 aggressively 过滤到 1.5 万个样本。

模型训练

对于我们的 Kaggle 提交，我们使用了与论文中详述的略有不同的训练配方。我们发现这个不同的配方产生的 token 比公开发布的模型少，而且我们没时间对最终模型进行更多减少 token 的实验，尽管它们表现更好。

我们首先在一个 220 万 CoT 解决方案子集上使用 SFT 训练了一个 Qwen2.5-14B-Base 模型，共 8 个 epoch。我们将基础 RoPE 改为 500k 以允许长推理。该模型使用 NVIDIA/Nemo-Skills 训练 8 个 epoch，学习率 1e-4，AdamW，权重衰减 0.01，10% 线性 warmup 衰减到 1e-7 的学习率。我们使用 1024 个样本的批处理大小，并利用 NVIDIA/NeMo-Aligner 的序列打包和上下文并行化技术，这显著加速了长推理数据的训练。训练在 512 个 H100 上持续了 48 小时（是的，512 个！）。我们可以用 20% 的计算量实现模型的大部分强度，但我们想扩大规模看看学习在哪里饱和。请参阅下方链接出版物中的图 3(b)，显示了不同训练阶段的指标。不同阶段的权重平均用于最终权重。

随后是在 1.5 万 TIR 样本上进行轻量级 TIR 微调。我们训练 TIR 模型 400 步，恒定学习率为 1e-5，并使用最后一个 checkpoint 而不进行平均。

如下详述，我们然后合并 CoT 和 TIR checkpoints，因为它既提高了准确性，又通过减少解决方案长度和代码执行次数加快了生成速度。

评估数据集

在比赛期间，我们主要使用 2024 年的美国邀请数学考试 (AIME) 和哈佛 - 麻省数学锦标赛 (HMMT) 的问题。后来我们添加了 2025 年的问题。我们的最终基准 Comp-Math-24-25 包含 256 个问题。

推理

模型合并

在这次比赛中，我们探索了合并具有 CoT 和 TIR 行为的两个 LLM 的各种方法。我们的主要目标是有效地结合这两个微调阶段的独特优势，以增强模型的性能。我们试验了 mergekit 包中的几种合并技术。令人惊讶的是，最有效的方法原来是两个 checkpoints 的简单线性组合：TIR 微调之前使用的 CoT checkpoint 和此后达到的最佳 TIR checkpoint。这种策略允许我们控制每个阶段对最终模型行为的影响程度。

在下表中，我们展示了合并模型在 Comp-Math-24-25 数据集上的准确性和生成统计信息。length 表示每个解决方案的平均 token 数，而 code 指的是每个解决方案的平均代码执行次数。

模型加速

我们优先考虑 Int8 仅权重 (W8A16) 和 FP8 量化，它们提供了比 BF16 更快的推理速度，且准确性损失最小。减小的权重大小也释放了内存用于更大的 key-value 缓存。
ReDrafter 是 Apple 开发的一种投机解码技术，并在 TensorRT-LLM 中实现。我们在 OpenMathReasoning-1 数据集问题的随机子集上训练了一个 ReDrafter 头。使用这些问题，我们用目标模型生成了 10 万个解决方案。生成的 redrafter 为每个 LLM 步骤生成 3 个 token，接受率为 65%，提供了约 1.8 倍的加速。TensorRT-LLM repo 的示例实现有一个非量化 llama 的示例，所以我们需要一些更改才能使其工作。我们尝试了其他起草方法，但无法让它们工作得这么好。特别是 LLM 起草应该可行且可能更好。

表中的准确率分数使用的是合并模型的 maj@12，平均超过 5 次运行。

TensorRT-LLM 推理

预训练模型使用 TensorRT-LLM 转换为 TensorRT 引擎。TensorRT 的 inflight batching 通过动态分组推理请求来提高吞吐量，一旦每个样本完成就释放它——减少延迟并优化 GPU 利用率。请参阅下方 vLLM 团队的一些最新基准测试。

由于样本是独立处理的，批处理可以无缝混合不同的提示或种子。TensorRT-LLM 包括许多其他优化，如自定义注意力 kernel 和分页 KV 缓存。

对于每个新问题，我们使用不同的种子通过 TensorRT 中的 in-flight batching 启动 12 个异步生成。监控每个样本的流以查找代码块、停止短语、最大 token 或超时。如果 LLM 生成代码，LLM 生成停止，代码块在沙盒中执行。沙盒输出（或部分错误跟踪）附加到 LLM 并继续生成。生成继续直到另一个代码块。当没有遇到其他代码块时，LLM 在最大 token、超时时间或停止短语处停止。
我们试验了各种采样参数，但观察到结果差异很小。然而，我们获胜的提交基于“几乎”贪婪搜索策略。尽管将温度设置为 0 并启用 redrafter_greedy_search 参数，TensorRT-LLM 仍然在单个批处理中产生略微不同的输出。我们没有详细调查这种行为，但推测它可能与舍入或批处理问题有关。最终，我们选择了这种方法，因为它在小批处理大小下提供了更稳定的结果，并在投机解码速度方面提供了 slight 改进。
在 Kaggle notebook 内，监控生成的完成情况。当初始答案相同时，我们早期停止，例如，如果前 5 个生成中有 4 个相同，我们取消 TensorRT 中的剩余生成。

在监控异步生成期间，我们还在完成 12 个生成中的 10 个后早期停止，以避免等待任何落后者过长。实施了一种缓冲策略，为每个问题分配 350 秒作为基本时间限制。如果问题提前完成，未使用的时间将添加到共享缓冲区。下一个问题然后可以从该缓冲区提取最多 210 秒的额外时间，总计 560 秒。

使用此推理管道和 deepseek-distill-14b，我们使用不同设置在公共榜单上的得分在 23 到 29 之间。分数的变异性令人担忧，我们自己训练的模型有助于稳定分数以及提高准确性。

对于最终选择的提交，选择了一个 14b CoT 模型和上述 MIX TIR 模型。MIX TIR 模型在 CV 上得分好得多，在公共榜单上也稍好一些（公共榜单分数：32,33,28）。那个 28 分的提交让我们感到紧张，可能是由于使用了更大的 maj@。最终，私有榜单比榜单结果更能反映我们的 CV 结果，这相当令人满意。

我们还有其他模型，如上所见，在 CV 上相当强（见报告），但相同大小的模型产生的 token 比以前的版本多，我们相信在公共榜单上它们在有机会properly 尝试问题之前就超时了。最终，我们没时间也没提交来properly 缩小 CV 和公共榜单之间看到的差异，特别是考虑到每天一次提交和只评分 50 个问题的方差。

未使用的方法

进行了一些实验，但没有包含在最终提交选择中，主要是由于时间有限，可能导致在做出好答案之前过早退出，

• 更强的模型 - 如论文中详述，我们有在难问题上训练的更强模型，但是生成了更多 token，我们不确信它能在 Kaggle 约束内完成。
• 生成解决方案重排序 - 这种类型的奖励模型需要额外的内存和时间资源。与为较小样本集服务额外的生成模型相比，增加批处理大小和最大 token proved 是更简单的替代方案。
• 结果奖励模型 - 初步实验显示该模型有一些优势。然而，对于长推理生成和 TIR，这些好处 diminished，特别是在评估较小子集（如 50 个问题）时。
• 7b 模型（TIR 和 CoT）与 redrafter 非常快，但 CV 和 LB 较低，尽管 batch size 32 完成得 comfortably。
• 32b CoT 模型 comfortably fit batch size 8 和 18k 生成长度，但分数接近 14b。
• 开发了一种可能更好的平滑缓冲策略，根据 5 小时剩余时间计算缓冲时间。然而，在进行上述提交时未使用。