【Bug已解决】[Bug] Zero-std reward groups produce spurious KL gradients when beta ＞ 0 in GRPO/RLOO 解决方案

【Bug已解决】[Bug] Zero-std reward groups produce spurious KL gradients when beta 0 in GRPO/RLOO 解决方案一、现象长什么样在 GRPO / RLOO 里开了 KL 正则beta 0后我们发现一个反直觉的现象有些样本组明明没学到任何相对信号整组 reward 完全相同却在持续产生梯度、缓慢把策略拉向 reference。具体表现某组 reward [1.0, 1.0, 1.0, 1.0] - 组内 std 0 该组优势被算成 00/0 经 eps 处理 但该组仍在更新参数且方向是靠近 ref也就是说零 std 组的优势是 0策略梯度项为 0但 KL 项beta·KL(policy‖ref)不受优势影响照样产生梯度。于是这些无信号组借着 KL 项偷偷推动策略等价于用了一批没有奖励依据的样本去微调模型——这叫spurious虚假KL 梯度。现象特征只在beta 0开 KL时暴露beta 0时零 std 组完全无梯度没问题训练后期很多组 reward 趋同std→0此时 spurious KL 梯度占比上升拖慢甚至干扰学习不报错但无信号样本在悄悄改参数评估和收敛曲线出现异常。二、背景GRPO / RLOO 的 loss 通常形如policy_grad_term -advantage * logp # 优势为 0 时此项为 0 kl_term beta * KL(policy || ref) # 与优势无关始终存在 loss policy_grad_term kl_term其中advantage是组内相对优势。当整组 reward 相同std0advantage全 0policy_grad_term 对整组是 0——这本意是这组没相对信号别用策略梯度推它。但kl_term是逐 token 的 KL 正则和优势无关。只要beta 0它就会对组内每个 token 产生把 policy 拉向 ref的梯度。于是策略梯度项说这组别动KL 项说这组往 ref 靠两者冲突且 KL 项单方面生效 → 零 std 组被 KL 悄悄推动。更糟的是KL 项里logp_policy - logp_ref的符号取决于当前策略相对 ref 的偏移若 policy 已偏离 refKL 梯度把它拉回这本身没错KL 正则本就这么工作但用在零 std、本应被整体跳过的组上就是 spurious——因为这组没有提供任何该往哪学的信号却因 KL 被强制往 ref 靠等于用无依据样本做了无谓的参数更新。三、根因根因一句话GRPO / RLOO 在beta 0时KL 正则项与优势无关、始终产生梯度当某组 reward std0整组优势为 0时策略梯度项为 0但该组仍被 KL 项推动产生没有奖励依据却改参数的 spurious KL 梯度。具体KL 项不随优势归零beta·KL对每个 token 独立生效不受advantage0影响零 std 组未被整体跳过优势0 只让 policy_grad_term 为 0没让 KL 项也归零冲突信号策略项别动 vs KL 项靠 refKL 单方面生效只在 beta0 暴露beta0时无解所以问题被关 KL 就好掩盖但开了 KL 想稳定训练时暴露静默不报错只是无信号样本在改参数。本质是KL 正则的缩放没有和组内是否有信号挂钩——零 std 组应整体退出优化包括 KL而非只退出策略梯度。四、最小可运行复现下面用纯 Python 模拟零 std 组 KL 仍产生梯度的机制def group_loss(rewards, logp_policy, logp_ref, beta, eps1e-8): n len(rewards) mean sum(rewards) / n var sum((r - mean) ** 2 for r in rewards) / n std (var eps) ** 0.5 advantage (rewards[0] - mean) / std # 零 std 组 - 0/sqrt(eps) 实际约 0 policy_term -advantage * logp_policy # 零 std - 0 kl_term beta * (logp_policy - logp_ref) # 与 advantage 无关始终非零 return policy_term, kl_term def demo(): rewards [1.0, 1.0, 1.0, 1.0] # 零 std 组 logp_policy, logp_ref 0.5, 0.2 beta 0.04 pt, kt group_loss(rewards, logp_policy, logp_ref, beta) print(f零 std 组: policy_term{pt:.3f} (应≈0), kl_term{kt:.3f} (spurious!)) print(f总梯度贡献含 spurious KL: {pt kt:.3f}) if __name__ __main__: demo()输出零 std 组: policy_term0.000 (应≈0), kl_term0.012 (spurious!) 总梯度贡献含 spurious KL: 0.012policy_term为 0正确无相对信号但kl_term0.012仍存在成为 spurious 梯度。复现了核心 bug零 std 组被 KL 项悄悄推动。五、解决方案第一层零 std 组整体屏蔽含 KL第一层最直接检测到组内 std≈0 时整组包括 KL 项都不参与梯度从根上消除 spuriousimport torch def grpo_loss_per_group(rewards, logp_policy, logp_ref, beta, eps1e-8): rewards torch.tensor(rewards, dtypelogp_policy.dtype) adv rewards - rewards.mean() std rewards.std(unbiasedFalse) eps adv adv / std # 零 std 标志组内 reward 全相同 is_zero_std (rewards.std(unbiasedFalse) 1e-6) policy_term -(adv * logp_policy).sum() kl_term beta * (logp_policy - logp_ref).sum() if is_zero_std: # 整组退出包括 KL消除 spurious return torch.zeros_like(policy_term), True return policy_term kl_term, False def demo(): rewards [1.0, 1.0, 1.0, 1.0] lp torch.tensor(0.5, requires_gradTrue) lr torch.tensor(0.2) loss, skipped grpo_loss_per_group(rewards, lp, lr, beta0.04) print(f零 std 组: loss{loss.item():.3f}, 跳过{skipped} (KL 不再 spurious)) if __name__ __main__: demo()核心是is_zero_std当整组 reward 相同整组 return 0 梯度含 KL不再让 KL 项单独推动。这和最自然的语义一致——这组没提供任何学习信号整组都不该改参数。六、解决方案第二层用 group 级 mask 缩放 KL而非硬跳过第一层是硬跳过整组 0 梯度但有时我们仍希望KL 正则对这些组轻度生效以维持分布不漂移。第二层用group 级权重缩放 KL有信号的组 KL 全效零 std 组 KL 乘 0介于两者间可平滑过渡import torch def grpo_loss_group_weighted(rewards, logp_policy, logp_ref, beta, std_floor1e-6): rewards torch.tensor(rewards, dtypelogp_policy.dtype) adv rewards - rewards.mean() std rewards.std(unbiasedFalse) adv adv / (std 1e-8) # group 级权重std 越小KL 权重越低零 std - 0 group_w torch.sigmoid((std - std_floor) * 1000.0) # stdfloor 时≈1, 否则≈0 policy_term -(adv * logp_policy).sum() kl_term beta * (logp_policy - logp_ref).sum() * group_w return policy_term kl_term def demo(): for rewards in ([1.0, 1.0, 1.0, 1.0], [0.0, 1.0, 0.5, 0.8]): lp torch.tensor(0.5) lr torch.tensor(0.2) loss grpo_loss_group_weighted(rewards, lp, lr, beta0.04) print(f组 reward{rewards} - loss{loss.item():.4f}) if __name__ __main__: demo()group_w用 sigmoid 把std 是否高于 floor平滑成权重零 std 组group_w≈0KL 被有效关掉有信号组group_w≈1KL 正常。比硬跳过更平滑且对接近零 std的组也温和降级避免边界突变。七、解决方案第三层断言无 spurious 不变量测试第三层加护栏确保零 std 组的 KL 梯度贡献为 0并锁进测试import torch def kl_gradient_for_zero_std_group(rewards, logp_policy, logp_ref, beta): 返回零 std 组的 KL 项对参数的梯度范数应≈0。 lp logp_policy.clone().requires_grad_(True) loss grpo_loss_group_weighted(rewards, lp, logp_ref, beta) loss.sum().backward() return lp.grad.abs().sum().item() if lp.grad is not None else 0.0 def test_zero_std_no_spurious_kl(): rewards [1.0, 1.0, 1.0, 1.0] lp torch.tensor(0.5) lr torch.tensor(0.2) g kl_gradient_for_zero_std_group(rewards, lp, lr, beta0.04) assert g 1e-6, f零 std 组仍产生 KL 梯度 {g}spurious! print(fOK: 零 std 组 KL 梯度{g:.2e} (≈0, 无 spurious)) def test_nonzero_std_has_signal(): rewards [0.0, 1.0, 0.2, 0.9] lp torch.tensor(0.5, requires_gradTrue) lr torch.tensor(0.2) loss grpo_loss_group_weighted(rewards, lp, lr, beta0.04) loss.sum().backward() assert lp.grad.abs().sum() 1e-6, 有信号组应产生梯度 print(OK: 有信号组正常产生梯度) if __name__ __main__: test_zero_std_no_spurious_kl() test_nonzero_std_has_signal()两个测试分别锁住零 std 组 KL 梯度≈0无 spurious和有信号组正常产生梯度KL 不误杀。任何破坏这个平衡的改动都会被 CI 拦下。八、落地建议如果你在 GRPO/RLOO 开 KL 后发现训练被无信号组干扰建议识别零 std 组组内 reward std 阈值如 1e-6即视为无信号。整组屏蔽含 KL第一层硬跳过整组含 KL归零梯度。或平滑降级第二层用 group 级权重缩放 KL零 std→0有信号→1。断言无 spurious加测试锁住零 std 组 KL 梯度≈0。不误杀有信号组确保非平凡组 KL 正常。监控占比训练后期零 std 组比例上升时看 spurious 是否被消除。九、排查清单如果开 KL 后训练被无信号组干扰按顺序查确认 beta 0spurious KL 只在开 KL 时暴露。看零 std 组是否仍在更新reward 全相同的组不应产生梯度。确认 KL 项是否被 group 信号缩放没有则加is_zero_std或 group_w。第一层硬跳过整组含 KL。第二层group 级 sigmoid 权重平滑降级。加断言/测试锁住零 std 组 KL 梯度≈0、有信号组正常。看训练后期零 std 组占比上升时确认 spurious 已消除。十、小结GRPO / RLOO 开 KL 后零 std reward 组整组 reward 相同产生 spurious KL 梯度根因是KL 正则项beta·KL(policy‖ref)与组内优势无关、始终产生梯度当某组优势因 std0 而归零时策略梯度项为 0但 KL 项仍单方面推动策略于是没提供任何学习信号的组借着 KL 悄悄改参数。它只在beta 0时暴露、不报错却是无依据样本在微调模型的隐患。修复分三层第一层检测零 std 组并整组含 KL屏蔽从根消除 spurious语义最干净第二层用 group 级 sigmoid 权重平滑缩放 KL零 std→0、有信号→1比硬跳过更平滑、对接近零 std的组也温和降级第三层加test_zero_std_no_spurious_kl零 std 组 KL 梯度≈0与test_nonzero_std_has_signal有信号组正常不变量测试。核心心法是KL 正则的缩放必须与组内是否有学习信号挂钩——零 std 组应当整体退出优化包括 KL否则 KL 会把无依据样本变成悄悄改参数的 spurious 梯度干扰甚至拖慢整轮训练。