Re-sync gspo_token _compute_loss with GRPOTrainer

[behroozazarkhalili]

Re-syncs trl/experimental/gspo_token/grpo_trainer.py::_compute_loss against the current GRPOTrainer._compute_loss, as reported in #6807.

Credit to @mmjerge, who found this and offered to send the patch.

What was wrong

The override is a copy of GRPOTrainer._compute_loss plus one branch. It stopped tracking its source, so four GRPOConfig options, four VLM forward kwargs and one metric silently did nothing on this trainer. All of it reproduces at 47d20ec6:

Setting	State in the copy
use_bias_correction_kl	KL never multiplied by coef_1
off_policy_mask_threshold	get_off_policy_mask call and its application both absent
entropy_coef and the adaptive controller	the whole _entropy_bonus_enabled block absent
router_aux_loss_coef	compute_aux_loss kwarg and the aux term absent
VLM forward kwargs	mm_token_type_ids, image_position_ids, spatial_shapes, num_tiles absent
the sequence_token error message	omitted the value the file itself accepts

The issue did not list the policy_loss metric, which GRPOTrainer logs and this copy does not.

use_bias_correction_kl failed differently from the rest. It multiplies the per-token KL by coef_1, and the copy computed the KL forty lines before coef_1 existed, so the option was unreachable rather than missing. Restoring canonical statement order is what fixes it.

What this changes

_compute_loss is now the canonical method at 47d20ec6, lines 3109 to 3379, with two divergences: the sequence_token branch of importance_sampling_level, and the ValueError message naming that third value. Checking that claim takes one command, and it returns those two hunks and nothing else:

SRC=trl/trainer/grpo_trainer.py
GSPO=trl/experimental/gspo_token/grpo_trainer.py
diff <(git show "47d20ec6:$SRC" | sed -n '3109,3379p') <(sed -n '22,297p' "$GSPO")

The diff is large relative to the bug because the method was rebuilt from its source instead of patched feature by feature. Patching is how it drifted: #6024, #6380, #6439 and #6083 each touched this file and each left drift behind.

Verification

• Tests (experimental) green in CI, and the 9 gspo_token tests pass on a compute node against the branch. test_matches_grpo_trainer_outside_the_sequence_token_branch compares the subclass with the unbound GRPOTrainer._compute_loss on the same trainer and inputs at importance_sampling_level="token", in loss, gradient and logged metrics, once per restored option (bias-corrected KL, off-policy mask, entropy bonus); the pre-resync file fails all three rows, while the earlier gradient test passes on it. Check code quality green, which is ruff@0.13.3 plus the pinned doc-builder; I ran both locally too, the doc-builder one with a negative control confirming it can still fail.
• The GSPO-token math is untouched. The branch is byte-identical to the one it replaces, torch.equal on the weight tensor returns True, and the stop-gradient identity d/dlogp exp(w) == exp(seq_weight) holds.
• Each defect was quantified by driving both code paths on identical synthetic tensors under importance_sampling_level="sequence_token":

option                      stale       re-synced      rel
use_bias_correction_kl    0.00516653   0.00515584    -0.21%
entropy bonus 0.01        0.00516653  -0.00414616  -180.25%
MoE aux loss 2.0          0.00516653   0.00716653   +38.71%
OPSM threshold 0.05       0.00516653   0.09534809  +1745.49%

• Metamorphic control: with every drifted option off, the two paths agree exactly at all three importance_sampling_level values, so the table isolates the options rather than the rewrite.
• sequence_token yields (B, T), the same shape as "token", so the newly reachable cispo, sapo, vespo and luspo branches broadcast correctly.

Two things reported rather than fixed

gmpo has the same gaps. trl/experimental/gmpo/gmpo_trainer.py is also missing the bias-correction KL, OPSM, the entropy bonus, the aux loss, and spatial_shapes/num_tiles. It is not a mechanical re-sync: GMPO's objective is sequence-level, coef is (B,), and there is no loss_type branching, so use_bias_correction_kl and OPSM have no direct translation. Happy to open a separate PR once you decide what those should mean there.

vespo ignores importance_sampling_level. get_gamma_weights receives log_ratio_per_token=log_ratio, so phi_seq never sees log_importance_weights or coef_1. Measured across all three levels the vespo loss is identical to twelve decimal places. GRPOTrainer behaves the same way for "sequence", so per the AGENTS.md consistency rule I reproduced it rather than adding a guard here that would diverge from the source.

Note on conflicts

#6637 and #6667 both touch this file's KL block, and this re-sync moves that block. Whichever merges second will conflict. I will rebase mine.

Resolves #6807

---

Note

Medium Risk
Large rewrite of RL training loss logic affects gradients and optional regularizers; risk is mitigated by parity tests against GRPOTrainer and unchanged sequence_token math.

Overview
Re-aligns GSPOTokenTrainer._compute_loss with the current GRPOTrainer implementation so drifted GRPOConfig behavior works again, while keeping the GSPO-token-only importance_sampling_level="sequence_token" branch unchanged.

The override again applies bias-corrected KL (use_bias_correction_kl), off-policy masking (off_policy_mask_threshold), entropy regularization (static and adaptive), MoE router aux loss, extra VLM forward kwargs, and the full set of loss types and metrics (e.g. policy_loss, CISPO/VESPO clip logging) that the stale copy had dropped or computed in the wrong order.

Tests add gradient checks that sequence_token routes per-token advantages differently from sequence, and a parametrized parity test that GSPOTokenTrainer matches GRPOTrainer on loss, gradients, and logged metrics when importance_sampling_level="token" and each restored option is enabled.

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

[mmjerge]

@behroozazarkhalili looks good thanks!

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