Warn when truncated sampling biases the vLLM importance-sampling ratio

[behroozazarkhalili]

What this does

Warns when vllm_importance_sampling_correction=True is combined with top_p < 1, top_k > 0, or min_p > 0.

Refs #6789. This is direction 4 from that issue, the warning, not the math fix. It forecloses none of the other directions.

Why

vLLM is asked for processed_logprobs, which are renormalized over the support that survives truncation, while the trainer takes a full-vocab log-softmax. grpo_trainer.py:2680 differences the two:

per_token_logps_diff = (old_per_token_logps - sampling_per_token_logps) * mask

so the result is log p_hf - log p_vllm + log S, where S is the surviving mass. The first two terms are the train/inference mismatch the correction exists to measure. log S is not, and it flows straight into vllm_importance_sampling_ratio, which multiplies the per-token loss.

Measured on one H100, GRPOTrainer in colocate mode, trl-internal-testing/small-Qwen2ForCausalLM-2.5, max_completion_length=16. The predicted values were fixed before the run:

top_p	predicted abs(log S)	measured	abs err
1.0	0.0000	0.0007	0.0007
0.9	0.1054	0.1054	0.0000
0.8	0.2231	0.2232	0.0001

A second logged quantity confirms the same mechanism independently. With 16 completion tokens the sequence-level correction is S^16, and sampling/importance_sampling_ratio/mean reads 0.1855 against a predicted 0.185302 at top_p=0.9, and 0.02814 against 0.028147 at top_p=0.8. Both agree to four decimals, which fixes the per-token surviving mass at exactly S = top_p.

At top_p=0.8 that weights a 16-token completion at 2.8 percent of its uncorrected value, with no error and no log line.

Why warn rather than fix the math

A correct fix needs the kept-token set the sampler actually used. Both frameworks the issue cites do exactly that: prime-rl #3235 transports kept ids out of vLLM's engine-core worker, and slime v0.3.1 requires rollout_top_p_token_ids whenever rollout_top_p is not 1.0. vLLM exposes the set only in #49577, which is merged but lands after 0.27.1 and so falls outside the current vllm>=0.18.0,<=0.27.1 pin.

Recomputing the mask from training-side logits is possible today but is an approximation, and it has a sharp edge: the sampled token can fall outside the recomputed nucleus, giving -inf and a NaN batch unless it is force-included. I left that for the design discussion on the issue.

Scope

The guard is in GRPOConfig.__post_init__, so GRPOWithReplayBufferConfig inherits it and the gspo_token variant picks it up through the shared config. RLOO is untouched because it has no importance-sampling correction at all.

Verification

• Test-first: the four should_warn=True cases fail before the guard, all six pass after.
• Mutation-checked: disabling each of the top_p, top_k, and min_p conditions, and removing the use_vllm and correction gate, each kills a test. The file is byte-identical after restore.
• ruff 0.13.3 check and format clean, version read from .pre-commit-config.yaml.
• doc-builder at the pinned 2430c1e with --max_len 119 clean on both files, with trl/trainer/sft_trainer.py as a positive control.

---

Note

Low Risk
Config-time UserWarning only; training math and sampling behavior are unchanged.

Overview
Warns when vLLM importance-sampling correction is used with truncated sampling (top_p < 1, top_k > 0, or min_p > 0).

vLLM logprobs are renormalized over the surviving support while the trainer uses a full-vocab log-softmax, so the IS ratio includes log S (surviving mass) on top of the train/inference mismatch. GRPOConfig.__post_init__ now emits a UserWarning in that combination and stays silent when correction is off or sampling is untruncated. Tests cover the truncation knobs and the no-correction case.

^{Reviewed by Cursor Bugbot for commit 6cca24b. 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.