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

Add training-time MoE load balancing and per-block dropout to the custom
model implementation.

MoE load balancing via [model.moe_load_balance] (mode/coeff), off by default.
mode='loss_free' applies the DeepSeek auxiliary-loss-free bias update
(arXiv:2408.15664) once per optimizer step: expert_bias += coeff * sign(load
error), from a per-step usage accumulator, with the delta recentered to zero
mean (matches torchtitan) so the bias vector doesn't drift a global offset.
Expert usage is aggregated across the dp_cp group so balancing targets the
whole global batch. Default coeff 1e-3. Before this, prime-rl only allocated
the expert_bias buffer (loaded from checkpoints and frozen) with no update.

model.dropout applies dropout at each block's output before the residual add
(attention + FFN/MoE), intended for SFT/self-distillation (e.g. 0.15);
default 0.0 so RL is unaffected. SFT validation switches the model to eval
mode when dropout is active so val/loss stays deterministic (skipped when
dropout is off to avoid torch.compile recompilation).

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: faresoPrime

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
@@ -446,6 +447,39 @@ 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
+
+
+@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:

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)

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

@@ -494,7 +494,6 @@ def forward(
             min=0,
             max=self.num_experts,
         )
-
         return top_scores, selected_experts_indices, num_tokens_per_expert, routing_confidence_sum
 
     def init_weights(self, init_std: float):
@@ -592,6 +591,8 @@ def __init__(self, moe_args: MoEArgs, dim: int, hidden_dim: int):
         )
         self.score_before_experts = moe_args.score_before_experts
         self.deepep_token_chunk_size: int | None = None
+        # Pre-residual block-output dropout. Set via set_block_dropout(); 0.0 is a no-op.
+        self.dropout_p = 0.0
 
         # define fields for auxiliary-loss-free load balancing (https://arxiv.org/abs/2408.15664)
         # NOTE: tokens_per_expert is accumulated in the model forward pass.
@@ -614,6 +615,13 @@ def __init__(self, moe_args: MoEArgs, dim: int, hidden_dim: int):
             persistent=False,
         )
         self.register_buffer("routing_confidence_sum", torch.tensor(0.0, dtype=torch.float32), persistent=False)
+        # Separate usage accumulator for the loss-free bias update (consumed once per optimizer step by
+        # update_expert_bias), kept distinct from tokens_per_expert which logging resets every micro-step.
+        self.register_buffer(
+            "expert_load_acc",
+            torch.zeros(num_experts, dtype=torch.float32),
+            persistent=False,
+        )
 
     def set_ep_comm_backend(self, backend: EPCommBackend) -> None:
         self.ep_comm_backend = backend
@@ -743,9 +751,15 @@ def forward(
         with torch.no_grad():
             self.tokens_per_expert.add_(num_tokens_per_expert)
             self.routing_confidence_sum.add_(routing_confidence_sum)
+            # Loss-free bias-update usage: training forwards only, so validation traffic (which runs
+            # before update_expert_bias in the SFT loop) is excluded. Full-block AC reruns this in
+            # backward, doubling counts uniformly per layer, which the sign-based update is invariant to.
+            if self.expert_bias is not None and self.training:
+                self.expert_load_acc.add_(num_tokens_per_expert)
 
         if self.ep_comm_backend == "deepep":
             routed_output = self._run_deepep_routed_experts(x, selected_experts_indices, top_scores)
+            routed_output = F.dropout(routed_output, p=self.dropout_p, training=self.training)
             return routed_output.reshape(bs, slen, dim)
 
         # top_scores and token_indices_experts_sorted shape (bs*slen*top_k,)
@@ -775,6 +789,7 @@ def forward(
 
         routed_indices = token_indices_experts_sorted.reshape(-1, 1).expand(-1, dim)
         out = out.scatter_add(dim=0, index=routed_indices, src=routed_output)
+        out = F.dropout(out, p=self.dropout_p, training=self.training)
         out = out.reshape(bs, slen, dim)
         return out
 
@@ -791,6 +806,7 @@ def init_weights(
         with torch.device(buffer_device):
             self.tokens_per_expert = torch.zeros(self.experts.num_experts, dtype=torch.float32)
             self.routing_confidence_sum = torch.tensor(0.0, dtype=torch.float32)
+            self.expert_load_acc = torch.zeros(self.experts.num_experts, dtype=torch.float32)
             if self.load_balance_coeff is not None:
                 self.expert_bias = torch.zeros(self.experts.num_experts, dtype=torch.float32)
 
@@ -1005,7 +1021,6 @@ def forward(
             min=0,
             max=self.num_experts,
         )
-
         return top_scores, selected_experts_indices, num_tokens_per_expert, routing_confidence_sum
 
     def init_weights(self, init_std: float):
@@ -1073,6 +1088,8 @@ def __init__(
         self.reorderer = TokenReorderer(num_experts=num_experts, top_k=top_k)
         self.shared_expert = BCNonGatedFeedForward(dim=dim, hidden_dim=shared_expert_intermediate_size)
         self.deepep_token_chunk_size: int | None = None
+        # Pre-residual block-output dropout. Set via set_block_dropout(); 0.0 is a no-op.
+        self.dropout_p = 0.0
 
         if latent_dim is not None:
             self.fc1_latent_proj = nn.Linear(dim, latent_dim, bias=False)
@@ -1098,6 +1115,13 @@ def __init__(
             persistent=False,
         )
         self.register_buffer("routing_confidence_sum", torch.tensor(0.0, dtype=torch.float32), persistent=False)
+        # Separate usage accumulator for the loss-free bias update (consumed once per optimizer step by
+        # update_expert_bias), kept distinct from tokens_per_expert which logging resets every micro-step.
+        self.register_buffer(
+            "expert_load_acc",
+            torch.zeros(num_experts, dtype=torch.float32),
+            persistent=False,
+        )
 
     def set_ep_comm_backend(self, backend: EPCommBackend) -> None:
         self.ep_comm_backend = backend
@@ -1204,9 +1228,15 @@ def forward(self, x: torch.Tensor, routed_experts: torch.Tensor | None = None) -
         with torch.no_grad():
             self.tokens_per_expert.add_(num_tokens_per_expert)
             self.routing_confidence_sum.add_(routing_confidence_sum)
+            # Loss-free bias-update usage: training forwards only, so validation traffic (which runs
+            # before update_expert_bias in the SFT loop) is excluded. Full-block AC reruns this in
+            # backward, doubling counts uniformly per layer, which the sign-based update is invariant to.
+            if self.expert_bias is not None and self.training:
+                self.expert_load_acc.add_(num_tokens_per_expert)
 
         if self.ep_comm_backend == "deepep":
             routed_output = self._run_deepep_routed_experts(x_flat, selected_experts_indices, top_scores)
+            routed_output = F.dropout(routed_output, p=self.dropout_p, training=self.training)
             return routed_output.reshape(bs, slen, dim)
 
         (
@@ -1226,6 +1256,7 @@ def forward(self, x: torch.Tensor, routed_experts: torch.Tensor | None = None) -
 
         token_indices_full = token_indices_experts_sorted.reshape(-1, 1).expand(-1, dim)
         out = out.scatter_add(dim=0, index=token_indices_full, src=routed_output)
+        out = F.dropout(out, p=self.dropout_p, training=self.training)
         out = out.reshape(bs, slen, dim)
         return out
 
@@ -1236,5 +1267,6 @@ def init_weights(self, init_std: float, buffer_device: torch.device):
         with torch.device(buffer_device):
             self.tokens_per_expert = torch.zeros(self.experts.num_experts, dtype=torch.float32)
             self.routing_confidence_sum = torch.tensor(0.0, dtype=torch.float32)
+            self.expert_load_acc = torch.zeros(self.experts.num_experts, dtype=torch.float32)
             if self.load_balance_coeff is not None:
                 self.expert_bias = torch.zeros(self.experts.num_experts, dtype=torch.float32)

src/prime_rl/trainer/rl/train.py

@@ -42,6 +42,7 @@
     setup_model,
     is_tt_moe_model,
     get_load_balance_stats,
+    update_expert_bias,
 )
 from prime_rl.trainer.parallel_dims import get_parallel_dims, resolve_ep
 from prime_rl.trainer.perf import get_perf_counter
@@ -543,6 +544,9 @@ def load_run_checkpoint(_optimizer, idx: int) -> None:
 
         # Update the model parameters
         optimizer.step()
+        if config.model.moe_load_balance.mode == "loss_free":
+            lb = config.model.moe_load_balance
+            update_expert_bias(model, lb.coeff, dp_cp_group)
         optimizer.zero_grad()
 
         # Update learning rate scheduler

src/prime_rl/trainer/sft/train.py

@@ -29,8 +29,10 @@
     forward,
     get_load_balance_stats,
     is_tt_moe_model,
+    set_block_dropout,
     setup_tokenizer,
     setup_model,
+    update_expert_bias,
 )
 from prime_rl.trainer.parallel_dims import get_parallel_dims, resolve_ep
 from prime_rl.trainer.perf import get_perf_counter
@@ -142,6 +144,10 @@ def train(config: SFTConfig):
     fused_cross_entropy: bool | str = {"liger_fused": "liger", "quack_fused": "quack"}.get(config.loss_impl, False)
     model = setup_model(config.model, parallel_dims, loading_from_ckpt_later, fused_cross_entropy=fused_cross_entropy)
 
+    if config.model.dropout > 0:
+        logger.info(f"Enabling per-block pre-residual dropout {config.model.dropout}")
+        set_block_dropout(model, config.model.dropout)
+
     if parallel_dims.cp_enabled:
         from prime_rl.utils.cp import assert_cp_style_supports_model
 
@@ -324,8 +330,19 @@ def run_validation(step: int) -> None:
         )
         val_dataloader = setup_dataloader(val_dataset, config.val.data)
 
-        # No train/eval switch: no dropout in these models, and toggling would trigger torch.compile recompilation
-        mean_loss, nan_count = run_eval_loop(val_dataloader)
+        # Toggling train/eval can trigger torch.compile recompilation, so only switch when it
+        # matters: block dropout must be disabled for a deterministic val metric, and loss-free load
+        # balancing accumulates expert usage in training mode only (so validation, which runs before
+        # update_expert_bias, doesn't skew the balance). With both off (the default) there is nothing
+        # to disable, so we keep the model in train mode.
+        toggle_eval = config.model.dropout > 0 or config.model.moe_load_balance.mode != "off"
+        if toggle_eval:
+            model.eval()
+        try:
+            mean_loss, nan_count = run_eval_loop(val_dataloader)
+        finally:
+            if toggle_eval:
+                model.train()
         if nan_count > 0:
             logger.warning(f"Validation at step {step}: {nan_count} batches had NaN loss")
         if mean_loss != mean_loss:
@@ -448,6 +465,9 @@ def run_validation(step: int) -> None:
 
         logger.debug("Optimizer step")
         optimizer.step()
+        if config.model.moe_load_balance.mode == "loss_free":
+            lb = config.model.moe_load_balance
+            update_expert_bias(model, lb.coeff, dp_cp_group)
         optimizer.zero_grad()
 
         # Update learning rate scheduler