diff --git a/src/maxtext/configs/base.yml b/src/maxtext/configs/base.yml index d538a4bf6d..06efe1e2bf 100644 --- a/src/maxtext/configs/base.yml +++ b/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) diff --git a/src/maxtext/configs/models/deepseek4-1600b.yml b/src/maxtext/configs/models/deepseek4-1600b.yml index a92762f196..578cd828b6 100644 --- a/src/maxtext/configs/models/deepseek4-1600b.yml +++ b/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 diff --git a/src/maxtext/configs/models/deepseek4-284b.yml b/src/maxtext/configs/models/deepseek4-284b.yml index 5623878b6c..65b1e8f36b 100644 --- a/src/maxtext/configs/models/deepseek4-284b.yml +++ b/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 diff --git a/src/maxtext/configs/models/deepseek4-tiny.yml b/src/maxtext/configs/models/deepseek4-tiny.yml index da81e937a8..6bb32720bd 100644 --- a/src/maxtext/configs/models/deepseek4-tiny.yml +++ b/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 --- diff --git a/src/maxtext/configs/types.py b/src/maxtext/configs/types.py index b13f58afbd..d980bc3589 100644 --- a/src/maxtext/configs/types.py +++ b/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( diff --git a/src/maxtext/experimental/lineage/lineage_adapter.py b/src/maxtext/experimental/lineage/lineage_adapter.py index 77d0d3607c..bf52bc8447 100644 --- a/src/maxtext/experimental/lineage/lineage_adapter.py +++ b/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) ) diff --git a/src/maxtext/layers/moe.py b/src/maxtext/layers/moe.py index 01e6ce4445..f99c7c99dd 100644 --- a/src/maxtext/layers/moe.py +++ b/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,14 +4403,13 @@ 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) lb_loss = ( self.load_balance_loss(top_k_indices, softmax_probs) if self.config.load_balance_loss_weight > 0.0 else None ) - # TODO(dipakg-lang, b/521990776): Add sequence-wise balance loss * 0.0001 else: lb_loss = None diff --git a/tests/unit/deepseek_v4_aux_loss_spec_test.py b/tests/unit/deepseek_v4_aux_loss_spec_test.py new file mode 100644 index 0000000000..648a3b519e --- /dev/null +++ b/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() diff --git a/tests/unit/moe_megatron_aux_loss_test.py b/tests/unit/moe_megatron_aux_loss_test.py index 90d3528c60..1da46a5ad1 100644 --- a/tests/unit/moe_megatron_aux_loss_test.py +++ b/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)))