Skip to content

Log GRPO policy_loss at the optimizer-window scale for every loss type - #7042

Closed
behroozazarkhalili wants to merge 4 commits into
mainfrom
fix/7011-policy-loss-scale
Closed

behroozazarkhalili wants to merge 4 commits into
mainfrom
fix/7011-policy-loss-scale

Conversation

@behroozazarkhalili

@behroozazarkhalili behroozazarkhalili commented Sep 4, 2026 •

Copy link
Copy Markdown
Collaborator

GRPO logged policy_loss at different scales depending on the loss type. I restored the optimizer window scale for CISPO, DAPO, and VESPO during training while leaving the returned loss unchanged. The regression tests independently compute the objective and check both values, failing on the base code and passing with the fix.
Fixes #7011.


Note

Low Risk
Metrics-only adjustment with no change to the backward loss; limited to GRPO logging for three loss types plus a unit test.

Overview
Fixes inconsistent policy_loss metrics for cispo, dapo, and vespo in GRPOTrainer._compute_loss. Those branches folded gradient accumulation and steps_per_generation into the loss normalizer before recording policy_loss, so logs reflected a single micro-batch contribution while other loss types logged the pre–gradient-accumulation (optimizer-window) scale.

The change multiplies the detached policy_loss by current_gradient_accumulation_steps in train mode (1.0 in eval) after that normalization. Returned loss and optimization are unchanged—only what gets appended to _metrics["train"]["policy_loss"] is rescaled.

Adds test_policy_loss_uses_optimizer_window_scale, parametrized over the three loss types, with hand-built inputs and explicit expectations for both _compute_loss output and the logged policy_loss under gradient_accumulation_steps=2 and steps_per_generation=4.

Reviewed by Cursor Bugbot for commit bd5fb1f. Bugbot is set up for automated code reviews on this repo. Configure here.

…ss type

The grpo, sapo, bnpo, dr_grpo and luspo branches capture policy_loss
before dividing by the gradient accumulation normalizer, so the logged
value is the optimizer-window loss. The cispo, dapo and vespo branch
folds that normalizer into its token-count denominator before the
capture, so its policy_loss was one micro-batch contribution, half the
value of the other branches at gradient_accumulation_steps=2. The
capture now multiplies back by the accumulation factor in train mode;
the returned loss is unchanged.

The regression test recomputes each of the three objectives and the
entropy term from the model's own log probabilities and checks both the
returned micro-batch loss and the logged policy_loss. On the previous
code every row reports exactly half the expected value.

Fixes #7011
@bot-ci-comment

bot-ci-comment Bot commented Sep 4, 2026

Copy link
Copy Markdown

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.

@behroozazarkhalili

Copy link
Copy Markdown
Collaborator Author

Closing this one myself. It came from reading the code, not from a failure in a real training run, and per the discussion on #6862 that isn't the bar for this repo. If someone hits this in practice, the branch is still here and I'm happy to reopen it with a reproduction.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

policy_loss is logged at two different scales depending on loss_type

1 participant