Warn when truncated sampling biases the vLLM importance-sampling ratio

[behroozazarkhalili]

What this does

Warns when something in GRPO consumes vLLM's logprobs, the importance-sampling correction or the off-policy mask, and the sampling vLLM runs with reshapes the distribution: top_p < 1, top_k > 0, min_p > 0, repetition_penalty != 1, or a generation_kwargs override of temperature. The check reads the effective values, so a generation_kwargs override of any of these counts and an override back to the default does not.

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 requests no vLLM logprobs at all. AsyncGRPO has the same bias with no correction to gate on; #6944 covers it.

The gate covers the off-policy mask as well as the correction because the trainer requests vLLM's logprobs unconditionally (grpo_trainer.py:1117), stores them whenever they arrive (:2931), and get_off_policy_mask reads them whenever off_policy_mask_threshold is set (:3162). Turning the correction off therefore does not remove the bias from the keep/drop decision, and the message no longer suggests it does on its own.

Three limits, stated in the code comment rather than checked. A temperature below vLLM's greedy threshold (1e-5 in its SamplingParams) makes vLLM skip temperature scaling while the trainer still divides, so the cancellation the check relies on breaks there; the constant is vLLM's, so it is not hardcoded here. top_k at or above the vocabulary size is disabled by vLLM but still warns, since the config does not know the vocabulary size. Other sampler-side modifiers passed through generation_kwargs (logit_bias, bad_words, frequency_penalty, ...) reshape the logits the same way and are not checked; the generation_kwargs docstring says so.

Verification

• Test-first: the four should_warn=True cases fail before the guard, all six pass after. On the previous commit four of the new cases fail (override alone, override back to the default, temperature override, off-policy mask with the correction off) and pass after; the matching-temperature and no-consumer cases pass on both.
• 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 warnings only; no training math or runtime sampling behavior changes beyond alerting users to a known metric bias.

Overview
Adds a GRPOConfig.__post_init__ guard that emits a UserWarning when vLLM is used and something actually consumes vLLM’s returned logprobs—vllm_importance_sampling_correction or a set off_policy_mask_threshold—while effective sampling settings would reshape those logprobs relative to the trainer’s full-vocab normalization.

The check uses merged effective values from config fields and generation_kwargs overrides for top_p, top_k, min_p, repetition_penalty, and temperature (warns when an override disagrees with the config temperature). Other sampler modifiers in generation_kwargs are called out in docs as unchecked but similarly biasing. The warning stays silent when neither consumer is enabled, even under truncated sampling.

Tests cover default vs truncated sampling, generation_kwargs overrides, off-policy masking without correction, and the no-consumer case. generation_kwargs docstring is expanded to describe this limitation (refs #6789).

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

[verma8076]

Nice work — the measured-vs-predicted |log S| table is exactly the right way to pin this down, and scoping to a warning (rather than guessing at the truncation mask) is the right call given vLLM's kept-token set isn't exposed until #49577, which is past the current pin.

One gap: this only guards GRPOConfig, so AsyncGRPOTrainer gets no warning at all. AsyncGRPOConfig doesn't subclass GRPOConfig — it extends _BaseConfig directly — but it has the same top_p/top_k/min_p fields, and async_grpo_trainer.py's log_ratio = log_probs - old_log_probs (line 1113) is the same mechanism: old_log_probs is collated straight from the rollout's vLLM logprobs (collator, line 572) with no vllm_importance_sampling_correction-style toggle to opt out of. So the bias isn't just uncovered by this warning there, it's unconditional — every AsyncGRPO run with truncated sampling has a biased coef_1 = torch.exp(log_ratio) feeding the clipped loss, with nothing to disable and no message pointing at why.

Given AsyncGRPO is arguably the more exposed case (no correction flag to turn off), might be worth either duplicating this guard into AsyncGRPOConfig.__post_init__ (if it has one) or at least noting the scope gap explicitly in the PR description so it doesn't read as "all vLLM-backed RL trainers covered." Happy to help wire that up if useful.

[behroozazarkhalili]

Confirmed, all three, and the gap turns out to be worse than a missing warning.

Checked against this PR's head and vLLM 0.22.0:

• AsyncGRPOConfig is at async_grpo_config.py:22 and GRPOConfig at grpo_config.py:23, both extending _BaseConfig directly. They are siblings, so nothing is inherited.
• top_p:241, top_k:248, min_p:255.
• log_ratio at async_grpo_trainer.py:1113, old_log_probs collated at :559-572, and grep importance_sampling returns nothing in either async file.

You hedged on whether it has a __post_init__. It does, at :384, so there was somewhere to hang the guard.

What I had not appreciated until working through it is that the async bias lands in a different place than GRPO's. GRPO computes its own old_per_token_logps (grpo_trainer.py:2664) and differences it against vLLM's, so the bias sits on a correction term that vllm_importance_sampling_correction=False switches off. AsyncGRPO uses the vLLM logprob as the PPO denominator itself, so for an unchanged policy coef_1 reads the surviving mass instead of 1.0:

top_p	coef_1	error vs 1.0	outside (0.8, 1.2)?
1.00	1.0000	0.0000	no
0.90	0.9000	0.1000	no
0.80	0.8000	0.2000	no
0.70	0.7000	0.3000	yes

Below 1 - epsilon_low the clip engages on every negative-advantage token whether or not the policy moved.

Two things came out of it. #6944 adds the warning on the AsyncGRPO side, and its CI is green. #6945 raises the design question underneath, which is that TRL imposes --logprobs-mode processed_logprobs on every server it prescribes and justifies it by the importance-sampling correction, a correction AsyncGRPO does not have.

Your comment also caught something in this PR itself. repetition_penalty reaches the same logprobs, because penalties are applied at sampler.py:403 before sample() computes them, and the guard here listed only the three truncation knobs. Fixed in 5107c7d with a test.

On the offer to help, the useful piece is measurement. That table is derived from the two normalizations and computed exactly, not observed on hardware, and confirming it end to end needs a two-GPU async run against a real vLLM server that I have not managed to schedule. If you have a lane for that, #6945 is the place to put 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 d51684c. Configure here.}

[cursor]

AsyncGRPO misses truncation warning

Medium Severity

The new truncation-bias warning lives only in GRPOConfig.__post_init__, but AsyncGRPOConfig extends _BaseConfig directly and never inherits it. AsyncGRPOTrainer still differences trainer log-softmax against vLLM old_log_probs in log_ratio, so top_p, top_k, min_p, and repetition_penalty bias that ratio the same way with no warning. Repository guidance requires duplicated trainer logic to stay aligned across copies.

 

^{Triggered by project rule: ../.ai/AGENTS.md}

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

[behroozazarkhalili]

Correct that AsyncGRPOConfig does not inherit this block. That is deliberate: the async trainer has no vllm_importance_sampling_correction, so the message here ("biases sampling/sampling_logp_difference and the correction built on it") would be false for it, and its own guard needs different wording. That half is #6944, which adds the warning to AsyncGRPOConfig.__post_init__ with the async-specific message and covers top_p, top_k, min_p and repetition_penalty; #6945 tracks the design question of whether AsyncGRPO should get the correction itself. Keeping the two PRs separate lets each be reviewed against its own trainer's semantics.