[bug] KL Divergence explodes to billions in GRPO training

[gitlost-murali]

Hello,

🐛 Describe the bug

I found very high KL divergence scores on the main branch with

python -m apps.grpo.main --config apps/grpo/qwen3_1_7b.yaml

and

group_size: 8
local_batch_size: 2 # per-device batch size
max_req_tokens: 1024
max_res_tokens: 16384

(complete config file: qwen3_1_7b_kl_checks.yaml )

Here are the logged numbers with KL divergence scores

grpo_loss/kl_divergence_mean: 1,568,987.96
grpo_loss/kl_divergence_max: 6,821,680,128.0
grpo_loss/policy_gradient_loss: -0.198
grpo_loss/total_loss: 1.77
buffer/sample/avg_sampled_policy_age: 1.0 

beta=1e-6 is suppressing these high values but something seems really off

Versions

Tested latest on e940fd89ad9b262e690988f2f1a5ee6cb0d25574

Config: apps/grpo/qwen3_1_7b.yaml with 16k response length and batch size 2

Step: 9
torch: 2.9.0+cu128
torchtitan: 0.2.0
vllm: 0.10.1.dev0+g6d8d0a24c.d20251219

[felipemello1]

hey @gitlost-murali , thanks for reporting! I am working on improving our losses. The way we have it is extremely naive, no clipping, vanilla GRPO (no normalization from dr GRPO), it uses KL while most of the modern losses dont.

Let me try to land the losses PR soon and see if this improves.

If not, i can think if a few culprits:

• Some token mismatch between inference/training, but thats unlikely since we dont retokenize
• Some masking error on loss/metric computation (e.g. compute the KL on padding tokens)
• some severe difference in top-p/top-k in inference/train

[gitlost-murali]

Thanks @felipemello1

I looked at the numbers in loss function (diff = ref_logprobs - logprobs):

Step	diff_max	diff_min
1	1.7	-1.8
2	8.9	-8.5
3	11.7	-14.6
4	15.3	-16.9
6	12.2	-23.7

1. As expected, the divergence grows as training progresses - policy drifts from ref_model. Hence, bigger diffs as we go further
2. The big numbers makes sense as exp(15) ~= 3269017.
3. Learning rate 1e-5 might be too aggressive. I'll check few runs with 1e-6 and lower learning rates
4. Clipping could be a good idea

My colleague @BastianBollAA also suspects there could be numerical precision issues or attention mask related issues

[felipemello1]

i think that the current beta is very very small. If that is true, then in practice there is no KL penalty being added to the loss. Could you double check? Perhaps increase the value?

but what i read in all papers is that, if we do "clip higher", then KL is controlled, and we dont even need the KL penalty.

I am sorry about the bug, let me get back to you when i land the loss PR, hopefully later today or tomorrow.

[gitlost-murali]

current beta is very very small. If that is true, then in practice there is no KL penalty being added to the loss. Could you double check? Perhaps increase the value

True. I'll increase the value and check

Thanks! No problem. not in a hurry though

[gitlost-murali]

Regarding the current simple_grpo_loss function:

Any reason why we don't pass/use the old_logprobs (vllm_logprobs) in the GRPO loss function. We are off-policy by 1 (almost on-policy). So, I think it makes sense to use old_logprobs (or vllm_logprobs).

Nit: The vllm logprobs are not enabled by default. We have to enable the flag

sampling_params:
  top_p: 1.0
  logprobs: true  # ← Added now!

What do you think about this loss function:

def simple_grpo_loss(
    logits: torch.Tensor,
    response: torch.Tensor,
    old_logprobs: torch.Tensor,
    ref_logprobs: torch.Tensor,
    advantages: torch.Tensor,
    padding_mask: torch.Tensor,
    beta: float = 0.0,
    epsilon: float = 0.2,
) -> torch.Tensor:

    logprobs: torch.Tensor = compute_logprobs(logits, response)
    
    # KL divergence from frozen reference model (prevents overall drift)
    kl = torch.exp(ref_logprobs - logprobs) - (ref_logprobs - logprobs) - 1

    # Importance sampling ratio: current policy / old policy (that generated the samples)
    ratio = torch.exp(logprobs - old_logprobs)
    clipped_ratio = torch.clamp(ratio, 1 - epsilon, 1 + epsilon)

    # PPO/GRPO-style clipped objective: take the min to be conservative
    per_token_policy_loss = torch.min(ratio * advantages, clipped_ratio * advantages)

    # Compute mean KL per valid token
    mean_kl = (
        ((kl * padding_mask).sum(dim=1)) / (padding_mask.sum(dim=1).clamp(min=1.0))
    ).mean()

    # Compute mean policy loss per valid token
    mean_policy_loss = (
        ((per_token_policy_loss * padding_mask).sum(dim=1))
        / (padding_mask.sum(dim=1).clamp(min=1.0))
    ).mean()

    # Compute loss using the means (mathematically equivalent)
    loss = -(mean_policy_loss - beta * mean_kl)
    return loss

[felipemello1]

ratio should probably be clipped before the exp

log_ratio = log_probs - old_log_probs.detach()
log_ratio = torch.clamp(log_ratio, min=-20.0, max=20.0)
ratio = torch.exp(log_ratio)
clipped_ratio = torch.clamp(ratio, 1 - epsilon, 1 + epsilon)

i also do the kl like this:

log_diff = log_ref.detach() - log_policy
log_diff = torch.clamp(log_diff, min=-20.0, max=20.0)
ratio = torch.exp(log_diff)
kl = ratio - log_diff - 1
kl = torch.clamp(kl, min=-10.0, max=10.0)

we currently dont provide a padding mask, but you can get it by doing target==PAD_ID. My PR will have a loss mask. Just make sure that your target is 1 off of inputs. I think in main.py we currently just slice the inputs and do the loss on the response instead. I dont recall. My PR will address that though. Like i said, its very naive as it is.

[felipemello1]

sorry, edited something that was wrong in my response

[zeocax]

I've also encountered this KL spike issue in my own experiments (a multi-turn implementation by myself).

Even with verl-style numerical stability handling like this:

log_diff = log_ref.detach() - log_policy
log_diff = torch.clamp(log_diff, min=-20.0, max=20.0)
ratio = torch.exp(log_diff)
kl = ratio - log_diff - 1
kl = torch.clamp(kl, min=-10.0, max=10.0)

the issue still persists.

I took a closer look at the token level and found that most of the abnormal KL values are concentrated on a very small number of tokens, so this likely explains why just numerical clamping alone isn't very effective.

Adding a learning rate warmup may help, the KL spike disappeared in my case, but I also made some other changes so I can't be certain it was the warmup specifically.

[zeocax]

I took a closer look at the token level and found that most of the abnormal KL values are concentrated on a very small number of tokens, so this likely explains why just numerical clamping alone isn't very effective.

What I mean is, even after clamping to the upper bound, there's still a spike.

[felipemello1]

@zeocax, mind sharing which other changes you made, if you recall them?