feat(trainer): MoE loss-free load balancing and block dropout

[faresobeid]

Summary

Adds training-time MoE load balancing and per-block dropout to prime-rl's custom model implementation. Both are opt-in and off by default.

MoE load balancing

New [model.moe_load_balance] config (mode = off | loss_free, coeff), wired into both SFT and RL. Before this, load_balance_coeff only allocated the router expert_bias buffer (loaded from checkpoints and frozen) — there was no active balancing mechanism during training. (torchtitan implements this same path as an optimizer step-pre-hook; prime-rl never wired it up.)

loss_free — DeepSeek auxiliary-loss-free expert_bias update (arXiv:2408.15664), applied once per optimizer step: expert_bias += coeff · sign(load_error), from a dedicated per-step usage accumulator. Details:

• The accumulator only grows on training-mode forwards, so SFT validation traffic isn't folded into the update.
• The bias delta is recentered to zero mean (matching torchtitan) so the bias vector doesn't drift a global offset over long runs. This also makes the update invariant to the uniform per-layer count doubling caused by full-block activation checkpointing (the router reruns in backward).
• Expert counts are summed across the dp_cp group so balancing targets the whole global batch. Default coeff = 1e-3.
• get_model forces the expert_bias buffer to exist when loss_free is requested (several models default load_balance_coeff=None), and raises a clear error if a model can't provide one (e.g. MiniMax with use_routing_bias=False) rather than silently no-op'ing.

Block dropout

model.dropout (default 0.0) applies dropout at each block's output before the residual add — attention block and FFN/MoE block — set via a set_block_dropout walker. Covers all model families: the shared attention/MLP/MoE layers, plus the bespoke blocks (Laguna gated attention + decoder-level shared expert; Qwen3.5-MoE gated/linear-attn mixers + shared expert; GLM-MoE-DSA sparse MLA; AFMoE attention; GPT-OSS attention + MoE; NemotronH Mamba). Intended for SFT/self-distillation (e.g. 0.15); wired into the SFT trainer only.

SFT validation switches the model to eval() when dropout or loss-free balancing is active, so val/loss is deterministic and validation forwards don't skew the load-balance counts. With both off (the default), the model stays in train mode (no torch.compile recompilation).

Files

• configs/trainer.py — MoELoadBalanceConfig, ModelConfig.dropout
• trainer/model.py — set_block_dropout, update_expert_bias, loss-free expert_bias enable + validation
• trainer/models/layers/{moe,mlp,attn}.py — per-step load accumulator (train-only), dropout_p output dropout
• trainer/models/laguna/…, trainer/models/qwen3_5_moe/… — dropout on bespoke attention and decoder-level shared experts
• trainer/{sft,rl}/train.py — bias-update wiring; SFT dropout enable + eval-mode validation
• docs/advanced.md — MoE Load Balancing + Block Dropout sections

Testing

Unit-tested on CPU: update_expert_bias lowers over-loaded / raises under-loaded experts, its delta sums to zero, and the update is bit-identical for 1× vs 2× (AC-doubled) counts; dropout is a no-op in eval / at p=0; config surface validates (mode ∈ {off, loss_free}, coeff default 1e-3). ruff check + ruff format --check clean. Not yet exercised in a live distributed run.

🤖 Generated with Claude Code

---

Note

Medium Risk
Changes MoE routing via mutable expert_bias and touches many model forward paths; misconfiguration could skew expert usage, though features are off by default and unsupported models fail at load time.

Overview
Adds opt-in MoE load balancing ([model.moe_load_balance], loss_free) and per-block dropout (model.dropout) for the custom model stack, with docs and trainer wiring.

Loss-free balancing runs once per optimizer step in SFT and RL: update_expert_bias nudges router expert_bias from a train-only expert_load_acc counter, all-reduced over dp_cp, with a zero-mean delta. MoE layers gain moe_lb_active / expert_load_acc; get_model forces expert_bias allocation and errors if the architecture has no bias buffer.

Block dropout applies F.dropout on attention/MLP/MoE outputs before residuals across shared layers and bespoke blocks (Laguna, Qwen3.5-MoE, GLM sparse MLA, AFMoE, GPT-OSS, NemotronH Mamba); SFT calls set_block_dropout when dropout > 0.

SFT validation uses model.eval() when dropout or load balancing is on so val loss stays deterministic and validation forwards do not feed the balance accumulator.

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

[cursor]

Cursor Bugbot has reviewed your changes and found 2 potential issues.

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

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