[AmineDiro]

What does this PR do?

Port the MoE load-balancing auxiliary loss (already supported in SFTTrainer) to GRPOTrainer, RLOOTrainer, and the experimental AsyncGRPOTrainer.

When a model is configured with output_router_logits=True and a router_aux_loss_coef, the policy loss now gets auxilary router loss added to it, matching the model's own reference forward and SFTTrainer. This keeps experts balanced when post-training MoE models.

Enable it via model_init_kwargs, e.g.:

GRPOConfig(..., model_init_kwargs={"output_router_logits": True, "router_aux_loss_coef": 0.001})

How it works

• GRPO / RLOO: the gradient-carrying forward in _get_per_token_logps_and_entropies now optionally requests output_router_logits=True and returns the model-computed outputs.aux_loss (averaged across micro-batch chunks). _compute_loss folds it into the loss and logs aux_loss.
• AsyncGRPO: its patched chunked-LM-head forward bypasses the MoE causal-LM forward, so it recomputes
  load_balancing_loss_func from router_logits directly (same as SFT's chunked path). Adds model_init_kwargs to AsyncGRPOConfig.
• Grad-accum scaling: the aux term is scaled to match each trainer's loss-normalization regime. Divided by gradient-accumulation steps for GRPO/AsyncGRPO (which disable HF auto-scaling), added unscaled for RLOO (which relies on HF auto-scaling)
• A new aux_loss metric is logged when enabled.

Notes

• The liger-kernel path (compute_liger_loss) is unchanged and does not apply the aux loss, consistent
  with SFT.

---

Note

Medium Risk
Touches shared _get_per_token_logps_and_entropies return shape and training loss math across multiple RL trainers; incorrect aux scaling or MoE forward flags could affect MoE training stability, while dense models are unaffected when aux is disabled.

Overview
Adds MoE router load-balancing auxiliary loss to RL trainers so post-training on mixture-of-experts models can keep expert usage balanced, aligned with existing SFTTrainer behavior.

GRPOTrainer / RLOOTrainer: _get_per_token_logps_and_entropies now optionally requests output_router_logits and returns a third value, aux_loss. Init reads output_router_logits and router_aux_loss_coef from the text config (VLMs via get_text_config()), warns when logits are on but coef is 0, and _compute_loss adds router_aux_loss_coef * aux_loss with scaling matched to each trainer’s grad-accum regime (GRPO divides by accumulation steps; RLOO adds unscaled). aux_loss is logged when enabled.

AsyncGRPOTrainer: New model_init_kwargs on AsyncGRPOConfig (passed into from_pretrained). The chunked LM-head path recomputes aux loss via load_balancing_loss_func in patch_chunked_lm_head when MoE is enabled; policy loss includes the aux term with the same grad-accum scaling as GRPO.

Call-site updates: Experimental GRPO variants (DPPO, GFPO, GMPO, etc.) unpack the new 3-tuple from _get_per_token_logps_and_entropies. SFTTrainer only gains a warning when output_router_logits is set but coef is zero.

Liger-kernel loss paths are unchanged (no aux term), consistent with SFT.

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

[cursor]

DAPO aux grad-accum mismatch

Medium Severity

For cispo, dapo, and vespo, policy loss is normalized with num_items_in_batch and is not divided by current_gradient_accumulation_steps, but the MoE auxiliary term is always divided by that step count. With gradient accumulation greater than one, the router aux contribution is scaled down relative to the policy loss and no longer matches the intended router_aux_loss_coef calibration.

 

^{Reviewed by Cursor Bugbot for commit 7d97669. Configure here.}

[qgallouedec]

The /grad_accum on the aux term is correct here too; it's what keeps it G-invariant like the
policy loss.

cispo/dapo/vespo skip /grad_accum on the policy loss because their normalizer num_items_in_batch is already the global accumulated-batch token count (same value on every micro-batch, Trainer auto-scaling disabled). So summing the G micro-batch losses is already G-invariant.

The aux loss is a per-micro-batch O(1) scalar, so to add it once per effective batch (not G times) it must be /G. Quick check (G=2, per-batch token-loss sum S over n items, aux A):

• single forward: S/n + coef·A
• accumulated: policy S/2n + S/2n = S/n, aux coef·A/2 + coef·A/2 = coef·A → same

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

[qgallouedec]

not for this pr, but I think it would be nicer to have a dedicated arg. I'll open a separate pr

[qgallouedec]

#6085

[AmineDiro]

Agreed, just followed the how SFTTrainer does it

[AmineDiro]

Should we wait on #6085 before merging this ?

[qgallouedec]

Agreed, just followed the how SFTTrainer does it

you did right!

Should we wait on #6085 before merging this ?

no, this one should be merged first

[AmineDiro]

🫡

[qgallouedec]

cool!!

[cursor]

Cursor Bugbot has reviewed your changes using default effort and found 1 potential issue.

There are 2 total unresolved issues (including 1 from previous review).

^{❌ Bugbot Autofix is OFF. To automatically fix reported issues with cloud agents, enable autofix in the Cursor dashboard.}

^{Reviewed by Cursor Bugbot for commit fbb9ba7. Configure here.}

[cursor]

Missing detach on aux_loss before distributed metrics gathering

Low Severity

The aux_loss tensor participates in the computation graph (it's added to loss for backpropagation) but is passed directly to gather_for_metrics without calling .detach(). The async GRPO trainer correctly uses aux_loss.detach().to(torch.float32) before its reduce call. Passing a grad-tracking tensor through distributed communication is wasteful and inconsistent with the pattern used elsewhere in this same PR.

Additional Locations (1)

• trl/trainer/rloo_trainer.py#L1468-L1469

 

^{Reviewed by Cursor Bugbot for commit fbb9ba7. Configure here.}

[DaoyuanLi2816]

Thanks for adding this @AmineDiro! A scope question: the experimental GRPO variants that override _compute_loss — GMPOTrainer, DPPOTrainer, PAPOTrainer, and gspo_token's GRPOTrainer — were updated here only to unpack the new 3-tuple from _get_per_token_logps_and_entropies (discarding the aux loss as _), so their overridden _compute_loss doesn't add the aux term.

They still inherit aux_loss_enabled / router_aux_loss_coef from GRPOTrainer.__init__, so a MoE model with output_router_logits=True is accepted (and the router_aux_loss_coef == 0 warning logic runs), but the auxiliary loss is never actually applied for these trainers.

Is that intentional for now (MoE aux loss scoped to GRPO/RLOO/AsyncGRPO), or would a follow-up PR propagating it to these variants be welcome? Happy to send one.