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

docs/advanced.md

@@ -6,6 +6,8 @@ This page covers the specialized features layered on top of the core training st
 
 - [Custom Modeling](#custom-modeling)
   - [Expert Parallelism Backends](#expert-parallelism-backends)
+  - [MoE Load Balancing](#moe-load-balancing)
+  - [Block Dropout](#block-dropout)
 - [Multimodal Training](#multimodal-training)
   - [Supported Families](#supported-families)
   - [Enabling VLM Mode](#enabling-vlm-mode)
@@ -50,6 +52,22 @@ DeepEP requires some careful tuning to achieve optimal performance, tuning param
 
 With DeepEP, gradient clipping is currently not supported. (`optim.max_norm` is set to `None` automatically.)
 
+### MoE Load Balancing
+
+Custom MoE models can rebalance expert usage during training via `[model.moe_load_balance]`. This is off by default (a pretrained router bias, if present, is still used for routing).
+
+```toml
+[model.moe_load_balance]
+mode = "loss_free"   # "off" (default) | "loss_free"
+coeff = 1e-3         # bias update step size
+```
+
+**`loss_free`** applies the DeepSeek auxiliary-loss-free update ([arXiv:2408.15664](https://arxiv.org/abs/2408.15664)): each optimizer step it nudges the router `expert_bias` toward balanced usage by `coeff · sign(load_error)` (the delta is recentered to zero mean so the bias vector doesn't drift a global offset), with no extra loss term. Expert token counts are summed across the data-parallel (`dp_cp`) group so balancing targets the whole global batch rather than each rank's shard.
+
+### Block Dropout
+
+`model.dropout` (default `0.0`) applies dropout at each block's output before the residual add — attention block and FFN/MoE block. It is intended for SFT/self-distillation (e.g. `0.15`); leave it at `0.0` for RL. Requires the custom model implementation.
+
 ## Multimodal Training
 
 ### Supported Families

packages/prime-rl-configs/src/prime_rl/configs/trainer.py

@@ -116,10 +116,31 @@ class DebugModelConfig(BaseConfig):
     """Replace MoE token-choice routing with a round-robin assignment so every expert sees an equal share. Intended for fake-data smoke tests where untrained routing would otherwise OOM under severe imbalance. Gating scores are still gathered from the override indices so the forward pass stays consistent."""
 
 
+class MoELoadBalanceConfig(BaseConfig):
+    """Training-time MoE load balancing (only affects the custom MoE implementation).
+
+    Auxiliary-loss-free bias balancing (DeepSeek, https://arxiv.org/abs/2408.15664). Expert usage
+    is aggregated across the data-parallel group so balancing targets the whole global batch rather
+    than each rank's local shard.
+    """
+
+    mode: Literal["off", "loss_free"] = "off"
+    """``off`` disables training-time balancing (a pretrained ``expert_bias`` is still used for routing if present). ``loss_free`` updates the router ``expert_bias`` each optimizer step from expert usage."""
+
+    coeff: float = Field(1e-3, gt=0)
+    """Bias update step size for ``loss_free`` (``expert_bias += coeff * sign(load_error)``)."""
+
+
 class ModelConfig(BaseModelConfig):
     seq_len: int = 2048
     """Sequence length the model is trained on."""
 
+    dropout: float = Field(0.0, ge=0, le=1)
+    """Dropout applied at each block's output before the residual add (attention block and FFN/MoE block). Intended for SFT/self-distillation (e.g. 0.15); leave at 0 for RL. Requires the custom model implementation."""
+
+    moe_load_balance: MoELoadBalanceConfig = MoELoadBalanceConfig()
+    """Training-time MoE load-balancing configuration."""
+
     attn: AttnImplementation = "flash_attention_2"
     """Attention implementation. With CP enabled, ring attention uses the matching kernel family (FA2/FA3/FA4)."""
 

src/prime_rl/trainer/model.py

@@ -11,6 +11,7 @@
 
 import torch
 import torch._dynamo
+import torch.distributed as dist
 import torch.nn as nn
 from huggingface_hub import snapshot_download
 from jaxtyping import Int
@@ -469,6 +470,47 @@ def get_load_balance_stats(
     }
 
 
+def set_block_dropout(model: nn.Module, dropout_p: float) -> None:
+    """Set the pre-residual block-output dropout probability on every supported block (attention,
+    dense MLP, MoE, LatentMoE). Dropout is applied to each block's output before the residual add.
+    ``dropout_p == 0.0`` is a no-op at runtime."""
+    for module in model.modules():
+        if hasattr(module, "dropout_p"):
+            module.dropout_p = dropout_p
+
+
+def set_moe_load_balance_active(model: nn.Module, active: bool) -> None:
+    """Enable/disable the loss-free bias-update usage accumulator on every MoE layer. Called once at
+    setup so ``expert_load_acc`` only grows when ``update_expert_bias`` is actually being run."""
+    for module in model.modules():
+        if hasattr(module, "moe_lb_active"):
+            module.moe_lb_active = active
+
+
+@torch.no_grad()
+def update_expert_bias(model: nn.Module, update_rate: float, dp_group=None) -> None:
+    """Auxiliary-loss-free load balancing (DeepSeek, https://arxiv.org/abs/2408.15664).
+
+    Nudges each MoE layer's router ``expert_bias`` toward balanced usage by ``update_rate`` in the
+    direction of the sign of the load error, using the usage accumulated since the last call
+    (``expert_load_acc``) summed across ``dp_group`` for the global batch. Call once per optimizer
+    step. The bias affects routing selection only, not gating.
+    """
+    language_model = get_language_model(model)
+    for transformer_block in language_model.layers:
+        block_mlp = getattr(transformer_block, "mlp", None)
+        if block_mlp is None or getattr(block_mlp, "expert_bias", None) is None:
+            continue
+        counts = block_mlp.expert_load_acc
+        if dp_group is not None:
+            counts = counts.clone()
+            dist.all_reduce(counts, op=dist.ReduceOp.SUM, group=dp_group)
+        delta = update_rate * torch.sign(counts.mean() - counts)  # positive where under-loaded
+        delta = delta - delta.mean()  # zero-mean so the bias vector doesn't drift a global offset
+        block_mlp.expert_bias.add_(delta)
+        block_mlp.expert_load_acc.zero_()
+
+
 def get_model(
     config: ModelConfig, device: torch.device = torch.device("cpu"), dtype: torch.dtype = torch.bfloat16
 ) -> nn.Module:
@@ -529,6 +571,11 @@ def get_model(
     model_config.use_grouped_mm = config.moe_use_grouped_mm
     model_config.fp8 = config.fp8
 
+    # Loss-free load balancing updates the router expert_bias buffer, which several models gate
+    # behind load_balance_coeff (defaults to None). Force it on so the buffer is constructed.
+    if config.moe_load_balance.mode == "loss_free" and getattr(model_config, "load_balance_coeff", None) is None:
+        model_config.load_balance_coeff = config.moe_load_balance.coeff
+
     if config.index_cache is not None:
         model_config.use_index_cache = True
         model_config.index_topk_freq = config.index_cache.topk_freq
@@ -644,6 +691,18 @@ def get_model(
     assert model.lm_head.weight.dtype == dtype, (
         f"LM head dtype wasnt loaded correctly {model.lm_head.weight.dtype} != {dtype}"
     )
+
+    # Fail loudly rather than silently no-op if loss-free balancing was requested but the model has
+    # no router expert_bias (e.g. MiniMax with use_routing_bias=False, which ignores load_balance_coeff).
+    if config.moe_load_balance.mode == "loss_free":
+        language_model = get_language_model(model)
+        if not any(
+            getattr(getattr(block, "mlp", None), "expert_bias", None) is not None for block in language_model.layers
+        ):
+            raise ValueError(
+                "moe_load_balance.mode='loss_free' requires the model's MoE routers to have an expert_bias "
+                "buffer, but this model was built without one. Loss-free balancing cannot run for it."
+            )
     return model
 
 

src/prime_rl/trainer/models/afmoe/modeling_afmoe.py

@@ -84,6 +84,8 @@ def __init__(self, config: AfmoeAttentionConfig):
         self.k_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
         self.v_proj = nn.Linear(config.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)
         self.o_proj = nn.Linear(self.num_heads * self.head_dim, config.hidden_size, bias=False)
+        # Pre-residual block-output dropout. Set via set_block_dropout(); 0.0 is a no-op.
+        self.dropout_p = 0.0
 
         # Output gating
         self.gate_proj = nn.Linear(config.hidden_size, self.num_heads * self.head_dim, bias=False)
@@ -102,7 +104,7 @@ def output_proj(
             attn_output = attn_output.transpose(1, 2).contiguous()
         attn_output = attn_output.contiguous().view(*input_shape, -1)
         attn_output = attn_output * torch.sigmoid(gate_states)
-        return self.o_proj(attn_output)
+        return F.dropout(self.o_proj(attn_output), p=self.dropout_p, training=self.training)
 
 
 class AfmoeSDPAAttention(AfmoeAttentionBase):

src/prime_rl/trainer/models/glm_moe_dsa/sparse_mla_attention.py

@@ -158,6 +158,8 @@ def __init__(self, args: SparseMlaAttentionArgs):
         )
 
         self.o_proj = nn.Linear(self.num_heads * self.v_head_dim, args.hidden_size, bias=args.attention_bias)
+        # Pre-residual block-output dropout. Set via set_block_dropout(); 0.0 is a no-op.
+        self.dropout_p = 0.0
         # IndexShare (GLM-5.2): layers that reuse cached indices carry no indexer weights.
         self.indexer = Indexer(args) if not args.skip_topk else None
         self.use_index_cache = args.use_index_cache
@@ -231,7 +233,7 @@ def output_proj(self, attn_output: torch.Tensor, w_v: torch.Tensor) -> torch.Ten
         attn_output = self._mla_unabsorb(attn_output, w_v)
         batch_size, total_tokens = attn_output.shape[:2]
         attn_output = attn_output.reshape(batch_size, total_tokens, -1)
-        return self.o_proj(attn_output)
+        return nn.functional.dropout(self.o_proj(attn_output), p=self.dropout_p, training=self.training)
 
     def forward(
         self,

src/prime_rl/trainer/models/gpt_oss/modeling_gpt_oss.py

@@ -135,6 +135,9 @@ def __init__(self, config: GptOssConfig, layer_idx: int):
         self.hidden_size = config.hidden_size
         self.self_attn = GptOssAttention(config=config, layer_idx=layer_idx)
         self.mlp = GptOssMoE(config)
+        # Pre-residual block-output dropout (attention and MoE are bespoke, not the shared layers).
+        # Set via set_block_dropout(); 0.0 is a no-op.
+        self.dropout_p = 0.0
         self.input_layernorm = GptOssRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
         self.post_attention_layernorm = GptOssRMSNorm(config.hidden_size, eps=config.rms_norm_eps)
 
@@ -159,11 +162,13 @@ def forward(
             position_embeddings=position_embeddings,
             **kwargs,
         )
+        hidden_states = F.dropout(hidden_states, p=self.dropout_p, training=self.training)
         hidden_states = residual + hidden_states
 
         residual = hidden_states
         hidden_states = self.post_attention_layernorm(hidden_states)
         hidden_states, _ = self.mlp(hidden_states)
+        hidden_states = F.dropout(hidden_states, p=self.dropout_p, training=self.training)
         hidden_states = residual + hidden_states
         return hidden_states
 

src/prime_rl/trainer/models/laguna/modeling_laguna.py

@@ -135,7 +135,10 @@ def forward(
         attn_output = attn_output.view(*input_shape, self.num_heads, self.head_dim)
         gate = F.softplus(self.g_proj(hidden_states).float()).to(attn_output.dtype)
         attn_output = (attn_output * gate.unsqueeze(-1)).view(*input_shape, -1)
-        return self.o_proj(attn_output), None
+        # Pre-residual block-output dropout (see base FlashAttention/SDPAAttention); dropout_p is set
+        # by set_block_dropout and is 0.0 (no-op) unless model.dropout is enabled.
+        attn_output = F.dropout(self.o_proj(attn_output), p=self.dropout_p, training=self.training)
+        return attn_output, None
 
 
 class LagunaSDPAAttention(SDPAAttention):
@@ -185,7 +188,10 @@ def forward(
         attn_output = attn_output.view(*input_shape, self.num_heads, self.head_dim)
         gate = F.softplus(self.g_proj(hidden_states).float()).to(attn_output.dtype)
         attn_output = (attn_output * gate.unsqueeze(-1)).view(*input_shape, -1)
-        return self.o_proj(attn_output), None
+        # Pre-residual block-output dropout (see base FlashAttention/SDPAAttention); dropout_p is set
+        # by set_block_dropout and is 0.0 (no-op) unless model.dropout is enabled.
+        attn_output = F.dropout(self.o_proj(attn_output), p=self.dropout_p, training=self.training)
+        return attn_output, None
 
 
 def _get_laguna_attention(config: LagunaConfig, layer_idx: int):
@@ -273,6 +279,8 @@ def forward(
         if self.shared_expert is not None:
             bs, slen, dim = hidden_states.shape
             shared_output = self.shared_expert(mlp_input.view(-1, dim)).view(bs, slen, dim)
+            # Match the routed-path block-output dropout (applied inside MoE) on the shared-expert path.
+            shared_output = F.dropout(shared_output, p=self.mlp.dropout_p, training=self.training)
             hidden_states = hidden_states + shared_output
         hidden_states = residual + hidden_states
         return hidden_states

src/prime_rl/trainer/models/layers/attn.py

@@ -72,6 +72,8 @@ def __init__(self, config: AttentionConfig, flash_attn_version: int = 2):
             config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
         )
         self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.output_bias)
+        # Pre-residual block-output dropout. Set via set_block_dropout(); 0.0 is a no-op.
+        self.dropout_p = 0.0
         self.use_qk_norm = config.use_qk_norm
         self.qk_norm_type = config.qk_norm_type
         if self.use_qk_norm:
@@ -180,6 +182,7 @@ def forward(
             max_seqlen=max_seqlen,
         )
         attn_output = self.output_proj(attn_output)
+        attn_output = F.dropout(attn_output, p=self.dropout_p, training=self.training)
         return attn_output, None
 
 
@@ -203,6 +206,8 @@ def __init__(self, config: AttentionConfig):
             config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias
         )
         self.o_proj = nn.Linear(config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.output_bias)
+        # Pre-residual block-output dropout. Set via set_block_dropout(); 0.0 is a no-op.
+        self.dropout_p = 0.0
         self.use_qk_norm = config.use_qk_norm
         self.qk_norm_type = config.qk_norm_type
         if self.use_qk_norm:
@@ -277,6 +282,7 @@ def forward(
 
         attn_output = self._attention_core(query_states, key_states, value_states)
         attn_output = self.output_proj(attn_output)
+        attn_output = F.dropout(attn_output, p=self.dropout_p, training=self.training)
         return attn_output, None
 
 

src/prime_rl/trainer/models/layers/mlp.py

@@ -1,6 +1,7 @@
 from dataclasses import dataclass
 
 import torch
+import torch.nn.functional as F
 from torch import nn
 from transformers.activations import ACT2FN
 
@@ -23,7 +24,9 @@ def __init__(self, config: MLPConfig):
         self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)
         self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)
         self.gate_act_fn = ACT2FN[config.gate_act]
+        # Pre-residual block-output dropout. Set via set_block_dropout(); 0.0 is a no-op.
+        self.dropout_p = 0.0
 
     def forward(self, x, routed_experts: torch.Tensor | None = None):
         down_proj = self.down_proj(self.gate_act_fn(self.gate_proj(x)) * self.up_proj(x))
-        return down_proj
+        return F.dropout(down_proj, p=self.dropout_p, training=self.training)