[GRPO] Apply the completion mask elementwise in the luspo loss aggregation

[YaseenBashaT]

Description

Current behavior: the luspo branch in GRPOTrainer._compute_loss is the only loss
type that never multiplies per_token_loss by mask elementwise. It uses mask purely
as a per-sequence length:

elif self.loss_type == "luspo":
    loss = (per_token_loss * mask.sum(1, keepdim=True)).mean()

That is only safe while per_token_loss has shape (B, 1). Four reachable code paths
broaden it to (B, T) before this line: beta != 0.0 (the KL term is (B, T)),
use_vllm with a token-level importance sampling mode, top_entropy_quantile < 1.0
(the entropy mask), and importance_sampling_level="token", which is the config
default and, for luspo, only ever produces a logger.warning, never an error.

The more serious consequence, not just padding leaking into the loss: once
per_token_loss is (B, T), the old formula quietly changes what's being weighted.
Writing x for per_token_loss, L_b for the completion length, and T for the
padded width:

old: (x * mask.sum(1, keepdim=True)).mean() = mean_b( L_b * mean_over_T(x_b) )
new: (x * mask).sum(-1).mean()               = mean_b( L_b * mean_over_valid(x_b) )

so the old form weights each sequence by L_b^2 / T instead of L_b, a quadratic
length bias, which is exactly the length bias LUSPO exists to remove. With beta > 0
the KL penalty also ends up dominated by padding noise, and its strength varies with
the batch's padding fraction instead of being a stable regularizer.

Changes: mask elementwise before aggregating:

elif self.loss_type == "luspo":
    # `per_token_loss` is (B, 1) only in the recommended sequence-level setup;
    # importance_sampling_level="token" (the config default), the KL term, token-level
    # vLLM IS ratios, and the entropy mask all broadcast it to (B, T), so mask before
    # aggregating.
    loss = (per_token_loss * mask).sum(-1).mean()

This is numerically equivalent, not bit-identical, to the previous formula in the
recommended (B, 1) setup: the reduction order differs, so on a random (B, 1) input
in fp32 the old formula and the new one can land a few ULPs apart (e.g.
-5.6097946167 vs -5.6097936630 on one such input). It's also only unaffected in
that setup while beta == 0 (the default). With the paper's recommended
importance_sampling_level="sequence" and beta > 0, per_token_loss is already
(B, T) before this line, so the KL term's effective scale moves too: from a mean
over all T slots including padding, to a mean over only the valid ones. That's the
point of the fix, but it means training runs using LUSPO with a KL penalty will see a
real (intended) change in the loss's magnitude, not just its correctness.

Test plan

Before

No test constructed _compute_loss inputs directly for luspo, so there was no way
to isolate padded positions and confirm the loss doesn't depend on them.

After

Added test_luspo_loss_ignores_padding, parametrized over two configurations that
each broadcast per_token_loss to (B, T) by a different path:
importance_sampling_level="sequence" with beta=0.1 (the KL term), and
importance_sampling_level="token" with beta=0.0 (the config default, and the case
with the L^2/T weighting bug, previously uncovered). Both build _compute_loss
inputs directly, perturb ref_per_token_logps/old_per_token_logps only at padded
(mask == 0) positions, and assert the loss is unchanged. old_per_token_logps and
ref_per_token_logps are anchored to the model's real per-token log-probs (plus a
small fixed offset) rather than arbitrary noise: for this tiny, near-uniform model,
unrelated noise either makes the importance-sampling ratio astronomically small or the
KL term astronomically large, and PPO-style clipping (widened here since it isn't what
this test is exercising) or assert_close's tolerance can then mask the exact
regression being tested for regardless of whether the fix is correct. Both
parametrized cases fail on unmodified code and pass after the fix; verified by
temporarily reverting just the trainer hunk and confirming both fail.

Existing luspo coverage (test_train_loss_types[luspo],
test_entropy_bonus_scale[luspo-*]) still passes locally in the recommended
(B, 1), beta=0 setup, consistent with the fix being a no-op there.

Merge note: please land this before #6648, not after. That PR adds
top_entropy_quantile=0.2 to test_entropy_bonus_scale[luspo-*], which is precisely
one of the broadcast paths this PR's fix covers, so that new case would otherwise run
against the broken aggregation. There's also a textual conflict: both PRs insert new
code right after test_entropy_bonus_scale in tests/test_grpo_trainer.py.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline, Pull Request section?
• Was this discussed/approved via a GitHub issue? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?

AI writing disclosure

• No AI usage: the PR was written entirely by a human.
• AI-assisted: some parts were suggested or improved by AI, but the PR was written and reviewed by a human.
• AI-generated: the PR was mostly or fully generated by an AI tool.

Who can review?

Anyone in the community is free to review the PR once the tests have passed.

---

Note

Medium Risk
Changes GRPO training objective math for LUSPO whenever per-token loss is broadcast to (B, T), so loss magnitude and gradients can shift (especially with KL); fix is correctness-focused but affects live training runs.

Overview
Fixes LUSPO loss aggregation in GRPOTrainer._compute_loss so the completion mask is applied elementwise over tokens, matching other loss types. The old path multiplied per_token_loss by mask.sum(1, keepdim=True) (sequence length only), which is wrong when per_token_loss is (B, T)—e.g. default importance_sampling_level="token", beta > 0 KL, token-level vLLM IS, or entropy masking. That let padding affect the loss and introduced an unintended (L^2/T) length weighting instead of averaging over valid tokens.

Aggregation is now (per_token_loss * mask).sum(-1).mean() (still divided by the existing gradient-accumulation normalizer). Adds test_luspo_loss_ignores_padding, which perturbs log-probs only on padded positions and asserts the scalar loss is unchanged, for both KL-driven and token-level IS setups.

^{Reviewed by Cursor Bugbot for commit 5ccc312. Bugbot is set up for automated code reviews on this repo. Configure here.}

[Copilot]

Pull request overview

This PR fixes a correctness issue in GRPOTrainer._compute_loss for loss_type="luspo" where padded completion positions could leak into the loss/gradients once per_token_loss becomes token-shaped (B, T) (e.g., via the KL term), and adds a regression test to ensure padding is ignored.

Changes:

• Update the luspo loss aggregation to apply the completion mask elementwise before reducing.
• Add a regression test that perturbs ref_per_token_logps only on padded tokens and asserts the loss is unchanged.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 2 comments.

File	Description
trl/trainer/grpo_trainer.py	Masks per_token_loss elementwise for luspo before aggregation to prevent padding leakage.
tests/test_grpo_trainer.py	Adds regression coverage to ensure luspo loss is invariant to changes in padded positions.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[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.

[albertvillanova]

Thanks — the bug is real and the fix is the right generalization. I verified locally that the new test passes on this branch and fails when the trainer hunk alone is reverted, so it's a genuine regression test. A few things before merging (details inline).

"bit-identical in the recommended case" isn't accurate. It's numerically equivalent, not bit-identical — the reduction order differs. On a random (B, 1) input in fp32 I get -5.6097946167 (old) vs -5.6097936630 (new). Please reword in the description; someone will eventually rely on that sentence.

The unaffected-scope claim needs a caveat. "the recommended (B, 1) setup is unchanged" holds only while beta == 0 (the default). With the paper's importance_sampling_level="sequence" and beta > 0 you are already in the (B, T) case, and the KL term's scale moves (from a mean over all T slots including padding, to a mean over the valid ones). That is the point of the PR, but people running LUSPO with a KL penalty should be able to read it off the description.

Merge order. Please land this before #6648. That PR adds top_entropy_quantile=0.2 to test_entropy_bonus_scale[luspo-*], which is precisely one of the broadcast paths your new comment names, so the new case would otherwise run against the broken aggregation. There is a textual conflict too — both insert right after test_entropy_bonus_scale.

Non-blocking:

• Consistency checked: luspo exists only in GRPOTrainer; RLOO has no loss_type branches and none of the experimental _compute_loss overrides carry a luspo branch. Nothing to mirror.
• This fixes the eager path only. With use_liger_kernel=True the aggregation happens inside LigerFusedLinearGRPOLoss via compute_liger_loss, which never reaches _compute_loss — worth a follow-up check that Liger masks before aggregating for luspo.
• I could not reproduce "existing luspo coverage still passes" as stated: test_entropy_bonus_scale[luspo] passes in isolation but was flaky in a batch run here, and test_train_loss_types[luspo-True] fails on the base commit in my environment (liger-kernel too old). Both unrelated to your change — just don't lean on that line as evidence.

[albertvillanova]

The description frames this as "padded positions leak in". True, but the more serious consequence is the weighting, and it is the strongest argument for the change — please put it in the description.

Writing x for per_token_loss, L_b for the completion length and T for the padded width:

old: (x * L).mean()            = mean_b( L_b * mean_over_T(x_b) )
new: (x * mask).sum(-1).mean() = mean_b( L_b * mean_over_valid(x_b) )

So once x is (B, T), the old form weights each sequence by L_b^2 / T instead of L_b — a quadratic length weighting, i.e. exactly the length bias LUSPO exists to remove.

[albertvillanova]

This drops the case the old comment named: importance_sampling_level="token", which is the config default and only warned about for luspo (grpo_trainer.py L901). Please keep it in the list.

(off_policy_mask is (B, 1), so it is correctly absent.)

[albertvillanova]

Dead line — remove it. trainer.model.eval() on the previous line makes _compute_loss take the mode == "eval" branch, so normalizer is 1.0 and current_gradient_accumulation_steps is never read.

[albertvillanova]

some CI environments bind Qwen2's RMSNorm to a Triton kernel that assumes CUDA and rejects CPU tensors outright

Did you actually hit this? The plain reason is sufficient and certainly true: Trainer.__init__ moves the model to args.device, so on a GPU runner the inputs have to be built there. Please state that instead — comments here should say what is true, not what might be.

(Dropping use_cpu=True is right, by the way: tests/conftest.py::force_use_cpu_without_accelerator already forces it on CPU-only machines.)

[albertvillanova]

Coverage is thinner than it looks. old_per_token_logps is not in the inputs, so it falls back to per_token_logps.detach(), coef_1 == 1 exactly, and the policy term collapses to -A. The test therefore only exercises the KL broadcast path.

importance_sampling_level="token" — the config default, and the case with the L^2/T weighting — is not covered at all. Cheap fix: parametrize over ("sequence", beta=0.1) and ("token", beta=0.0).

[albertvillanova]

This is the first test in this file to call _compute_loss with hand-built inputs. I'm fine with it — there is no other way to isolate padded positions — but please add a line to the test comment saying why, so a future refactor knows what it is pinning.

[YaseenBashaT]

Thanks for going through this so carefully, this is a lot of good catches.

Fixed the description on all the wording points: it now says numerically equivalent instead of bit-identical, with the actual fp32 numbers you gave, and I added the caveat that the (B, 1) case is only unaffected when beta is 0. I also pulled in your weighting explanation, the L_b^2/T vs L_b point is a much stronger way to explain why this matters than just "padding leaks in", so that's now the lead argument in the description instead of an afterthought.

On the code comment, you're right that dropping importance_sampling_level="token" from the list was a mistake, since it's the actual default and the most direct cause. Put it back.

Removed the dead current_gradient_accumulation_steps line, good catch, model.eval() already routes _compute_loss through the eval branch so that line was never doing anything.

Reworded the CI comment too. You're right that I was stating a guess instead of the actual reason, so it just says what's true now: Trainer.init moves the model to args.device, that's why the inputs need to go there too.

On coverage: I parametrized the test to also cover importance_sampling_level="token" with beta=0 like you suggested, but it turned out the naive version of that doesn't actually catch anything. Without old_per_token_logps set explicitly it defaults to per_token_logps.detach(), which makes the importance ratio exactly 1 everywhere regardless of padding, so the new case would have passed on both fixed and buggy code. Had to add old_per_token_logps explicitly and anchor it (and ref_per_token_logps) close to the model's real log-probs instead of raw noise, otherwise the ratio either vanishes or blows up at this model's log-prob scale and the difference falls under assert_close's tolerance either way. Also had to widen epsilon since PPO clipping was saturating and hiding the signal too. Verified both parametrized cases actually fail if I revert just the trainer fix, and pass with it in.

Added the merge-before-#6648 note and the textual-conflict heads up to the description too.

On the non-blocking points, agreed on all three, nothing to do on my end for the Liger path or RLOO, and noted your point about not leaning on the "still passes" line as strong evidence given what you saw in your environment, softened that wording as well.

[albertvillanova]

All the review points are addressed — description, comment, dead line, device rationale, coverage note. I re-verified locally: both parametrizations pass on this head and both fail when the trainer hunk alone is reverted, so the new case genuinely bites. CI is green too, which also explains away the two luspo failures I mentioned from my environment.

One thing left, inline.

[albertvillanova]

This comment was true of the previous revision, but the same commit made it false: L1943 now sets old_per_token_logps = baseline_logps + 0.05. It also contradicts the comment 25 lines below, which argues the opposite ("must be close in scale to the model's actual per-token log-probs, not arbitrary noise").

And the knob isn't needed: with the anchoring, the ratio is exp(-0.05) ≈ 0.95 unperturbed (inside the default [0.8, 1.2]) and exp(-1.05) ≈ 0.35 at perturbed padded positions — clipped to 0.8, but 0.8 ≠ 0.95, so the signal survives clipping anyway. I checked: with epsilon=1e6 removed, both parametrizations still fail against the pre-fix formula and pass with the fix.

So please just drop epsilon=1e6 and these four lines rather than rewording them — one less knob and no contradiction to maintain. (Worth re-confirming on your side; the token ids are unseeded, so that was one run.)

[albertvillanova]

Suggested change

	# Wide open on purpose: old_per_token_logps below is random and unrelated to per_token_logps, so the
	# importance-sampling ratio can easily land outside the default clip range. PPO-style clipping would
	# then saturate and mask the exact signal this test perturbs, independently of whether the padding fix
	# under test is correct. This test isolates masking, not clipping, so clipping is disabled here.
	epsilon=1e6,

[YaseenBashaT]

Good catch, and sorry about that, that comment was leftover from an earlier version of the test before I anchored old_per_token_logps to the real log-probs, and I didn't go back and clean it up once the reasoning behind it stopped being true. Should have caught that contradiction myself before pushing.

Dropped both the comment and the epsilon override like you suggested. Reran it 5 times each way since you flagged the token ids are unseeded: 5/5 pass with the fix in, 5/5 fail against the pre-fix formula with the default clip range, so it's solid without the extra knob.

Thanks again for how thorough this review was, genuinely caught real gaps.

[albertvillanova]

All points addressed. Verified locally on 5ccc312: both parametrizations pass with the fix and fail (3/3 runs) with the trainer hunk reverted under default clipping, so dropping the epsilon override was safe. Trainer diff is minimal and the comment now covers all four broadcast paths.

Please land this before #6648. Thanks for the careful iterations.

[YaseenBashaT]

Thanks for sticking with this through all the rounds, the padding fix, the wording fixes, and then catching that leftover comment right at the end. This is a much better PR for it than what I first opened. Appreciate the time you put into it.