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

[behroozazarkhalili]

Fixes #7011.

What

policy_loss was logged at two scales depending on loss_type. The grpo, sapo, bnpo, dr_grpo and luspo branches capture it before the gradient accumulation rescale, so the logged number is the optimizer-window loss. The cispo, dapo, vespo branch folds current_gradient_accumulation_steps / steps_per_generation into its token-count normalizer before the capture, so it logged one micro-batch contribution instead. The capture in that branch now multiplies back by the accumulation factor in train mode (1.0 in eval). The returned loss, and therefore training, is unchanged.

Test

test_policy_loss_uses_optimizer_window_scale, parametrized over cispo, dapo, vespo, builds a 3-sample batch by hand, recomputes the per-token objective and the entropy term from _get_per_token_logps_and_entropies, and checks both the returned loss and the logged policy_loss at gradient_accumulation_steps=2, steps_per_generation=4.

Verification

On a compute node (job 58031566):

• fixed tree: 3 passed;
• base grpo_trainer.py: 3 failed, each logged value exactly half the expected one;
• mutant using steps_per_generation as the factor: 3 failed;
• mutant rescaling the returned loss as well: 3 failed;
• ruff 0.13.3 check and format: clean.

Scope

• policy_loss is appended only when the entropy bonus is enabled, so the issue's reproduction as written (no entropy_coef) raises IndexError on this revision; the test sets entropy_coef=0.2.
• gspo_token carries the same normalizer without a policy_loss capture; Re-sync gspo_token _compute_loss with GRPOTrainer #6856 resyncs its _compute_loss from this file and will pick the corrected block up. GMPO and the replay-buffer trainer have no affected branch.

---

Note

Low Risk
Observability-only fix for logged metrics; training loss computation and gradients are unchanged.

Overview
Fixes inconsistent policy_loss metrics for cispo, dapo, and vespo: those branches folded gradient accumulation and steps_per_generation into the loss normalizer before logging, so the metric reflected a micro-batch slice while other loss types logged the full optimizer-window value.

GRPOTrainer._compute_loss now multiplies the detached policy_loss by current_gradient_accumulation_steps in train mode (1.0 in eval). The scalar loss returned for backprop is unchanged.

Adds test_policy_loss_uses_optimizer_window_scale, parametrized over the three loss types, to assert the returned loss and logged policy_loss match hand-computed expectations with gradient_accumulation_steps=2 and steps_per_generation=4.

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

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