OpenReasoner项目中PRM模型Step Tag处理Bug分析与修复

问题背景

在OpenReasoner项目的数学推理训练过程中，研究人员发现当运行train_math.py脚本时，程序会在PRM(Preference Reward Model)模型的奖励计算阶段抛出数组越界异常。具体表现为在获取推理步骤得分时，尝试访问空张量的最后一个元素导致IndexError。

错误现象分析

错误发生在PRM模型的get_reward方法中，当代码尝试获取step_score[-1]时，发现step_score实际上是一个空张量。通过调试发现，关键问题出在对推理步骤标记(step tag)的处理上。

根本原因

经过深入排查，发现问题源于step tag的tokenization过程。在原始代码中，step tag("ки")前面被意外添加了空格，导致tokenizer将其识别为两个不同的token：

• 无空格"ки" → token ID: 16748
• 有空格" ки" → token IDs: 7665, 1802

这种差异使得模型无法正确识别和处理step tag，最终导致step_score计算结果为空。

解决方案

修复方案非常简单直接：确保step tag前没有多余空格。具体修改是将PRM模型中的输入拼接代码从：

inputs_for_prm.append(f"{o}{a} {self.step_tag}")

改为：

inputs_for_prm.append(f"{o}{a}{self.step_tag}")

即删除{a}和{self.step_tag}之间的空格。

技术影响

这个修复确保了：

1. step tag能够被正确tokenize为预期的单个token
2. PRM模型能够准确识别推理步骤边界
3. 奖励计算能够正常获取每个推理步骤的得分

对于数学推理训练任务而言，正确的step tag处理至关重要，因为它直接关系到强化学习过程中对中间推理步骤的奖励信号计算。

经验总结

这个案例提醒我们，在使用基于token的模型时，需要特别注意：

1. 特殊标记的tokenization一致性
2. 空格等不可见字符对tokenization的影响
3. 输入拼接时的格式控制

特别是在多模块协作的系统中，一个模块的输出格式必须严格匹配另一个模块的输入预期，否则可能导致难以察觉的错误。