fix(grpo): clip the KL log-ratio before exp to avoid inf overflow

[behroozazarkhalili]

What

The K3 KL estimator overflows to inf when the policy and reference distributions drift far apart during training (a large positive log-ratio), as reported in #3015:

per_token_kl = torch.exp(ref_per_token_logps - per_token_logps) - (ref_per_token_logps - per_token_logps) - 1

Fix

Add an opt-in GRPOConfig.kl_log_ratio_clip (float, optional, default None) that clips the log-ratio before the exponential:

kl_log_ratio = ref_per_token_logps - per_token_logps
if self.args.kl_log_ratio_clip is not None:
    clip = torch.tensor(self.args.kl_log_ratio_clip, dtype=kl_log_ratio.dtype, device=kl_log_ratio.device)
    if not torch.isfinite(torch.exp(clip)):
        raise ValueError(...)  # the clip itself must keep exp finite in the working dtype
    kl_log_ratio = kl_log_ratio + (kl_log_ratio.clamp(max=clip) - kl_log_ratio).detach()
per_token_kl = torch.exp(kl_log_ratio) - kl_log_ratio - 1

The clip is upper-only (a large negative log-ratio underflows exp to zero and leaves K3 finite) and straight-through: the value is clamped but the gradient passes as if it were not. A plain clamp zeroes the K3 slope for clipped tokens, and with use_bias_correction_kl=True (the default) the term reduces to K3(clip) * ratio, whose gradient rewards a lower policy log-prob, so minimizing the loss would push the policy further from the reference. Measured with plain torch at x = 20, clip = 10: d(KL)/d(log-prob) is +2.2e4 with a plain clamp, -10 straight-through, and -32 unclipped in float32, where the exact -x * ratio = -20 is already lost to cancellation.

• When None (the default) the estimator is bit-identical to the current one, so there is no behavior change for existing users.
• When set, it keeps per_token_kl finite while preserving the K3 shape (exp(x) - x - 1 stays non-negative for the clamped x). The value must be positive and finite: GRPOConfig rejects the rest, since a non-positive clip invents KL at an exact match and -inf reaches the estimator as inf past the trainer's overflow guard.

This is approach (a) from @albertvillanova's guidance on #3015 (clip the log-ratio before exp behind an opt-in field).

Consistency

The K3 block is duplicated in two experimental trainers that subclass GRPOTrainer (gmpo, gspo_token); the same guard is applied to both so the trainers stay aligned per the repo's duplication policy.

Tests

Seven kl_log_ratio_clip tests in tests/test_grpo_trainer.py, each building a tiny GRPOTrainer on CPU and calling _compute_loss directly on hand-built inputs (so they exercise the real K3 branch, not a standalone helper): the unclipped estimator overflows to inf above the dtype's exp ceiling and clipping tames it; a clip above that ceiling raises; a large negative log-ratio is left intact; a normally scaled log-ratio is unaffected; the gradient of the loss with respect to a clipped token's log-prob stays negative with the bias correction on and off; the KL term with zero advantages stays finite and positive; and five non-positive or non-finite clips are rejected. A copy of the tree with a plain clamp fails both gradient cases, and one with the validation removed fails all five rejection cases. ruff check + ruff format clean.

Since this is a numerical-stability guard rather than a new paper method, no paper_index.md entry was added.

Resolves #3015

cc @qgallouedec @kashif @albertvillanova

---

Note

Medium Risk
Changes KL penalty numerics and gradients only when users opt in via kl_log_ratio_clip, but mis-set clips or Liger + KL combinations now fail fast instead of silently diverging.

Overview
Fixes #3015 by adding an opt-in GRPOConfig.kl_log_ratio_clip (default None) that upper-bounds ref_per_token_logps - per_token_logps with a straight-through clamp before the K3 term exp(x) - x - 1, so large positive log-ratios no longer blow up to inf. When unset, behavior stays unchanged for typical log-ratios.

GRPOConfig validates the clip as positive and finite. GRPOTrainer (and the duplicated paths in GMPO / GSPO token trainers) also reject clip values whose own exp overflows the working dtype. use_liger_kernel=True with beta != 0 and a clip set now raises NotImplementedError because the fused Liger loss cannot apply the guard.

Tests add a shared _kl_clip_setup harness and regression coverage for overflow recovery, invalid config/clip values, no-op when ratios are small or negative, gradient direction under bias correction, and the Liger incompatibility.

^{Reviewed by Cursor Bugbot for commit ef66a5a. Bugbot is set up for automated code reviews on this repo. Configure here.}

[bot-ci-comment]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[qgallouedec]

Thanks, I'm not strongly opposed, but afaik, no-one is using a ref model regularization anymore. So it'm not sure if adding a new parameter again is really worth it

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit 707808f. Configure here.}

[behroozazarkhalili]

Fair point, and I would rather not add public config surface you have to keep stable if the KL path is winding down.

The crash itself is still real for anyone running beta != 0 (the K3 exp(ref - cur) overflows to inf on entropy collapse), but it can be fixed without a new parameter. A few lines up, the sequence-importance path already clamps a log-ratio unconditionally:

# grpo_trainer.py, existing
log_ratio_clamped = torch.clamp(log_ratio_per_token, -20.0, 20.0)

So I can drop kl_log_ratio_clip and instead clamp the K3 log-ratio to a fixed [-20, 20], consistent with the above. exp(20) is about 4.8e8, so it is a no-op for any legitimate per-token KL and only bites in the overflow regime. That fixes the crash, adds no new public API, and keeps the three duplicated K3 blocks (grpo, gmpo, gspo_token) consistent.

Want me to switch #6637 to that (unconditional clamp, no new field)? If you would rather not touch the KL path at all given its declining use, I am equally happy to just close this. cc @albertvillanova

[behroozazarkhalili]

Follow-up: Bugbot just flagged that kl_log_ratio_clip is silently ignored on the Liger path. The fused LigerFusedLinearGRPOLoss computes the K3 KL inside the kernel, so a Python-side clip never runs, and there is no init-time guard for the combination.

That is actually another reason to prefer the unconditional-clamp route: with no public kl_log_ratio_clip field there is no option that can silently no-op on Liger. The non-Liger path still gets the overflow fix, and the Liger kernel's internal K3 stays a separate upstream concern. So dropping the parameter resolves both your config-surface point and the Liger finding in one go. Happy to push that version whenever you give the nod.

[behroozazarkhalili]

Verification on GPU, since the Liger path cannot be exercised on a CPU runner.

Ran on one H100 at the pushed head 390356b9, torch 2.11.0 with triton 3.6.0, in two passes.

The 13 Liger nodes: 13 passed, exit 0, in 232s. These are the tests that go red locally without a GPU, where triton raises "0 active drivers" before any assertion runs, so a red there says nothing about the code. On real hardware they pass.

The four kl_log_ratio_clip cases (5 selected): 5 passed, exit 0, in 53s. That covers the clip staying finite above the dtype's exp threshold and the one-sided clamp leaving a large negative log-ratio alone.

This is the same hardware question raised on the Bugbot thread about the Liger path. The fused loss still raises NotImplementedError when kl_log_ratio_clip is set, and that guard is among the 13.