GRPO: vLLM importance-sampling ratio is biased when top_p/top_k/min_p truncate sampling

[qgallouedec]

When top_p < 1, top_k > 0 or min_p is set, GRPO's vLLM importance-sampling correction compares two log-probabilities computed over different distributions, so the ratio is biased. vLLM returns log-probs renormalized over the truncated (nucleus) distribution, while the trainer recomputes them over the full vocabulary.

Two other RL frameworks have now shipped a training-side fix for exactly this, and the native vLLM fix is merged but not yet in a release TRL supports.

Where the mismatch is

Generation side asks vLLM for post-processor log-probs, in both colocate and server mode:

https://github.com/huggingface/trl/blob/main/trl/generation/vllm_generation.py#L365

Training side divides by temperature and takes a full-vocab log-softmax:

https://github.com/huggingface/trl/blob/main/trl/trainer/grpo_trainer.py#L1530-L1534

processed_logprobs means "after every logits processor", which includes top-p/top-k truncation followed by renormalization over the surviving tokens. Temperature is handled symmetrically (line 1532 mirrors it, and that symmetry was the fix for #4159); truncation is not. For a sampled token that survives the nucleus, vLLM's log-prob is therefore higher than the trainer's by -log(sum of surviving probability mass), which is a per-token, systematically negative bias in log π_θ - log π_behaviour.

importance_sampling_ratio multiplies the per-token loss directly:

https://github.com/huggingface/trl/blob/main/trl/trainer/grpo_trainer.py#L3208-L3209

so the effect is a silent, token-dependent down-weighting of the policy gradient rather than a crash.

Scope

• Defaults are safe: top_p=1.0, top_k=0, min_p=None, so nothing is truncated and the two distributions agree.

• AsyncGRPOTrainer is affected more severely. Server launch passes --logprobs-mode processed_logprobs, and those generator log-probs are not an optional correction there: they are the denominator of the policy ratio itself, with no flag guarding them.

  log_ratio = log_probs - old_log_probs   # async_grpo_trainer.py#L945
  coef_1 = torch.exp(log_ratio)

  AsyncGRPOConfig exposes top_p/top_k/min_p like GRPO, and it logs no sampling_logp_difference equivalent, so there is nothing to notice the drift with.

• Affected as soon as a user sets any of them, which top_p=0.9/top_k=50 recipes sometimes do.

• Only active with vllm_importance_sampling_correction=True.

• Same logprobs_mode="processed_logprobs" is set in trl/scripts/vllm_serve.py#L361, so server mode is affected identically.

It is already observable

The existing metric spikes under truncation, since it measures precisely this gap:

sampling/sampling_logp_difference/mean and .../max (grpo_trainer.py, _generate_and_score_completions)

A quick check: run GRPO+vLLM with top_p=1.0 and then top_p=0.8, and compare that metric. If it jumps with no other change, the bias is real. Worth doing before we pick a fix.

Prior art

• vLLM #49577 "Mask Replay" is merged (2026-08-13) but landed after v0.27.1 (2026-08-11), so it is not in any release covered by TRL's vllm>=0.18.0,<=0.27.1 pin. We cannot rely on it yet.
• slime shipped a training-side mask in v0.3.1.
• prime-rl #3235 "feat: top-p/top-k train sampling with sampling replay" is building the same thing independently.

Possible directions

1. Apply the same truncation to the training-side logits before the log-softmax, so both sides normalize over the same support. Mirrors what slime and prime-rl do, and keeps the temperature-symmetry precedent from vLLM not computing correct log probs when using GRPO with =/=1 temperature #4159.
2. Ask vLLM for raw log-probs and apply temperature and truncation ourselves on both sides.
3. Wait for Mask Replay and lift the ceiling to vLLM 0.28. Doesn't help anyone on a supported version today.
4. Do nothing to the math, but warn when vllm_importance_sampling_correction=True is combined with truncation, so the bias is at least not silent.

I'd like a repro of the metric jump before committing to one.

Possibly related

• vLLM not computing correct log probs when using GRPO with =/=1 temperature #4159
• vllm_importance_sampling_correction causes vllm_importance_sampling_ratio explosion #4772

[Bhoomi-jain]

Hi, I would like to work on this issue.

[behroozazarkhalili]

The bias has a closed form. The importance ratio is multiplied by exactly S, the probability mass that survives truncation.

The trainer takes a full-vocab log-softmax of logits / temperature (grpo_trainer.py:1539), no truncation. vLLM is configured with logprobs_mode="processed_logprobs" (vllm_generation.py:365, vllm_serve.py:361), which renormalizes over the truncated support, so it returns log(p_t / S). The difference is therefore

log_ratio = log p_t - log(p_t / S) = log S

for every token, independent of which token was sampled. exp(log_ratio) = S, so every per-token loss is scaled by the surviving mass.

What that predicts numerically

Simulated over a 32k-vocab distribution, with S measured directly and the closed form checked against it:

setting	S	log S (the bias)	IS ratio	gradient kept
top_p=1.0 (default)	1.000000	0.000000	1.0000	100.0%
top_p=0.95	0.950004	-0.051289	0.9500	95.0%
top_p=0.9	0.900009	-0.105350	0.9000	90.0%
top_p=0.8	0.800020	-0.223119	0.8000	80.0%
top_k=50	0.171154	-1.765189	0.1712	17.1%
top_k=20	0.108383	-2.222085	0.1084	10.8%
top_p=0.9 + top_k=50	0.167632	-1.785983	0.1676	16.8%

Control: at top_p=1.0, S == 1 exactly, so the bias vanishes. That matches the default config being unaffected.

The two knobs behave differently, and the difference matters for how alarming each one is.

top_p is model-independent. By construction nucleus sampling keeps just enough mass to reach the threshold, so S ≈ top_p regardless of how peaked the distribution is, and the bias is log(top_p). A sweep over logit scales from 1.0 to 12.0 (entropy 9.87 nats down to 0.76) gives S at top_p=0.9 between 0.9000 and 0.9513.

top_k depends on how confident the model is. Same sweep, top_k=50:

logit scale	entropy (nats)	S @ top_p=0.9	S @ top_k=50	S @ top_k=20
1.0	9.873	0.9000	0.0254	0.0130
2.0	8.355	0.9000	0.1726	0.1118
4.0	4.011	0.9005	0.7813	0.6750
6.0	1.970	0.9114	0.9703	0.9299
8.0	1.432	0.9226	0.9951	0.9804
12.0	0.756	0.9513	0.9999	0.9987

So top_k is nearly harmless on a well-trained, confident model and severe on an uncertain one. Early in RL training, when the policy is still flat, is exactly when the bias is largest.

Measured on GPU

The three numbers were fixed before the run, so this is a pre-specified prediction rather than a fit. Setup: GRPOTrainer with use_vllm=True, vllm_mode="colocate", vllm_importance_sampling_correction=True, on trl-internal-testing/small-Qwen2ForCausalLM-2.5, one H100, max_completion_length=16, two steps per config, one process per top_p.

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 through a different route. With 16 completion tokens the sequence-level correction is S^16, and sampling/importance_sampling_ratio/mean reads:

top_p	S^16	measured ratio (two steps)
1.0	1.000000	1.0000, 0.9999
0.9	0.185302	0.1855, 0.1849
0.8	0.028147	0.02814, 0.0281

Both quantities agree to four decimals, which fixes the per-token surviving mass at exactly S = top_p. At top_p=0.8 the correction weights a 16-token completion at 2.8 percent of its uncorrected value, and the exponent is the completion length, so the factor shrinks geometrically as completions get longer.

Scope of the measurement: one GPU, colocate mode, a small test model, short completions. The derivation is model-independent for top_p, because nucleus sampling keeps whatever mass the threshold names, so the same reading should hold at any scale. I have measured it only in the setting above.

[behroozazarkhalili]

Two corrections to the issue text, both from reading the source, plus a note on what the recompute option costs.

The bias enters at one subtraction, and it is not the policy ratio

For GRPOTrainer the policy ratio is not affected. grpo_trainer.py:3172 computes log_ratio = per_token_logps - old_per_token_logps, and both operands are trainer forwards over the full vocabulary. old_per_token_logps comes from _get_per_token_logps_and_entropies at grpo_trainer.py:2664, not from vLLM.

vLLM's logprobs enter at exactly one place, grpo_trainer.py:2680:

per_token_logps_diff = (old_per_token_logps - sampling_per_token_logps) * mask

That is log p_hf - log p_vllm + log S. The first two terms are the train/inference mismatch the correction exists to measure; log S is the contamination. So the effect is confined to vllm_importance_sampling_ratio, and the fix has to touch a tensor consumed only there. Masking old_per_token_logps in place would put log S into the policy ratio instead, which is worse than the current behavior.

The description is accurate for AsyncGRPOTrainer, where the server's logprobs really are the denominator at async_grpo_trainer.py:1113. That case needs a different fix and I would treat it separately.

Scope note: vllm_importance_sampling_correction exists in grpo_config.py, grpo_trainer.py, experimental/gspo_token, and experimental/grpo_with_replay_buffer. rloo_trainer.py has no importance-sampling correction at all, so it is unaffected despite having the same logits / temperature then selective_log_softmax shape at lines 983 to 985.

Direction 1 is not what slime and prime-rl do

Both do replay, carrying the kept-token set from the sampler to the trainer, rather than recomputing the nucleus from training-side logits.

prime-rl #3235 is titled "top-p/top-k train sampling with sampling replay" and adds src/prime_rl/inference/vllm/kept_tokens.py, 268 lines that append the kept ids to each LogprobsTensors row in the engine-core worker so the trainer can, in their words, "replay these sets to renormalize its own logprobs identically". They cite DeepSeek V3.2 "Keep Sampling Mask" (arXiv:2512.02556 3.1) and cap train-sampling top_k at 512 so capture never overflows and replay is exact at every position.

slime v0.3.1 shipped the same shape: PR 2102 added top-p masking "so training can reproduce the effective sampling distribution used during rollout", and PR 2167 made rollout_top_p_token_ids mandatory whenever rollout_top_p is not 1.0.

Both therefore need the kept set over the wire, which is what vLLM #49577 adds and why prime-rl patched vLLM internals rather than waiting.

What recompute would cost

Recomputing the mask from the trainer's logits is implementable under the current vllm>=0.18.0,<=0.27.1 pin, but it has a failure mode worth naming before anyone builds it: the sampled token can fall outside the recomputed nucleus. vLLM kept it and sampled it, but the trainer's logits differ numerically, so a token at the boundary can be masked to -inf, giving an infinite difference and a NaN batch. Any recompute implementation needs to force-include the sampled token in the kept set. That is not fatal, but it makes recompute an approximation that both reference implementations chose to avoid.

To match vLLM exactly the recomputed mask would also have to follow v1/sample/sampler.py: temperature first, then min_p (it is registered as argmax-invariant, so it runs before the rest), then top_k, then top_p with the ascending cumsum <= 1 - p rule and the max force-kept.

Meanwhile

I opened #6880 for direction 4 only, warning when vllm_importance_sampling_correction=True is combined with top_p < 1, top_k > 0, or min_p > 0. It is one guard in GRPOConfig.__post_init__, so it reaches the experimental GRPO variants through inheritance, and it forecloses none of 1, 2, or 3. Happy to close it if you would rather ship the math fix directly.

On the direction itself, if the pin can move to a vLLM release carrying #49577 then direction 3 is the same choice both reference projects made. If it cannot move for a while, recompute with the force-include guard is the only correct-today option, and I am glad to build it.

@Bhoomi-jain, you said you wanted to work on this. The warning PR is deliberately the small half. If you are taking the math fix, the numbers and source references above should save you some time, and I am happy to review.