自定义 GRPO 训练循环为何不学习并出现梯度异常？

2026-11-06AI 小顾

结论

这段循环有几处实现问题，会让 GRPO 目标偏离预期。首先，你把当前策略的 log probabilities 当作参考值，因此前向计算得到的概率比恒为 1。其次，因果语言模型的 logits 和目标 token 没有错位对齐。生成时使用 temperature=2.0，计算 loss 时却直接使用未按该温度调整的模型分布，也会让 rollout 策略与训练时计算的策略不一致。

这些问题足以让训练信号失真。不过，仅凭贴出的代码和日志，无法确定它们是不是模型在约 500 步后停止学习的唯一原因。

先修正概率比的基准

GRPO 的策略概率比应比较当前策略与生成这批 completion 时的旧策略：

[ r_t(\theta)=\exp(\log\pi_\theta(a_t|s_t)-\log\pi_{\text{old}}(a_t|s_t)) ]

log_probs 应由正在更新的策略计算；old_log_probs 则应在采样时保存，或由更新前的策略在同一批 rollout 上计算。你的代码把 ref_logits 设为 policy_model(batch_ids).logits，因此两组 log probabilities 相等，coef1 恒为 1。梯度仍可能通过 log_probs 传播，但这个前向概率比没有体现策略相对于 rollout 时的变化。

ref_log_probs 通常用于 KL 正则项，作用不同于 old_log_probs。如果 KL 系数为零，可以暂时省略 KL 项；概率比仍应使用旧策略的 log probabilities。

# rollout 阶段保存生成 token 在采样策略下的 log probability：
old_log_probs = ...  # shape: (B, T)

# 更新阶段：当前策略的 log probability
current_log_probs = ...

ratio = torch.exp(current_log_probs - old_log_probs)

每批 rollout 更新一次或多次时，都应保留这批数据对应的 old_log_probs。如果每批只更新一次，也可以在生成完成后、参数更新前计算并保存旧策略的 log probabilities。

对齐 causal LM 的 logits 和 token

对标准 causal LM，位置 t 的 logits 用来预测位置 t + 1 的 token。当前代码在同一位置取 logits 和 batch_ids：

policy_logprobs = policy_logprobs.gather(-1, batch_ids.unsqueeze(-1)).squeeze(-1)

应将 logits 和目标 token 错开一位，并对 mask 做相同切片。否则，算出的就不是生成 token 在前一上下文下的概率。

# logits[:, t] 预测 input_ids[:, t + 1]
shift_logits = logits[:, :-1, :]
target_ids = input_ids[:, 1:]
shift_mask = response_mask[:, 1:]

log_probs = F.log_softmax(shift_logits, dim=-1)
token_log_probs = log_probs.gather(
    dim=-1,
    index=target_ids.unsqueeze(-1),
).squeeze(-1)

当前 response_mask 标出 response token 的位置。计算 loss 时也要使用 response_mask[:, 1:]，使它与 target_ids 对齐。训练批次还应提供 attention mask，避免左侧 padding token 参与模型上下文计算。response mask 和 attention mask 用途不同，不能互相替代。

保持生成分布与概率计算一致

采样时使用了 temperature=2.0，但后续 log probability 是从未缩放的 logits 计算的。如果 rollout 是从温度调整后的分布采样，保存的 old_log_probs 也应针对同一个分布计算，也就是使用 logits / temperature 后的分布。计算当前策略概率时，也要明确训练目标对应哪种策略分布，并保持一致。

此外，reward 使用解码文本的字符长度，而 GRPO loss 作用在生成 token 上。这可以作为自定义奖励目标，但字符长度和 token 数并不相同。要复现“鼓励更长 completion”的示例，应先确认奖励函数衡量的长度符合预期。

核对 loss 和归一化

你用负号最小化策略目标，方向是合理的：正 advantage 会提高概率，负 advantage 会降低概率。接下来应基于正确的旧策略概率比计算，并确认 batch 中的 reward 和 mask 对应同一组 completion。

当前 advantages = (rewards - rewards.mean()) / rewards.std() 是对本批 4 条 completion 做标准化。注意，torch.std() 默认使用无偏估计；样本组较小时，结果会不同于使用总体标准差。loss.sum() / 4 则把 batch 大小写死了，建议改用 .mean()，或明确实现所选 GRPO loss 的 reduction 规则。不同 reduction 会改变梯度尺度，因此不能只凭梯度范数判断实现是否正确。

建议的排查顺序

修正 logits、目标 token 和 response mask 的一位错位。
在 rollout 时保存生成策略的 old_log_probs，并用它计算概率比。不要把同一模型的输出同时当作当前策略和旧策略。
统一温度处理，并检查 padding 的 attention mask。
对一小批固定 rollout 做数值检查：更新前 ratio 应接近 1；检查有效 token 的数量、advantage 是否有正有负，以及 loss 和梯度是否为有限值。
确认优化器每次更新前都清零梯度，并观察 reward、completion token 长度和策略概率比的变化。clip_grad_norm_ 返回的是裁剪前的总梯度范数，所以日志中出现很大的值，不代表裁剪后仍以该范数更新。

先修正 token 对齐和 rollout 概率比，再根据固定批次的检查结果调整 KL、reward 标度和 reduction。仅凭现有日志，无法断定某个梯度范数阈值一定异常。