[GRPO] Fix entropy bonus normalization inconsistency across loss types

[YaseenBashaT]

Description

Current behavior: cispo/dapo/vespo divide the entropy bonus by a
global token count borrowed from the policy-loss normalizer, while every
other loss type divides by the tokens that actually contributed
(effective_mask.sum()). These only agree when nothing gets filtered out.
Once top_entropy_quantile < 1.0 drops most tokens, the global count
doesn't shrink with it, so the bonus gets silently crushed for those three
loss types — measured ~4.8x smaller at top_entropy_quantile=0.2 in one
worked example (ratio ρ ≈ top_entropy_quantile ≈ 0.208; a 10x crush would
need top_entropy_quantile≈0.1, not 0.2 — correcting the earlier estimate
in this description).

This also contradicts what the docs already promise
(docs/source/grpo_trainer.md, "Entropy regularization" section):

The bonus is always the mean per-token entropy regardless of loss_type;
it is not rescaled to match a loss type's policy normalization... so
entropy_coef has the same meaning for every loss type.

Changes: Drop the per-loss_type branch. The entropy term now always
divides by effective_mask.sum(), then by the gradient-accumulation
factor. One formula, and it now matches what the docs already said.

Accepted regression (intentional): cispo/dapo/vespo previously
got an exact token-weighted global average of entropy across the whole
gradient-accumulation window, because the old denominator summed real
token counts across every accumulated micro-batch. After this change they
get an unweighted mean of per-micro-batch, per-rank means instead — the
same aggregation every other loss type already used. This is deliberate:
uniform treatment across loss types is what the docs promise, and the two
aggregations only diverge meaningfully when token counts vary a lot across
micro-batches/ranks.

Merge note: please merge this after #6654, not just rebase on top of
it — the entropy mask is one of several things that broadcast
per_token_loss to (B, T) for luspo (see #6654), and this PR's new
luspo @ top_entropy_quantile=0.2 case depends on that fix being in to
be meaningful.

Test plan

Before

test_entropy_bonus_scale only ran with top_entropy_quantile=1.0 — the
one setting where the bug can't show up, since nothing gets filtered.

After

Added top_entropy_quantile=0.2 to the same test's parametrization, across
all 4 tested loss_types. Also changed the assertion itself: comparing
contrib / entropy to a fixed entropy_coef is only exact at
top_entropy_quantile=1.0 (where effective_mask == mask, so H_eff ==
H_full). Under filtering it's entropy_coef * H_eff / H_full, and the
logged entropy metric is H_full (global_masked_mean over the full
completion mask), not H_eff — so asserting the ratio equals entropy_coef
at top_entropy_quantile=0.2 was fixture-lucky: it only held because this
tiny, near-untrained model's entropy is ~uniform across tokens
(H_eff/H_full ≈ 1.0), which isn't true in general. The test now asserts
the actual invariant instead: every loss type's ratio must agree with
every other's, at the same quantile.

Verified in both directions: reverted just the trainer hunk (restoring the
old per-loss_type branching) and reran. top_entropy_quantile=0.2 fails
as expected (1 failed, 1 passed) — that's the case that actually
exercises the bug. top_entropy_quantile=1.0 still passes even on the
reverted code, which is correct and not a gap in the test: at quantile
1.0 nothing is filtered, so the old and new formulas coincide by
construction, matching the "Before" note above that this is the one
setting where the bug can't show up. Restored the fix afterward and
confirmed both cases pass (2 passed). Full suite (145 passed, 41
skipped, 0 failed) was run before this test was rewritten to its current
parametrization; the rewrite collapsed what used to be 8 separate
parametrized cases into 2 (each now looping over all 4 loss types
internally), so that count is no longer the right one to cite.
test_entropy_bonus_scale itself: 2 passed, 0 failed, directly re-run
after the rewrite.

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

---

Note

Medium Risk
Changes training loss math for cispo/dapo/vespo (especially with top_entropy_quantile < 1), which can shift exploration strength in production GRPO runs.

Overview
Unifies GRPO entropy regularization so cispo/dapo/vespo no longer scale the bonus with the policy loss’s global token normalizer. The entropy term is always the mean over effective_mask tokens (including top_entropy_quantile filtering), divided only by the gradient-accumulation factor—matching documented behavior that entropy_coef means the same for every loss_type.

Regression test updates: test_entropy_bonus_scale now runs at top_entropy_quantile 1.0 and 0.2, trains each of grpo/dr_grpo/dapo/luspo in one loop, and checks that (policy_loss - loss) / entropy is consistent across loss types rather than equaling entropy_coef (which only lines up with logged entropy when no quantile filtering).

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

[Copilot]

Pull request overview

Unifies GRPO entropy bonus normalization so the entropy regularizer is computed as a mean over the tokens it acts on (via effective_mask) and scaled only by gradient accumulation, removing the prior loss_type-specific normalization behavior for cispo/dapo/vespo.

Changes:

• Simplifies entropy bonus computation in GRPOTrainer._compute_loss by removing the loss_type branch and always normalizing by effective_mask.sum().
• Updates test_entropy_bonus_scale to also exercise entropy masking by parametrizing top_entropy_quantile with 1.0 and 0.2.

Reviewed changes

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

File	Description
trl/trainer/grpo_trainer.py	Removes loss_type-dependent entropy bonus normalization and applies a single mean-over-effective_mask formula scaled by gradient accumulation.
tests/test_grpo_trainer.py	Extends the entropy scaling regression test to cover cases where top_entropy_quantile < 1.0 activates entropy masking.

Suppressed comments (1)

tests/test_grpo_trainer.py:1845

• With top_entropy_quantile < 1.0, the loss’s entropy regularizer is computed over effective_mask (filtered tokens), but this test still divides by the logged entropy metric, which in GRPOTrainer is computed with global_masked_mean(entropies) using the unfiltered mask. That makes contrib / entropy == entropy_coef generally false under entropy masking, so this parametrization is likely to fail or not test the intended invariant. Consider either logging an “effective entropy” metric matching effective_mask (and using that here), or adjusting the test to compute the same entropy quantity used in the loss.

    @pytest.mark.parametrize("top_entropy_quantile", [1.0, 0.2])
    @pytest.mark.parametrize("loss_type", ["grpo", "dr_grpo", "dapo", "luspo"])
    def test_entropy_bonus_scale(self, loss_type, top_entropy_quantile):
        # Regression test: the entropy bonus is the mean per-token entropy H for every loss type (documented
        # objective L = L_policy - entropy_coef * H), so it must not inherit any loss-type-specific policy
        # normalization. A previous "unified" formula divided H by a global token count for the
        # cispo/dapo/vespo family, making the bonus ~1/sequence_length too small; conversely, scaling the
        # bonus like the dr_grpo (fixed budget) or luspo (sequence-weighted) policy term would also be wrong.
        # With gradient_accumulation_steps=1 the per-step entropy contribution to the loss is
        # contrib = policy_loss - loss = entropy_coef * entropy_loss, so contrib / entropy must equal
        # entropy_coef for all loss types.

---

💡 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]

I investigated this against #6140, where the current two-branch split comes from. Conclusion: the bug is real and the fix is the right direction — and it does not contradict #6140. Both rest on the same principle (the bonus is H, the mean per-token entropy, independent of loss_type); they differ only on whether the cispo/dapo/vespo branch implements it. It does — but only when top_entropy_quantile == 1.0.

What the branch was actually doing. It was not applying policy normalization. It existed because normalizer means different things: the grad-accumulation step count for the grpo family, a global token count for cispo/dapo/vespo (built from num_items_in_batch, i.e. the same loss_mask as mask). Summing the entropies there accumulates over the optimizer step to the same H. That is exact whenever effective_mask == mask.

Where it breaks. In #6140, 2f34d156 restricted the numerator to effective_mask (addressing "entropy bonus ignores quantile mask"), and 76255d30 reintroduced the branch a day later. The cispo/dapo/vespo denominator still counts every token, so under filtering that branch optimizes ρ·H_eff instead of H_eff, with ρ = N/M ≈ top_entropy_quantile. The interaction was never re-examined. So: a real bug, but strictly scoped to loss_type ∈ {cispo, dapo, vespo} × top_entropy_quantile < 1.0, and the magnitude is 1/ρ — I measured ρ = 0.208 at top_entropy_quantile=0.2, i.e. ~4.8×, not the 10× in the description (that would be quantile=0.1). Please rescope the description accordingly: the deleted comment was making a true claim in the unfiltered case.

What settles the direction is #6140's own documentation and controller, not the general principle:

• docs/source/grpo_trainer.md (the Tip added by Add entropy regularization to GRPO #6140): "When using top_entropy_quantile < 1.0, entropy_target applies to the high-entropy token subset — that subset's entropy will be higher than the logged full-token entropy."
• The adaptive controller accumulates [(entropies * effective_mask).sum(), effective_mask.sum()], so world_entropy = H_eff.

The buggy branch therefore desyncs the optimized quantity from the gated one — the same failure mode #6140 rejected the Dr. GRPO rescaling for. This PR restores that invariant.

Verified locally. Pre-fix, only dapo-0.2 fails; dapo-1.0 and grpo-0.2 pass, confirming no behaviour change at quantile=1.0 and that the grpo family was already correct. Post-fix, 8/8 pass.

Three things before merging:

1. The new parametrization is fixture-lucky. It divides by the logged entropy metric (global_masked_mean over the full mask) while the bonus uses effective_mask, so the ratio is entropy_coef · H_eff/H_full. On the tiny model I measured H_eff/H_full = 1.0000 (entropy ≈ log|V| ≈ 11.93 at essentially every token), so it passes by accident; on a real model that ratio is 2–5× and the assertion is false. This is @copilot's suppressed comment and it is correct. It still catches the dapo bug, but the comment's "contrib / entropy must equal entropy_coef for all loss types" is no longer true under filtering. Better assertion: ratio(dapo, q) ≈ ratio(grpo, q) — cross-loss-type consistency is the actual invariant.
2. Merge after [GRPO] Apply the completion mask elementwise in the luspo loss aggregation #6654. Not just a rebase: [GRPO] Apply the completion mask elementwise in the luspo loss aggregation #6654's own comment names the entropy mask as one of the things that broadcasts per_token_loss to (B, T), and the new luspo-0.2 case here is exactly that configuration.
3. Accepted regression, worth naming. cispo/dapo/vespo lose exact token-weighted global averaging (now an unweighted mean of per-micro-batch, per-rank means, like every other loss type). Fine — uniformity is what the docs promise — and the follow-up commit correctly drops the "matches world_entropy" claim. Keep that commit.

[YaseenBashaT]

Thanks for going this deep on it, seriously appreciate the effort here.

You're right on the 4.8x vs 10x, that was just sloppy math on my part in the original writeup, fixed it. Also agree on the test, comparing against a fixed entropy_coef only worked because this tiny model's entropy is basically flat across every token, so it was passing for the wrong reason. Switched it to check that every loss type lands on the same ratio as grpo instead, which is the thing that's actually supposed to hold.

Added a note in the description calling out the cispo/dapo/vespo averaging change on purpose so it doesn't look like a random regression later, and flagged that this should merge after #6654 rather than just rebasing on it.

And yeah, good catch that Copilot's suppressed comment was actually right, I'd glossed over that one.

[albertvillanova]

#6654 is merged (2396dfe5, now the tip of main), so the blocker on this one is cleared.

The new test_entropy_bonus_scale is the right call — asserting that every loss type lands on the same ratio is the invariant this PR actually restores, and it no longer depends on the fixture's near-uniform entropy. Description corrections look good too.

One thing left before I approve: please merge main into this branch — don't rebase. Rebasing rewrites the commits on the public PR branch and loses the review history that's attached to them; a merge commit keeps it intact.

git fetch upstream main
git merge upstream/main
git push

I trial-merged it locally and it applies cleanly — no conflicts, despite both PRs touching the area right around test_entropy_bonus_scale. I also ran the affected tests on that trial merge — test_entropy_bonus_scale (both quantiles) and test_luspo_loss_ignores_padding (both params): 4 passed. So the merge should be all that's needed, no follow-up fixes.

This matters beyond staying current: your new luspo @ top_entropy_quantile=0.2 case runs the entropy mask through the luspo aggregation, which is exactly the (B, T) broadcast path #6654 fixed. Until that fix is on this branch, that parametrization exercises the old, broken aggregation.

Once it's merged and CI is green I'll approve.

[YaseenBashaT]

Done, merged main in (not rebased) and pushed. Ran test_entropy_bonus_scale and test_luspo_loss_ignores_padding on the merge myself too, all 4 pass, matches what you found on your trial merge.