Use the sequence-wise balance loss for DeepSeek-V4 routers

src/maxtext/configs/base.yml

@@ -309,7 +309,7 @@ log_step_diagnostics: false # append grad/param norms, max |grad|, routed-bias c
 moe_expert_input_dim: -1 # feature dimension of the tokens entering the MoE expert blocks.
 base_moe_mlp_dim: -1 # intermediate dimension at MoE layer. For a fully MoE model, base_mlp_dim must be equal to base_moe_mlp_dim.
 load_balance_loss_weight: 0.0 # weight for the load balance loss
-moe_use_megatron_seq_aux_loss: false # when true, sigmoid routers (routed_score_func="sigmoid") use Megatron-LM's sequence-wise aux loss summed over MoE layers; false (default) uses the Switch-style loss averaged over MoE layers.
+moe_use_megatron_seq_aux_loss: false # when true, sigmoid and sqrtsoftplus routers use Megatron-LM's sequence-wise aux loss summed over MoE layers; false (default) uses the Switch-style loss averaged over MoE layers.
 use_random_routing: false # whether to use random routing for debug/test purpose
 use_custom_sort_vjp: true # whether to use a custom VJP sort for efficient backward pass processing in sparse matmul
 router_topk_matmul_vjp: false # deepseek routing: dense one-hot VJP for the top-k weight gather (no scatter-add in the backward)

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

@@ -51,6 +51,8 @@ norm_topk_prob: true
 routed_bias: true
 routed_bias_update_rate: 0.001
 load_balance_loss_weight: 0.0001
+# DeepSeek-V3 sequence-wise balance loss (V3 Eqs 17-20) on the Sqrt(Softplus) affinities.
+moe_use_megatron_seq_aux_loss: true
 log_moe_bias_norms: false
 routed_scaling_factor: 2.5
 float32_gate_logits: true

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

@@ -54,6 +54,8 @@ norm_topk_prob: true
 routed_bias: true
 routed_bias_update_rate: 0.001
 load_balance_loss_weight: 0.0001
+# DeepSeek-V3 sequence-wise balance loss (V3 Eqs 17-20) on the Sqrt(Softplus) affinities.
+moe_use_megatron_seq_aux_loss: true
 log_moe_bias_norms: false
 routed_scaling_factor: 1.5
 

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

@@ -53,6 +53,8 @@ routed_score_func: "sqrtsoftplus"
 routed_bias: true
 routed_bias_update_rate: 0.001
 load_balance_loss_weight: 0.0001
+# DeepSeek-V3 sequence-wise balance loss (V3 Eqs 17-20) on the Sqrt(Softplus) affinities.
+moe_use_megatron_seq_aux_loss: true
 log_moe_bias_norms: false
 
 # --- Attention configuration ---

src/maxtext/configs/types.py

@@ -1276,10 +1276,11 @@ class MoEGeneral(BaseModel):
   moe_use_megatron_seq_aux_loss: bool = Field(
       False,
       description=(
-          "When True, sigmoid routers (routed_score_func='sigmoid', with te_moe_block=False) use"
-          " Megatron-LM's sequence-wise auxiliary load-balancing loss (fp32 per-token normalized sigmoid scores,"
+          "When True, sigmoid and Sqrt(Softplus) routers (routed_score_func='sigmoid' or 'sqrtsoftplus', with"
+          " te_moe_block=False) use Megatron-LM's sequence-wise auxiliary load-balancing loss (fp32 per-token"
+          " normalized affinity scores,"
           " token fractions from an unbiased, ungrouped top-k over the full sequence, summed over MoE layers)."
-          " When False (default), sigmoid routers use the Switch Transformer-style loss averaged over MoE layers."
+          " When False (default), these routers use the Switch Transformer-style loss averaged over MoE layers."
       ),
   )
   use_custom_sort_vjp: bool = Field(

src/maxtext/experimental/lineage/lineage_adapter.py

@@ -135,7 +135,7 @@ def uses_megatron_seq_aux_loss(cfg: Any) -> bool:
   """Mirrors `moe.uses_megatron_seq_aux_loss`."""
   return bool(
       getattr(cfg, "moe_use_megatron_seq_aux_loss", False)
-      and getattr(cfg, "routed_score_func", "") == "sigmoid"
+      and getattr(cfg, "routed_score_func", "") in ("sigmoid", "sqrtsoftplus")
       and not getattr(cfg, "te_moe_block", False)
   )
 

src/maxtext/layers/moe.py

@@ -454,17 +454,17 @@ def routing_tag_fn(config):
 
 
 def uses_megatron_seq_aux_loss(config) -> bool:
-  """Whether the MoE aux loss is Megatron-LM's sigmoid-router seq_aux_loss.
+  """Whether the MoE aux loss is Megatron-LM's sequence-wise seq_aux_loss (DeepSeek-V3 Eqs 17-20).
 
-  When moe_use_megatron_seq_aux_loss=True, sigmoid routers (DeepSeek-V3 family) use it:
-  fp32 per-token normalized sigmoid scores, token fractions from an unbiased, ungrouped top-k,
-  one loss per full sequence, summed over MoE layers. By default (moe_use_megatron_seq_aux_loss=False),
-  and for other routers or te_moe_block (whose aux loss TransformerEngine computes), MaxText uses
-  the per-layer Switch-style loss averaged over MoE layers.
+  When moe_use_megatron_seq_aux_loss=True, sigmoid routers (DeepSeek-V3 family) and Sqrt(Softplus)
+  routers (DeepSeek-V4) use it: fp32 per-token normalized affinity scores, token fractions from an
+  unbiased, ungrouped top-k, one loss per full sequence, summed over MoE layers. By default
+  (moe_use_megatron_seq_aux_loss=False), and for other routers or te_moe_block (whose aux loss
+  TransformerEngine computes), MaxText uses the per-layer Switch-style loss averaged over MoE layers.
   """
   return (
       getattr(config, "moe_use_megatron_seq_aux_loss", False)
-      and getattr(config, "routed_score_func", "") == "sigmoid"
+      and getattr(config, "routed_score_func", "") in ("sigmoid", "sqrtsoftplus")
       and not getattr(config, "te_moe_block", False)
   )
 
@@ -473,12 +473,13 @@ def megatron_aux_loss_inputs(scores, num_experts_per_tok):
   """Aux-loss probabilities and token assignments of the Megatron-LM sigmoid router.
 
   Mirrors Megatron-LM `compute_routing_scores_for_aux_loss(score_function="sigmoid")`: the
-  probabilities are the fp32 sigmoid scores normalized per token, and the token assignments
+  probabilities are the fp32 affinity scores (sigmoid, or Sqrt(Softplus) for DeepSeek-V4)
+  normalized per token, and the token assignments
   come from a plain top-k over the scores, i.e. without the expert bias and without the
   group-limited mask used for dispatch.
 
   Args:
-      scores: Sigmoid router scores before the expert bias, shape (batch, sequence, num_experts).
+      scores: Router affinity scores before the expert bias, shape (batch, sequence, num_experts).
       num_experts_per_tok: Router top-k.
 
   Returns:
@@ -2675,7 +2676,7 @@ def get_routed_moe_shardings(is_batch_sharded_by_expert, has_input_ids):
     # Exclude the EP axis so psum doesn't redundantly reduce and scale counts by num_ep
     # for calculate_load_balance_updates().
     roe_bias_axis_names = _filter_axis_names(bias_update_axis_names, self._expert_parallelism_name)
-    # Sigmoid routers use Megatron-LM's seq_aux_loss, computed once on the full sequences after the
+    # Sigmoid and Sqrt(Softplus) routers use Megatron-LM's seq_aux_loss, computed once on the full sequences after the
     # shard_map (below), so the per-chunk loss inside the ring-of-experts loop is skipped.
     use_megatron_seq_aux_loss = (
         self.config.load_balance_loss_weight > 0.0
@@ -4288,11 +4289,11 @@ def load_balance_loss(self, top_k_indices, logits, density_prob=None) -> jax.Arr
   def megatron_seq_aux_loss(self, scores) -> jax.Array:
     """Sequence-wise aux loss of the Megatron-LM sigmoid router (seq_aux_loss).
 
-    Used for sigmoid routers (see `uses_megatron_seq_aux_loss`). Pass whole sequences: the token
+    Used for sigmoid and Sqrt(Softplus) routers (see `uses_megatron_seq_aux_loss`). Pass whole sequences: the token
     fractions and probabilities are per-sequence statistics.
 
     Args:
-      scores: Sigmoid router scores before the expert bias, shape (batch, sequence, num_experts).
+      scores: Router affinity scores before the expert bias, shape (batch, sequence, num_experts).
 
     Returns:
       Scalar loss averaged over the batch, already scaled by load_balance_loss_weight.
@@ -4402,7 +4403,7 @@ def dense_matmul(
     # These layers route deterministically based on token IDs and do not generate auxiliary loss.
     if self.config.model_call_mode != "inference" and not self.is_hash_routing:
       if self.config.load_balance_loss_weight > 0.0 and uses_megatron_seq_aux_loss(self.config):
-        # Same loss as sparse_matmul: Megatron-LM's seq_aux_loss on the pre-bias sigmoid scores.
+        # Same loss as sparse_matmul: Megatron-LM's seq_aux_loss on the pre-bias affinity scores.
         lb_loss = self.megatron_seq_aux_loss(pre_bias_logits if pre_bias_logits is not None else gate_logits)
       else:
         softmax_probs = jax.nn.softmax(gate_logits.astype(jnp.float32), axis=-1).astype(self.dtype)

tests/unit/deepseek_v4_aux_loss_spec_test.py

@@ -0,0 +1,179 @@
+# Copyright 2023–2026 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# 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.
+
+"""Spec test of the DeepSeek-V4 sequence-wise balance loss on deepseek4-tiny.
+
+DeepSeek-V4 keeps the DeepSeek-V3 sequence-wise balance loss (V3 Eqs 17-20) with
+Sqrt(Softplus) affinities s:
+  L = alpha * sum_i f_i * P_i,
+  f_i = N / (K * T) * #{t : s_{i,t} in Topk(s_{:,t}, K)} (unbiased scores),
+  P_i = mean_t s_{i,t} / sum_j s_{j,t}.
+The reference below is float64 numpy written from the paper, independent of MaxText.
+"""
+
+import unittest
+
+from flax.linen import partitioning as nn_partitioning
+import jax
+import jax.numpy as jnp
+from jax.sharding import Mesh
+import numpy as np
+
+from maxtext.configs import pyconfig
+from maxtext.layers import moe
+from maxtext.layers.initializers import nd_dense_init
+from maxtext.utils import maxtext_utils
+from tests.utils import linen_wrappers
+from tests.utils.test_helpers import get_test_config_path
+
+
+def _sqrt_softplus(u):
+  return np.sqrt(np.logaddexp(0.0, u))
+
+
+def reference_balance_loss(u, alpha, k):
+  """Returns (loss, dloss/du) for router logits u of shape (batch, seq, experts), in float64."""
+  u = np.asarray(u, np.float64)
+  batch, seq, num_experts = u.shape
+  s = _sqrt_softplus(u)
+  topk = np.argsort(-s, axis=-1, kind="stable")[..., :k]
+  counts = np.zeros((batch, num_experts))
+  for b in range(batch):
+    counts[b] = np.bincount(topk[b].ravel(), minlength=num_experts)
+  f = num_experts / (k * seq) * counts
+  s_sum = np.sum(s, axis=-1, keepdims=True)
+  p = np.mean(s / s_sum, axis=1)
+  loss = alpha * np.mean(np.sum(f * p, axis=-1))
+  # f is an indicator count (zero gradient); the gradient flows through P only.
+  f_b = f[:, None, :]
+  dl_ds = alpha / (batch * seq) * (f_b / s_sum - np.sum(f_b * s, axis=-1, keepdims=True) / s_sum**2)
+  # d sqrt(softplus(u)) / du = sigmoid(u) / (2 s).
+  dl_du = dl_ds / (1.0 + np.exp(-u)) / (2.0 * s)
+  return loss, dl_du
+
+
+def _tiny_v4_config(run_name, **overrides):
+  """Returns a float32 deepseek4-tiny config on CPU; `overrides` replace the defaults below."""
+  kwargs = {
+      "run_name": run_name,
+      "enable_checkpointing": False,
+      "model_name": "deepseek4-tiny",
+      "attention": "dot_product",
+      "override_model_config": True,
+      "dtype": "float32",
+      "weight_dtype": "float32",
+      "matmul_precision": "highest",
+      "float32_gate_logits": True,
+      # megablox runs its Pallas kernels in interpret mode off-TPU.
+      "megablox": True,
+      "sparse_matmul": True,
+      "per_device_batch_size": 2,
+      "max_target_length": 16,
+  }
+  kwargs.update(overrides)
+  return pyconfig.initialize([None, get_test_config_path()], **kwargs)
+
+
+def _run_lb_loss(cfg):
+  """Returns (lb_loss, d lb_loss / d gate kernel, gate kernel, inputs) of one routed (non-hash) MoE block."""
+  mesh = Mesh(maxtext_utils.create_device_mesh(cfg), cfg.mesh_axes)
+  model = linen_wrappers.to_linen(
+      moe.RoutedMoE,
+      name="MoeBlock",
+      config=cfg,
+      num_experts=cfg.num_experts,
+      num_experts_per_tok=cfg.num_experts_per_tok,
+      mesh=mesh,
+      kernel_init=nd_dense_init(1.0, "fan_in", "truncated_normal"),
+      kernel_axes=("embed", "mlp"),
+      intermediate_dim=cfg.moe_mlp_dim,
+      dtype=cfg.dtype,
+  )
+  rng_model, rng_inputs = jax.random.split(jax.random.PRNGKey(11))
+  inputs = jax.random.normal(
+      rng_inputs, (cfg.per_device_batch_size * jax.device_count(), cfg.max_target_length, cfg.base_emb_dim)
+  )
+  with jax.set_mesh(mesh), nn_partitioning.axis_rules(cfg.logical_axis_rules):
+    params = model.init({"params": rng_model, "dropout": rng_model}, inputs)["params"]
+
+    def loss_fn(params):
+      (_, lb_loss, _), _ = model.apply({"params": params}, inputs, mutable=["intermediates"])
+      return lb_loss
+
+    lb_loss, grads = jax.jit(jax.value_and_grad(loss_fn))(params)
+  gate_kernel = np.asarray(params["gate"]["kernel"].value, np.float64)
+  gate_grad = np.asarray(grads["gate"]["kernel"].value, np.float64)
+  return float(lb_loss), gate_grad, gate_kernel, np.asarray(inputs, np.float64)
+
+
+class _RouterStub:
+  """The RoutedMoE attributes the aux-loss methods read."""
+
+  load_balance_loss = moe.RoutedMoE.load_balance_loss
+  megatron_seq_aux_loss = moe.RoutedMoE.megatron_seq_aux_loss
+
+  def __init__(self, cfg):
+    self.num_experts = cfg.num_experts
+    self.num_experts_per_tok = cfg.num_experts_per_tok
+    self.config = cfg
+
+
+class DeepSeekV4BalanceLossSpecTest(unittest.TestCase):
+  """deepseek4-tiny RoutedMoE balance loss vs the spec reference."""
+
+  @classmethod
+  def setUpClass(cls):
+    super().setUpClass()
+    cls.cfg = _tiny_v4_config("ds4_aux_spec")
+    cls.lb_loss, cls.gate_grad, gate_kernel, inputs = _run_lb_loss(cls.cfg)
+    cls.logits = np.einsum("bse,ex->bsx", inputs, gate_kernel)
+    cls.ref_loss, dl_du = reference_balance_loss(
+        cls.logits, cls.cfg.load_balance_loss_weight, cls.cfg.num_experts_per_tok
+    )
+    cls.ref_gate_grad = np.einsum("bse,bsx->ex", inputs, dl_du)
+
+  def test_config_matches_spec(self):
+    self.assertEqual(self.cfg.routed_score_func, "sqrtsoftplus")
+    self.assertAlmostEqual(self.cfg.load_balance_loss_weight, 1e-4)
+    self.assertTrue(moe.uses_megatron_seq_aux_loss(self.cfg))
+
+  def test_loss_value_matches_spec(self):
+    rel = abs(self.lb_loss - self.ref_loss) / abs(self.ref_loss)
+    self.assertLessEqual(rel, 1e-5, f"maxtext={self.lb_loss!r} spec={self.ref_loss!r}")
+
+  def test_router_gradient_matches_spec(self):
+    rel_l2 = np.linalg.norm(self.gate_grad - self.ref_gate_grad) / np.linalg.norm(self.ref_gate_grad)
+    cosine = np.sum(self.gate_grad * self.ref_gate_grad) / (
+        np.linalg.norm(self.gate_grad) * np.linalg.norm(self.ref_gate_grad)
+    )
+    self.assertLessEqual(rel_l2, 1e-4)
+    self.assertGreaterEqual(cosine, 0.99999)
+
+  def test_score_gradient_matches_spec(self):
+    """Gradient w.r.t. the pre-bias router logits of the loss on Sqrt(Softplus) scores."""
+    router = _RouterStub(self.cfg)
+
+    def loss(u):
+      return router.megatron_seq_aux_loss(jnp.sqrt(jax.nn.softplus(u)))
+
+    got = np.asarray(jax.grad(loss)(jnp.asarray(self.logits, jnp.float32)), np.float64)
+    _, want = reference_balance_loss(self.logits, self.cfg.load_balance_loss_weight, self.cfg.num_experts_per_tok)
+    rel_l2 = np.linalg.norm(got - want) / np.linalg.norm(want)
+    cosine = np.sum(got * want) / (np.linalg.norm(got) * np.linalg.norm(want))
+    self.assertLessEqual(rel_l2, 1e-4)
+    self.assertGreaterEqual(cosine, 0.99999)
+
+
+if __name__ == "__main__":
+  unittest.main()

tests/unit/moe_megatron_aux_loss_test.py

@@ -135,10 +135,11 @@ def config(routed_score_func, te_moe_block=False, moe_use_megatron_seq_aux_loss=
       )
 
     self.assertTrue(moe.uses_megatron_seq_aux_loss(config("sigmoid")))
+    self.assertTrue(moe.uses_megatron_seq_aux_loss(config("sqrtsoftplus")))
     self.assertFalse(moe.uses_megatron_seq_aux_loss(config("sigmoid", moe_use_megatron_seq_aux_loss=False)))
     # Missing attribute on a stub config defaults to False.
     self.assertFalse(moe.uses_megatron_seq_aux_loss(SimpleNamespace(routed_score_func="sigmoid")))
-    for score_func in ("softmax", "sqrtsoftplus", ""):
+    for score_func in ("softmax", ""):
       self.assertFalse(moe.uses_megatron_seq_aux_loss(config(score_func)), score_func)
     # TransformerEngine's MoEBlock computes its own aux loss.
     self.assertFalse(moe.uses_megatron_seq_aux_loss(config("sigmoid", te_moe_block=True)))