【Bug已解决】TPOTrainer.evaluate() returns NaN eval_loss while training loss is finite 解决方案一、现象长什么样用TPOTrainerToken-level Policy Optimization训练时训练 loss 一直在合理的有限值附近波动但调用trainer.evaluate()后日志里的eval_loss却是NaN{loss: 0.82, grad_norm: 1.3, epoch: 2} {eval_loss: nan, epoch: 2}更具体地有时是恒定的nan有时是inf但训练侧一切正常。现象指向evaluate()路径里某处计算与train()路径不一致产生了未定义值0/0、log(0)、或空组的平均而这条路径不影响参数更新所以训练 loss 看起来好好的只有评估指标废了。这种 bug 的危害是看不见的训练失效你以为模型在学train loss 在降但 eval 全 nan没法判断泛化还可能掩盖了真正的数据/数值问题。二、背景TPO 在 token 级做策略优化loss 通常形如每个 token 有一个优势advantageloss -advantage * logprob的某种加权。和 GRPO 类似它往往按prompt group组织样本组内做归一化。train()和evaluate()理论应走几乎相同的 loss 计算但实践中evaluate()常写成简化版比如直接用Trainer基类默认的 eval 行为它对因果 LM 用shift_labels算交叉熵而 TPO 的 loss 不是标准交叉熵——于是evaluate()用的是错误的 loss 公式再叠加一些边界情况空 group、全 padding 样本、advantage 全 0就产出 NaN。常见制造 NaN 的点0/0组内 advantage 归一化时std0整组 reward 相同除以零得 nan空 batch某 eval batch 全是 padding 样本num_tokens0平均时除零log(0)某 token 的 logprob 为-inf概率 0乘上非 0 优势得-inf*有限 nanloss 公式不一致evaluate没用 TPO 的 token-level loss而是基类交叉熵数值范围与预期不符。三、根因根因一句话TPOTrainer.evaluate()没有复用train()的 TPO token-level loss 计算而是走了基类的默认 eval 路径或一份有缺陷的简化版在空组 / 零 std / 零 token 等边界下产生 NaN而这组 NaN 不参与梯度更新所以训练 loss 正常、eval 全 nan。具体基类Trainer.evaluate默认会计算eval_loss基于模型输出 logits 的交叉熵但 TPO 的损失语义是 token-level 策略梯度损失两者不是一回事且 TPO 的归一化group std在 eval 的某些 batch 上触发 0/0。结果eval_loss既算错了公式又踩了除零稳定输出 NaN。四、最小可运行复现下面用纯 Python 复现两个核心 NaN 来源组内 std0 的 0/0以及空 batch 的除零平均def group_advantage(rewards): mean sum(rewards) / len(rewards) std (sum((r - mean) ** 2 for r in rewards) / len(rewards)) ** 0.5 return [(r - mean) / std for r in rewards] # std0 - 0/0 nan def mean_loss(losses): return sum(losses) / len(losses) # 空列表 - 0/0 nan/ZeroDivision def demo(): # 1) 整组 reward 相同 - std0 - 0/0 adv group_advantage([1.0, 1.0, 1.0]) print(零 std 组优势:, adv, 含 nan:, any(a ! a for a in adv)) # 2) 空 batch 平均 try: mean_loss([]) except ZeroDivisionError as e: print(空 batch 平均:, e) if __name__ __main__: demo()输出零 std 组优势: [nan, nan, nan] 含 nan: True 空 batch 平均: division by zero两处都精确对应线上现象组里 reward 全一样强化学习初期常见时 std0归一化出 NaNeval 的某个 batch 若全是 padding/无效样本平均除零。复现了eval 稳定 nan的机制。五、解决方案第一层归一化加 epsilon 空组跳过第一层修掉两个除零组内归一化加eps空组/空 batch 直接跳过不计入def group_advantage(rewards, eps1e-8): n len(rewards) if n 0: return [] mean sum(rewards) / n var sum((r - mean) ** 2 for r in rewards) / n std (var eps) ** 0.5 # 加 epsstd0 不再 0/0 return [(r - mean) / std for r in rewards] def safe_mean_loss(losses): if not losses: return 0.0 # 空 batch 返回 0不除零 return sum(losses) / len(losses) def demo(): adv group_advantage([1.0, 1.0, 1.0]) print(加 eps 后零 std 组优势:, adv, 含 nan:, any(a ! a for a in adv)) print(空 batch 平均:, safe_mean_loss([])) if __name__ __main__: demo()eps让std0时退化为全 0 优势整组一样本来就没相对信号给 0 正确safe_mean_loss对空 batch 返回 0.0 而非除零。两步消除两类 NaN。六、解决方案第二层evaluate 复用 train 的 TPO loss而非基类交叉熵第一层只是补丁但eval_loss仍可能是错公式算出来的有限值。第二层让evaluate()真正复用train()的 TPO token-level loss保证两者语义一致import torch import torch.nn.functional as F class TPOTrainer: def __init__(self, eps1e-8): self.eps eps def tpo_loss(self, logps, advantages, mask): TPO token-level loss只在有效 token 上加权平均。 if mask.sum() 0: return torch.tensor(0.0, requires_gradTrue) # 空组返回 0 weighted -(advantages * logps) * mask return weighted.sum() / mask.sum().clamp(minself.eps) def training_step(self, logps, adv, mask): return self.tpo_loss(logps, adv, mask) def evaluate(self, logps, adv, mask): # 关键evaluate 复用同一份 tpo_loss而不是基类交叉熵 with torch.no_grad(): return self.tpo_loss(logps, adv, mask) def demo(): t TPOTrainer() logps torch.randn(2, 3, requires_gradTrue) adv torch.randn(2, 3) mask torch.ones(2, 3) train_l t.training_step(logps, adv, mask) eval_l t.evaluate(logps.detach(), adv, mask) print(train/eval 用同一公式:, torch.isclose(train_l.detach(), eval_l)) # 空组不再 nan empty t.evaluate(logps.detach(), adv, torch.zeros(2, 3)) print(空组 eval_loss , empty.item(), is nan:, empty.isnan()) if __name__ __main__: demo()核心是evaluate调用self.tpo_loss(...)而非基类默认交叉熵且mask.sum()0时返回 0.0。这样eval_loss与train()的 loss 同构数值可比对且不再 NaN。七、解决方案第三层NaN 护栏 评估聚合去无效样本第三层在评估聚合时剔除无效样本并加 NaN 护栏保证eval_loss永远有限import torch def aggregate_eval(losses): 聚合各 batch eval_loss剔除 nan/inf 后再平均。 valid [l for l in losses if torch.isfinite(l)] if not valid: return 0.0 return sum(valid) / len(valid) def guard_finite(x: torch.Tensor, fallback: float 0.0) - torch.Tensor: 把 nan/inf 替换成 fallback避免污染后续聚合。 return torch.where(torch.isfinite(x), x, torch.tensor(fallback)) def demo(): raw [torch.tensor(0.8), torch.tensor(float(nan)), torch.tensor(0.9), torch.tensor(float(inf))] cleaned [guard_finite(r).item() for r in raw] print(护栏后:, cleaned) print(聚合 eval_loss , aggregate_eval([guard_finite(r) for r in raw])) if __name__ __main__: demo()guard_finite在每 batch 的 loss 上兜底nan/inf 变 0.0不污染聚合aggregate_eval再剔除仍异常的批次只对有限值平均保证最终eval_loss永远有限且有意义。八、落地建议如果你在TPOTrainer上遇到 eval nan建议确认 evaluate 是否复用 TPO loss不是就改成调同一份tpo_loss。归一化加 eps组内 advantage 除 std 时加eps1e-8防 0/0。空组/空 batch 返回 0mask.sum()0直接返回 0.0 tensor。加 NaN 护栏每 batchguard_finite聚合时aggregate_eval剔异常。对齐 train/eval 公式两者 loss 必须同构否则 eval_loss 数值不可比。加测试构造全相同 reward 组空 batch断言 eval_loss 有限。九、排查清单如果TPOTrainer.evaluate()返回 NaN 而 train loss 正常按顺序查确认 evaluate 用的 loss 公式是否复用train()的 TPO token-level loss还是基类交叉熵。看组内优势是否 0/0整组 reward 相同时 std0归一化出 NaN加eps。看是否有空 batcheval batch 全 padding 时平均除零返回 0.0。看 log(0)某 token logprob 为-inf乘非 0 优势得 nan加 mask 屏蔽。加 NaN 护栏每 batchguard_finite聚合aggregate_eval剔异常。对齐 train/eval两者 loss 同构eval_loss 才可比对。加边界测试锁住零 std 组空 batch下 eval_loss 有限。十、小结TPOTrainer.evaluate()返回 NaN 而训练 loss 正常根因是**evaluate()没复用train()的 TPO token-level loss而是走了基类默认 eval 路径或缺陷简化版在零 std 组0/0、空 batch除零、log(0) 等边界下产生未定义值而这组 NaN 不参与梯度更新所以训练侧毫无破绽只有评估指标废了**。修复分三层第一层给组内归一化加eps、空组/空 batch 返回 0.0消除两类除零第二层让evaluate()真正调用与train()同一份tpo_loss保证两者 loss 同构、数值可比第三层加guard_finite与aggregate_eval护栏剔除 nan/inf 再平均保证eval_loss永远有限。核心心法是eval 必须复用 train 的 loss 语义并对所有零分母/空集合边界显式兜底——否则评估指标会静默变成 NaN让你误以为训练正常、实则失去了对泛化的唯一观测窗口。