if we use GRPO and args.kl_coef is non-zero, is the KL computation incorrect?

slime/backends/megatron_utils/loss.py

    if args.kl_coef == 0:
        # when kl_coef is 0, we won't compute ref_log_prob
        xs = log_probs if log_probs is not None else values
        kl = [torch.zeros_like(x, dtype=torch.float32, device=x.device) for x in xs]
    else:
        kl = [
            compute_approx_kl(
                log_probs[i],
                ref_log_probs[i],
                kl_loss_type=args.kl_loss_type,
            )
            for i in range(len(log_probs))
        ]

    if args.advantage_estimator in ["grpo", "gspo"]:
        rewards = torch.tensor(rewards, dtype=torch.float32, device=kl[0].device)
        returns = get_grpo_returns(rewards, kl)
        # TODO: is the copy necessary?
        advantages = [r for r in returns]

slime/utils/ppo_utils.py

def get_grpo_returns(
    rewards: torch.Tensor,
    kl: list[torch.Tensor],
):
    returns = []
    for i in range(len(rewards)):
        returns.append(torch.ones_like(kl[i]) * rewards[i])
    return returns

If we use GRPO and args.kl_coef is non-zero, then kl_coef is effectively not used; the reward isn’t adjusted by args.kl_coef

Start in slime/backends/megatron_utils/loss.py and trace how kl_coef, kl, rewards, and get_grpo_returns interact for the GRPO path. Then inspect slime/utils/ppo_utils.py to determine whether the returned rewards incorporate args.kl_coef; done means the intended nonzero-coefficient behavior is established and covered by an appropriate regression check.