Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 48 additions & 0 deletions tests/test_grpo_trainer.py
Original file line number Diff line number Diff line change
Expand Up @@ -2688,6 +2688,54 @@ def dummy_reward_func(completions, **kwargs):
expected_warning = "weights each sequence by its completion length"
assert (expected_warning in caplog.text) == (loss_type != "grpo")

@pytest.mark.parametrize(
("top_p", "top_k", "min_p", "should_warn"),
[
(1.0, 0, None, False), # defaults: nothing is truncated, the two distributions agree
(0.9, 0, None, True),
(1.0, 50, None, True),
(1.0, 0, 0.05, True),
(0.9, 50, 0.05, True),
],
)
def test_warning_raised_truncated_sampling_with_importance_sampling_correction(
self, top_p, top_k, min_p, should_warn
):
"""Truncated sampling biases the vLLM importance-sampling ratio.

vLLM returns logprobs renormalized over the surviving support while the trainer takes a full-vocab log-softmax,
so their difference carries `log S`, the log of the mass that survives truncation, on top of the
train/inference mismatch the correction is meant to measure. Warn instead of correcting silently.
"""
with warnings.catch_warnings(record=True) as caught:
warnings.simplefilter("always")
GRPOConfig(
output_dir=self.tmp_dir,
use_vllm=True,
vllm_importance_sampling_correction=True,
top_p=top_p,
top_k=top_k,
min_p=min_p,
report_to="none",
)

messages = [str(w.message) for w in caught]
assert any("biases `sampling/sampling_logp_difference`" in m for m in messages) == should_warn

def test_no_truncation_warning_without_importance_sampling_correction(self):
"""The bias only exists when the correction consumes vLLM's logprobs, so truncation alone must stay silent."""
with warnings.catch_warnings(record=True) as caught:
warnings.simplefilter("always")
GRPOConfig(
output_dir=self.tmp_dir,
use_vllm=True,
vllm_importance_sampling_correction=False,
top_p=0.9,
report_to="none",
)

assert not any("biases `sampling/sampling_logp_difference`" in str(w.message) for w in caught)

def test_train_num_generations_larger_than_batch_size(self):
dataset = load_dataset("trl-internal-testing/zen", "standard_prompt_only", split="train")

Expand Down
28 changes: 28 additions & 0 deletions trl/trainer/grpo_config.py
Original file line number Diff line number Diff line change
Expand Up @@ -1134,6 +1134,34 @@ def __post_init__(self):
)
self.vllm_importance_sampling_clip_max = self.vllm_importance_sampling_cap

if self.use_vllm and self.vllm_importance_sampling_correction:
# vLLM is asked for `processed_logprobs`, which are renormalized over the support that survives
# top_p/top_k/min_p, while the trainer takes a full-vocab log-softmax. Their difference therefore carries
# log(S), where S is the surviving mass, on top of the train/inference mismatch the correction exists to
# measure. Measured on one H100: `sampling/sampling_logp_difference/mean` reads |log(top_p)| exactly, i.e.
# 0.105 at top_p=0.9 and 0.223 at top_p=0.8 against 0.001 at top_p=1.0.
truncating = [
name
for name, active in (
("top_p", self.top_p < 1.0),
("top_k", self.top_k > 0),
("min_p", self.min_p is not None and self.min_p > 0.0),
)
if active
]
if truncating:
warnings.warn(
f"{' and '.join(truncating)} truncates sampling, which biases "
"`sampling/sampling_logp_difference` and the importance-sampling ratio derived from it: vLLM "
"renormalizes its logprobs over the surviving tokens while the trainer normalizes over the full "
"vocabulary, so their difference includes the log of the surviving probability mass. Set "
"`vllm_importance_sampling_correction=False`, or keep the sampling defaults (`top_p=1.0`, "
"`top_k=0`, `min_p=None`), until the correction accounts for truncation. See "
"https://github.com/huggingface/trl/issues/6789.",
UserWarning,
stacklevel=3,
)

if (
self.vllm_importance_sampling_clip_min is not None
and self.vllm_importance_sampling_clip_max is not None
Expand Down
Loading