RL loss importance ratio can overflow to non-finite gradients

[morluto]

Summary

The RL loss computes an importance ratio as:

ratio = exp(trainer_logprob - inference_logprob)

When rollout logprobs are stale relative to trainer logprobs, this ratio can overflow and produce non-finite loss/gradients.

Example

A single token is enough:

trainer_logprob = -2.0
inference_logprob = -91.0
log_ratio = trainer_logprob - inference_logprob  # 89.0
ratio = exp(log_ratio)  # inf in float32/bfloat16

If the token contributes to the policy-gradient term, the loss can become inf or -inf. If the token is masked, 0 * inf can still become nan, so masking does not fully protect the step.

Impact

Before the optimizer step, the non-finite value can flow through backward into gradients. Without a finite-loss or finite-gradient guard, optimizer.step() can apply non-finite gradients to model weights or optimizer state.

This is especially reachable in async/off-policy RL, where the trainer policy and rollout-time inference policy can drift.

Expected behavior

The policy-gradient importance ratio should be bounded before it multiplies token advantages.

mismatch_kl should remain useful as an off-policy diagnostic, but its computation should be numerically safe so monitoring does not receive inf/nan.

Proposed fix

• Clip the policy-gradient importance ratio with a configurable cap.
• Keep diagnostic mismatch-KL separate from the clipped training ratio.
• Add a finite-gradient guard before optimizer.step(), including paths where gradient clipping is disabled.
• Add CPU tests for overflow, masked-token overflow, IPO/ref-KL paths, config validation, and clipping metrics.