Enabled auxilillary loss free load balancing and sequence wise load balancing for DSV4 [Replacement for PR-https://github.com/AI-Hypercomputer/maxtext/pull/4233]

docs/reference/core_concepts/moe_configuration.md

@@ -64,7 +64,12 @@ Dropping:
 
 `routed_bias`: If enabled, adds a learnable bias term to the gate logits to facilitate load balancing.
 
-`routed_bias_update_rate`: Defines the update rate to routed bias term above. Applicable only to the DeepSeek decoder block.
+`routed_bias_update_rate`: Defines the update rate to the routed bias term above. Applicable only to the DeepSeek decoder block. For DeepSeek V4, this enables a specialized, auxiliary-loss-free routing bias mechanism. This implementation utilizes a pure `nnx.Variable` (`MoEBiasVar`) instead of a standard `nnx.Param`, which completely isolates the bias update step from the global model optimizer state. The bias is updated directly at the end of the routing step to balance the token distribution mathematically across experts without compromising language modeling convergence.
+
+#### DeepSeek V4 Auxiliary-Loss-Free & Sequence-Wise Load Balancing
+MaxText implements an exact, paper-aligned version of DeepSeek V4's load balancing strategies (as specified in [the DeepSeek-V4 technical report](https://arxiv.org/html/2606.19348v1)). The architecture employs two distinct mechanisms:
+1. **Auxiliary-Loss-Free Strategy**: Handled via `MoEBiasVar`, this implementation utilizes a pure `nnx.Variable` instead of a standard `nnx.Param`. This strictly isolates the bias update step from the global model optimizer state. Unlike the Hugging Face reference implementation—which deviates from the paper by keeping the bias parameters coupled to the global optimizer—our implementation ensures the routing bias balances token distribution mathematically across experts without polluting the main model gradients.
+2. **Sequence-Wise Balance Loss**: An augmenting auxiliary loss (`load_balance_loss_weight`) applied to prevent extreme routing imbalance within individual sequences.
 
 `routed_score_func`: Defines the scoring function for the router.
 

src/maxtext/common/metric_logger.py

@@ -216,7 +216,7 @@ def _log_training_metrics(self, metrics, step):
 
     if self.config.num_experts > 1:
       moe_lb_loss = scalars.get("learning/moe_lb_loss", 0.0)
-      log_parts.append(f"moe_lb_loss: {moe_lb_loss:.3f}")
+      log_parts.append(f"moe_lb_loss: {moe_lb_loss:.6f}")
 
     if getattr(self.config, "mtp_num_layers", 0) > 0:
       mtp_loss = scalars.get("learning/mtp_loss", 0.0)

src/maxtext/configs/models/deepseek4-284b.yml

@@ -22,6 +22,7 @@ base_mlp_dim: 2048
 base_moe_mlp_dim: 2048
 vocab_size: 129280
 head_dim: 512
+qk_rope_head_dim: 64
 
 # --- Standard Defaults ---
 enable_dropout: false
@@ -51,6 +52,9 @@ shared_experts: 1
 routed_score_func: "sqrtsoftplus"
 norm_topk_prob: true
 routed_bias: true
+routed_bias_update_rate: 0.001
+load_balance_loss_weight: 0.0001
+log_moe_bias_norms: false
 routed_scaling_factor: 1.5
 
 

src/maxtext/configs/models/deepseek4-tiny.yml

@@ -11,16 +11,18 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-# Test Model config for DeepSeek-V4-Flash 284B (https://huggingface.co/deepseek-ai/DeepSeek-V4-Flash)
 
-base_emb_dim: 4096
-base_num_query_heads: 64
+# Tiny model config for DeepSeek V4 for CPU execution and testing
+
+base_emb_dim: 64
+base_num_query_heads: 4
 base_num_kv_heads: 1
-base_num_decoder_layers: 7
-base_mlp_dim: 2048
-base_moe_mlp_dim: 2048
+base_num_decoder_layers: 43
+base_mlp_dim: 64
+base_moe_mlp_dim: 64
 vocab_size: 129280
-head_dim: 512
+head_dim: 32
+qk_rope_head_dim: 32
 
 # --- Standard Defaults ---
 enable_dropout: false
@@ -31,40 +33,38 @@ normalization_layer_epsilon: 1.0e-6
 decoder_block: "deepseek4"
 mhc_expansion_rate: 4
 first_num_hash_layers: 3
-indexer_head_dim: 128
-indexer_n_heads: 64
-indexer_topk: 512
+indexer_head_dim: 32
+indexer_n_heads: 4
+indexer_topk: 16
 
 # Note: Layers (0, 1, 2) are prefix layers as `first_num_hash_layers=3`.
-# The 6th layer (MTP module with compress_ratio=0) has been explicitly dropped for now.
-# This leaves exactly 7 layers: 3 prefix [0,0,4] + 4 scanned.
+# The 44th layer (MTP module with compress_ratio=0) has been explicitly dropped for now.
+# This leaves exactly 43 layers: 3 prefix [0,0,4] + 40 scanned.
 # `compress_ratio=0` uses sliding window attention. In this case, layer (0, 1).
-# This is a tiny version of deepseek4 with fewer layers and less experts for debugging.
-compress_ratios: [0, 0, 4, 128, 4, 128, 4]
+compress_ratios: [0, 0, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4, 128, 4]
 
 # --- MoE configuration ---
 mlp_activations: ["silu", "linear"]
-num_experts: 8
-num_experts_per_tok: 3
+num_experts: 16
+num_experts_per_tok: 4
 mlp_activations_limit: 10
 shared_experts: 1
 routed_score_func: "sqrtsoftplus"
 routed_bias: true
-routed_scaling_factor: 1.5
-
+routed_bias_update_rate: 0.001
+load_balance_loss_weight: 0.0001
+log_moe_bias_norms: false
 
 # --- Attention configuration ---
 attention_type: 'compressed'
-attention: 'dot_product'
-q_lora_rank: 1024
-o_groups: 8
-o_lora_rank: 1024
-sliding_window_size: 128
+q_lora_rank: 16
+o_groups: 4
+o_lora_rank: 16
+sliding_window_size: 32
 
 # --- RoPE ---
+
 rope_type: "default"
 rope_max_timescale: 10000                  # Main RoPE theta
 compressed_rope_max_timescale: 160000      # Compressed RoPE theta
-max_position_embeddings: 1048576
-original_max_position_embeddings: 65536
-
+max_position_embeddings: 4096

src/maxtext/configs/types.py

@@ -1015,6 +1015,7 @@ class DeepSeekMoE(BaseModel):
   routed_score_func: str = Field("", description="Scoring function for routing (e.g., 'softmax', 'sigmoid').")
   routed_bias: bool = Field(False, description="Whether to add a bias term for routing.")
   routed_bias_update_rate: float = Field(0.0, description="Update rate applied to the router bias term.")
+  log_moe_bias_norms: bool = Field(False, description="Whether to log the norms of MoE router biases.")
   mlp_bias: bool = Field(
       False,
       description="Whether to add a learnable bias for MLP matmul, "
@@ -3531,8 +3532,17 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
           )
       if self.decoder_block == DecoderBlockType.GPT_OSS and not self.sparse_matmul and self.capacity_factor != -1:
         raise ValueError("GPT-OSS MoE only supports dropless (capacity_factor=-1) with dense matmul.")
-      if self.routed_bias and self.routed_bias_update_rate > 0.0 and self.decoder_block != DecoderBlockType.DEEPSEEK:
+      if (
+          self.routed_bias
+          and self.routed_bias_update_rate > 0.0
+          and self.decoder_block not in (DecoderBlockType.DEEPSEEK, DecoderBlockType.DEEPSEEK4)
+      ):
         raise ValueError("Loss-free load balancing is only supported for the DeepSeek decoder block.")
+      if not self.pure_nnx and self.routed_bias and self.decoder_block == DecoderBlockType.DEEPSEEK4:
+        raise ValueError(
+            "Auxiliary-loss-free routed bias for DeepSeek V4 is only supported in pure NNX mode. "
+            "Please set pure_nnx=True or disable routed_bias."
+        )
       if self.model_name.startswith("deepseek4") and self.first_num_hash_layers > 0 and self.use_ring_of_experts:
         raise ValueError("DeepSeek V4 hash routing is currently not supported with ring of experts.")
       self.validate_ragged_buffer_factor()

src/maxtext/layers/moe.py

@@ -232,6 +232,10 @@ class Tid2EidVar(nnx.Variable):
   """Custom variable to hold tid2eid without trainable param overhead."""
 
 
+class MoEBiasVar(nnx.Variable):
+  """Custom NNX Variable for Auxiliary-Loss-Free MoE Routing Bias (DSV4)."""
+
+
 class GateLogit(nnx.Module):
   """A layer used to compute gate logits, allowing to return the pre bias values for DeepSeek routing."""
 
@@ -307,10 +311,17 @@ def __init__(
     if self.use_bias:
       bias_axes = self.kernel_axes[-len(self.out_features_shape) :]
       bias_shape = kernel_shape[-len(self.out_features_shape) :]
+      # DSV3 was using nnx.Param and that code we are keeping the same
       self.bias = nnx.Param(
           default_bias_init(rngs.params(), bias_shape, self.weight_dtype),
           out_sharding=bias_axes,
       )
+      if self.model_name.startswith("deepseek4"):
+        # DSV4 uses MoEBiasVar to naturally isolate from sequence-wise updates
+        self.bias = MoEBiasVar(
+            default_bias_init(rngs.params(), bias_shape, self.weight_dtype),
+            out_sharding=bias_axes,
+        )
     else:
       self.bias = None
 
@@ -2573,13 +2584,29 @@ def generate_masks(self, top_k_indices, softmax_probs):
 
   # See Switch Transformer (https://arxiv.org/abs/2101.03961) for more details.
   def load_balance_loss(self, top_k_indices, logits) -> jax.Array:
-    """Compute the load balance loss."""
+    """Compute the sequence-wise load balance loss.
+
+    For DeepSeek V4 like models, standard load balancing across an entire batch can
+    be inadequate due to heterogeneous prompt lengths and varying sequence
+    characteristics. This method implements sequence-wise load balancing by
+    computing the token density and routing probabilities on a per-sequence basis.
+
+    The resulting loss is scaled by `self.config.load_balance_loss_weight`.
+    When this configuration value is set > 0, the computed loss is aggregated
+    into the total training loss. By minimizing this scaled auxiliary loss,
+    the optimizer updates the routing parameters to actively enforce an even
+    distribution of tokens to experts within each individual sequence.
+    """
     expert_mask = jax.nn.one_hot(top_k_indices, num_classes=self.num_experts, dtype=jnp.int32)
     summed_expert_mask = jnp.sum(expert_mask, axis=2)
     # Get fraction of tokens dispatched to each expert
+    # jnp.mean over axis=1 (sequence length) isolates the token density per sequence.
     density = jnp.mean(summed_expert_mask, axis=1)
     # get fraction of probability allocated to each expert
+    # jnp.mean over axis=1 isolates the routing probability per sequence.
     density_prob = jnp.mean(logits, axis=1)
+    # The sequence-wise densities and probabilities are multiplied and then averaged
+    # over the batch dimension, scaled by the required constant.
     loss = jnp.mean(density * density_prob) * (self.num_experts**2) * self.config.load_balance_loss_weight
     return loss
 

src/maxtext/layers/nnx_decoders.py

@@ -925,6 +925,7 @@ def _apply_layers_sequentially(
       kv_caches_stacked=None,
       skip_block_remat: bool = False,
       unroll: int = 1,
+      metadata_axis_name: str = "layers",
       **kwargs,
   ):
     """Runs the layer stack using nnx.scan.
@@ -947,6 +948,10 @@ def _apply_layers_sequentially(
         e.g. per-layer) remat internally, to avoid double rematerialization.
       unroll: Number of scan iterations to unroll into straight-line code
         (forwarded to jax.lax.scan). unroll >= length fully unrolls the loop.
+      metadata_axis_name: The name of the scan axis used during layer initialization.
+        This must perfectly match the string passed to `_create_scanned_layers`
+        (e.g., "layers", "scanned_blocks") to prevent strict JAX `pjit` PyTree
+        metadata mismatch errors when using custom `nnx.Variable` types (like `MoEBiasVar`).
       **kwargs: Keyword args forwarded to the layer (filtered by the layer signature).
 
     Returns:
@@ -1074,7 +1079,7 @@ def layer_fn(carry, scanned_vars):
 
       # Move the scan axis to each variable's param_scan_axis and restore its name
       # in the sharding metadata. jax.lax.scan emits it at position 0.
-      scanned_state = maxtext_utils_nnx.nnx_add_and_sync_scan_axis(scanned_state, "layers")
+      scanned_state = maxtext_utils_nnx.nnx_add_and_sync_scan_axis(scanned_state, metadata_axis_name)
 
       returned_kv_stacked = None
 
@@ -2053,6 +2058,7 @@ def _apply_deepseek4_scanned_blocks(
           deterministic,
           model_mode,
           length=num_full_blocks,
+          metadata_axis_name="scanned_blocks",
           **layer_call_kwargs,
       )
 

src/maxtext/trainers/pre_train/train.py

@@ -35,7 +35,7 @@
 import jax.numpy as jnp
 from jax.sharding import NamedSharding
 
-from flax import linen as nn, nnx
+from flax import linen as nn, nnx, traverse_util
 from flax.linen import partitioning as nn_partitioning
 from flax.nnx import variablelib
 
@@ -467,6 +467,7 @@ def diff_wrapper(curr_params, custom_params, rest, config, data):
   moe_bias_updates = aux.get("moe_bias_updates")
   mtp_loss = aux.get("mtp_loss", 0.0)
   new_opt_state = None
+  bias_metrics = {}
 
   if isinstance(model, nn.Module):
     if config.gradient_clipping_threshold > 0:
@@ -556,9 +557,37 @@ def move(path, value):
     new_state = state
 
     # Apply updates for Auxiliary-Loss-Free load balancing for DeepSeek family
-    if config.routed_bias and config.routed_bias_update_rate > 0.0 and moe_bias_updates is not None:
-      target_bias = new_state.model.decoder.moe_layers.DeepSeekMoeBlock_0.MoeBlock_0.gate.bias
-      target_bias.value = target_bias.value + jnp.array(moe_bias_updates[0]).transpose()
+    if config.routed_bias and config.routed_bias_update_rate > 0.0:
+      if config.model_name.startswith("deepseek4"):
+        max_logging.log("DeepSeek V4: Applying auxiliary-loss-free routing bias via pure NNX MoEBiasVar.")
+        flat_intermediates = traverse_util.flatten_dict(aux.get("intermediate_outputs", {}))
+        for path, update in flat_intermediates.items():
+          if path[-1] != "moe_bias_updates":
+            continue
+          target = new_state.model
+          prefix = path[1:-1] if path[0] == "intermediates" else path[:-1]
+          for key in prefix:
+            if hasattr(target, key):
+              target = getattr(target, key)
+            elif isinstance(target, dict) and key in target:
+              target = target[key]
+            else:
+              target = None
+              break
+          if target is None:
+            continue
+          for _, node in nnx.iter_graph(target):
+            if type(node).__name__ == "GateLogit" and hasattr(node, "bias") and node.bias is not None:
+              update_val = update[0] if isinstance(update, (tuple, list)) else update
+              name_prefix = "-".join(map(str, prefix))
+              if getattr(config, "log_moe_bias_norms", False):
+                bias_metrics[f"learning/moe_bias_before_norm_{name_prefix}"] = jnp.linalg.norm(node.bias.value)
+              node.bias.value = node.bias.value + jnp.array(update_val)
+              if getattr(config, "log_moe_bias_norms", False):
+                bias_metrics[f"learning/moe_bias_update_norm_{name_prefix}"] = jnp.linalg.norm(jnp.array(update_val))
+      elif moe_bias_updates is not None:
+        target_bias = new_state.model.decoder.moe_layers.DeepSeekMoeBlock_0.MoeBlock_0.gate.bias
+        target_bias.value = target_bias.value + jnp.array(moe_bias_updates[0]).transpose()
 
   lm_loss = xent_sum / (total_weights + EPS)
   scalar_metrics = {
@@ -571,6 +600,7 @@ def move(path, value):
       "learning/mtp_loss": mtp_loss,
       "learning/total_weights": total_weights,
   }
+  scalar_metrics.update(bias_metrics)
   if config.use_qk_clip:
     if isinstance(model, nn.Module):
       new_state = qk_clip_utils.apply_qk_clip(new_state, intermediate_outputs, config)

tests/unit/attention_compressed_test.py

@@ -37,6 +37,7 @@ def setUp(self):
             "qk_rope_head_dim=16",
             "v_head_dim=16",
             "qk_nope_head_dim=16",
+            "override_model_config=True",
         ]
     )
     self.mesh = Mesh(jax.devices(), ("data",))

tests/unit/nnx_decoders_test.py

@@ -1259,3 +1259,65 @@ def mock_donor_idx(lyr, layer_types, num_kv_shared):
               model_mode=MODEL_MODE_TRAIN,
               kv_caches=kv_caches,
           )
+
+
+class TestApplyLayersSequentiallyMetadataAxisName(unittest.TestCase):
+
+  def test_metadata_axis_name_parameterization(self):
+    from maxtext.layers.nnx_decoders import NNXDecoder
+    from maxtext.utils import maxtext_utils_nnx
+    import jax
+    from flax import nnx
+    from unittest.mock import MagicMock
+
+    cfg = _make_config()
+    cfg.param_scan_axis = 0
+    mesh = _make_mesh(cfg)
+    rngs = nnx.Rngs(params=0)
+
+    decoder = NNXDecoder(
+        config=cfg,
+        mesh=mesh,
+        model_mode=MODEL_MODE_TRAIN,
+        rngs=rngs,
+    )
+
+    class DummyLayer(nnx.Module):
+
+      def __init__(self, rngs):
+        self.p = nnx.Param(jax.numpy.zeros((2, 2)))
+
+      def __call__(self, x, **kwargs):
+        return x + self.p.value, None
+
+    # Manually create a stacked layer using NNX scan
+    stacked_layers = nnx.vmap(lambda: DummyLayer(rngs=rngs), in_axes=(), out_axes=0, axis_size=2)()
+
+    x_in = jax.numpy.zeros((2,))
+
+    # We mock maxtext_utils_nnx.nnx_add_scan_axis to ensure the custom name is passed
+    original_add_scan_axis = maxtext_utils_nnx.nnx_add_scan_axis
+    mock_add_scan_axis = MagicMock(side_effect=original_add_scan_axis)
+    maxtext_utils_nnx.nnx_add_scan_axis = mock_add_scan_axis
+
+    try:
+      # Use a custom metadata_axis_name
+      custom_axis_name = "custom_scanned_blocks"
+      out, layers, _ = decoder._apply_layers_sequentially(
+          layers=stacked_layers, x_in=x_in, length=2, metadata_axis_name=custom_axis_name
+      )
+
+      # Verify that the custom axis name was indeed passed down
+      found_custom_name = False
+      for call_args in mock_add_scan_axis.call_args_list:
+        if call_args[0][1] == custom_axis_name:
+          found_custom_name = True
+          break
+
+      self.assertTrue(found_custom_name, "The custom metadata_axis_name was not passed to nnx_add_scan_axis!")
+    finally:
+      maxtext_utils_nnx.nnx_add_scan_axis = original_add_scan_axis
+
+
+if __name__ == "__main__":
+  unittest.main()

tests/unit/train_compile_test.py

@@ -852,6 +852,8 @@ def test_deepseek4(self, scan_layers, enable_nnx):
             f"enable_nnx={enable_nnx}",
             f"pure_nnx={enable_nnx}",
             f"pure_nnx_decoder={enable_nnx}",
+            "routed_bias=False",
+            "override_model_config=True",
         )
     )