train_grpo.py中torch+grpo路径下old_per_token_logps未detach导致GRPO优势项梯度抵消

[zeta172417-design]

问题描述

我在使用train_grpo.py训练时，发现当使用如下配置：

--rollout_engine torch
--loss_type grpo
--use_moe 0

也就是torch+dense+grpo路径时，old_per_token_logps和per_token_logps可能来自同一个未detach的张量，导致GRPO中的ratio虽然数值为1，但梯度被抵消，进而使reward/advantage项无法有效更新actor。

---

相关代码

当前代码中大致是这样：

old_per_token_logps = rollout_result.per_token_logps.to(args.device)

...

if use_sglang or lm_config.use_moe:
    res = model_unwrapped(outputs)
    aux_loss = res.aux_loss if lm_config.use_moe else torch.tensor(0.0, device=args.device)
    logits = res.logits[:, :-1, :]
    per_token_logps = F.log_softmax(logits, dim=-1).gather(
        2, outputs[:, 1:].unsqueeze(-1)
    ).squeeze(-1)[:, -completion_ids.size(1):]
else:
    aux_loss = torch.tensor(0.0, device=args.device)
    per_token_logps = rollout_result.per_token_logps

在torch+dense+grpo路径下，会走else分支，于是：

old_per_token_logps = rollout_result.per_token_logps
per_token_logps = rollout_result.per_token_logps

然后GRPO中计算：

ratio = torch.exp(per_token_logps - old_per_token_logps)

等价于：

ratio = torch.exp(x - x)

这时ratio数值上等于1，但因为是同一个计算图张量相减，所以梯度也会被抵消：

d(x - x) / dx = 0

因此GRPO的reward/advantage项可能无法有效更新actor，训练主要只剩下KL项在起作用。

---

建议修改

可以将old_per_token_logps显式detach：

old_per_token_logps = rollout_result.per_token_logps.detach().clone().to(args.device)

这样ratio变成：

ratio = torch.exp(x - x.detach())

数值上初始仍然是1，但梯度不会被抵消：

d(x - x.detach()) / dx = 1

这样GRPO的advantage项才能正常通过当前策略的per_token_logps更新actor。

---

为什么这个问题主要出现在torch+dense+grpo路径

我理解是：

1. sglang路径下，old_per_token_logps来自外部推理服务，本身不连接本地PyTorch计算图；
2. use_moe路径下，会重新forward计算当前per_token_logps，不再是同一个张量直接相减；
3. 但torch+dense路径下，当前per_token_logps直接复用rollout_result.per_token_logps，如果old_per_token_logps不detach，就会出现exp(x - x)导致梯度抵消的问题。

---

观察到的现象

修改前，loss_type=grpo时模型似乎较难从reward/advantage中学习，kl_ref变化也较小。

将old_per_token_logps改为：

old_per_token_logps = rollout_result.per_token_logps.detach().clone().to(args.device)

之后，kl_ref变化明显增大，说明actor相对ref model开始发生更明显的策略偏移，这也符合GRPO优势项梯度恢复后的预期。

---

想确认的问题

请问这里old_per_token_logps未detach是有意设计，还是一个可能的实现bug？

从GRPO/PPO的理论上看，old_per_token_logps应该表示采样时旧策略的固定logprob，因此似乎应当作为常数处理，不应该继续参与反向传播。

[jingyaogong]

代码早已改为“一律重新 forward + 显式 .detach()”，该 Bug 不复存在，你没有把fork或者本地的代码对齐版本