feat: implement Muon optimizer support for DeepSeek V4 architecture

Author: snehalv2002

src/maxtext/configs/base.yml

@@ -927,7 +927,7 @@ mu_dtype: "" # data type to store "mu" of AdamW tracking the first moment. Inher
 # Muon optimizer parameters
 # https://github.com/google-deepmind/optax/blob/main/optax/contrib/_muon.py
 # "mu_dtype", "adam_eps" are shared by AdamW
-# "nesterov", "ns_coeffs", "ns_steps", "weight_decay_mask", "adaptive" use default
+# "nesterov", "weight_decay_mask", "adaptive" use default
 muon_beta: 0.95 # Decay rate for the exponentially weighted average of grads.
 muon_weight_decay: 0 # Strength of the weight decay regularization. This is multiplied with the learning rate.
 muon_consistent_rms: None # If None, apply width scaling to updates. If float, apply consistent rms scaling (recommend 0.2).

src/maxtext/configs/types.py

@@ -3376,6 +3376,7 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
 
     if self.opt_type == "muon" and self.decoder_block not in [
         DecoderBlockType.DEEPSEEK,
+        DecoderBlockType.DEEPSEEK4,
         DecoderBlockType.QWEN3,
         DecoderBlockType.GEMMA3,
         DecoderBlockType.LLAMA2,

src/maxtext/optimizers/optimizers.py

@@ -197,16 +197,29 @@ def get_optimizer(config, learning_rate_schedule, model=None):
       muon_weight_dimension_numbers = get_muon_weight_dimension_numbers(model, config)
     else:
       raise ValueError("Please specify model to extract muon dimension number.")
+
+    # pylint: disable=import-outside-toplevel
+    from maxtext.configs.types import DecoderBlockType
+
+    if config.decoder_block == DecoderBlockType.DEEPSEEK4:
+      ns_coeffs = [(3.4445, -4.7750, 2.0315)] * 8 + [(2.0, -1.5, 0.5)] * 2
+      ns_steps = 10
+    else:
+      ns_coeffs = (3.4445, -4.7750, 2.0315)
+      ns_steps = 5
+
     muon_kwargs = {
         # Shared parameters: "nesterov" uses default
         "learning_rate": learning_rate_schedule,
         "eps": config.adam_eps,
         "mu_dtype": config.mu_dtype,
-        # Muon-specific parameters: "ns_coeffs", "ns_steps", "weight_decay_mask", "adaptive" uses default
+        # Muon-specific parameters: "weight_decay_mask", "adaptive" uses default
         "beta": config.muon_beta,
         "weight_decay": config.muon_weight_decay,
         "muon_weight_dimension_numbers": muon_weight_dimension_numbers,
         "consistent_rms": config.muon_consistent_rms,
+        "ns_coeffs": ns_coeffs,
+        "ns_steps": ns_steps,
         # AdamW-specific parameters
         "adam_b1": config.adam_b1,
         "adam_b2": config.adam_b2,

src/maxtext/utils/muon_utils.py

@@ -70,7 +70,11 @@ def transform_logic(path: Tuple[str, ...]) -> Optional[mdn]:
   """
 
   # 1 Exclude parameters not suitable for Muon (scalar, embeddings, unembedding)
-  if _is_path_contain_any(("scale", "bias", "embedding", "logits_dense"), path):
+  if any(
+      any(x in segment for x in ("scale", "embedding", "logits_dense", "beta", "base", "sinks", "tid2eid", "inv_freq"))
+      or segment == "bias"
+      for segment in path
+  ):
     return None
 
   # 2 Special weights
@@ -86,9 +90,12 @@ def transform_logic(path: Tuple[str, ...]) -> Optional[mdn]:
     # Attention output projection: [0, L, -2, -1]
     if "out" in path:
       return mdn((0, -2), (-1,))
+    # Block-diagonal grouped linear layer: [n_groups, L, in_features_per_group, out_features_per_group]
+    elif "o_a_proj" in path:
+      return mdn((-2,), (-1,))
     # Attention qkv projection: [0, L, -2, -1]
     # MLA, exclude wq_a / wkv_a
-    elif _is_path_contain_any(("query", "key", "value", "wq_b", "wkv_b"), path):
+    elif _is_path_contain_any(("query", "key", "value", "wq_b", "wkv_b", "wkv"), path):
       return mdn((0,), (-2, -1))
 
   # 3 Standard weights, [0, L, -1]
@@ -182,7 +189,16 @@ def get_model_mdn(model_name, scan_layers=True, verbose=False, pure_nnx=False):
       f"scan_layers={scan_layers}",
       "attention=dot_product",
       f"pure_nnx={pure_nnx}",
+      "override_model_config=True",
+      "skip_jax_distributed_system=True",
   ]
+  if not pure_nnx:
+    argv.extend(
+        [
+            "enable_nnx=False",
+            "pure_nnx_decoder=False",
+        ]
+    )
   config = pyconfig.initialize(argv)
   # Setup model
   devices_array = maxtext_utils.create_device_mesh(config)

tests/unit/muon_utils_test.py

@@ -100,6 +100,25 @@ def test_self_attention_mla_wq_a_is_excluded_from_special(self):
   def test_standard_weight(self):
     self.assertEqual(muon_utils.transform_logic(("decoder", "mlp", "kernel")), mdn((0,), (-1,)))
 
+  # --- 4. DeepSeek V4 Specific ---
+  def test_deepseek_v4_exclusions(self):
+    # Scale, beta, base, sinks, and tid2eid nested segments should be excluded (None)
+    self.assertIsNone(muon_utils.transform_logic(("decoder", "hc_head", "hc_scale")))
+    self.assertIsNone(muon_utils.transform_logic(("decoder", "hc_head", "hc_base")))
+    self.assertIsNone(muon_utils.transform_logic(("decoder", "layers", "layers_0", "mhc_attention", "res_beta")))
+    self.assertIsNone(muon_utils.transform_logic(("decoder", "layers", "layers_0", "self_attention", "sinks")))
+    self.assertIsNone(
+        muon_utils.transform_logic(("decoder", "layers", "layers_0", "mlp", "MoeBlock_0", "gate", "tid2eid"))
+    )
+    self.assertIsNone(
+        muon_utils.transform_logic(("decoder", "layers", "layers_0", "self_attention", "rotary_embedding", "inv_freq"))
+    )
+
+  def test_deepseek_v4_self_attention_grouped_projection(self):
+    # o_a_proj projects with reduction on in_features_per_group (-2)
+    # and output on out_features_per_group (-1)
+    self.assertEqual(muon_utils.transform_logic(("decoder", "self_attention", "o_a_proj")), mdn((-2,), (-1,)))
+
 
 class TestGetTransformTree(unittest.TestCase):
   """Tests for get_transform_tree: recursive dict walk that applies transform_logic."""

tests/unit/optimizers_test.py

@@ -219,6 +219,176 @@
 }
 
 
+# deepseek4 building blocks
+_DEEPSEEK4_MHC_ATTENTION = {
+    "mhc_norm": {"scale": None},
+    "post_alpha": mdn(reduction_axis=(0,), output_axis=(-1,)),
+    "post_alpha_scale": None,
+    "post_beta": None,
+    "pre_alpha": mdn(reduction_axis=(0,), output_axis=(-1,)),
+    "pre_alpha_scale": None,
+    "pre_beta": None,
+    "res_alpha": mdn(reduction_axis=(0,), output_axis=(-1,)),
+    "res_alpha_scale": None,
+    "res_beta": None,
+}
+
+_DEEPSEEK4_MHC_MLP = {
+    "mhc_norm": {"scale": None},
+    "post_alpha": mdn(reduction_axis=(0,), output_axis=(-1,)),
+    "post_alpha_scale": None,
+    "post_beta": None,
+    "pre_alpha": mdn(reduction_axis=(0,), output_axis=(-1,)),
+    "pre_alpha_scale": None,
+    "pre_beta": None,
+    "res_alpha": mdn(reduction_axis=(0,), output_axis=(-1,)),
+    "res_alpha_scale": None,
+    "res_beta": None,
+}
+
+_DEEPSEEK4_MLP = {
+    "MoeBlock_0": {
+        "gate": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+        "wi_0": mdn(reduction_axis=(-2,), output_axis=(-1,)),
+        "wi_1": mdn(reduction_axis=(-2,), output_axis=(-1,)),
+        "wo": mdn(reduction_axis=(-2,), output_axis=(-1,)),
+    },
+    "shared_experts": {
+        "wi_0": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+        "wi_1": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+        "wo": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+    },
+}
+
+_DEEPSEEK4_MLP_SCANNED = {
+    "MoeBlock_0": {
+        "gate": {
+            "bias": None,
+            "kernel": mdn(reduction_axis=(0,), output_axis=(-1,)),
+        },
+        "wi_0": mdn(reduction_axis=(-2,), output_axis=(-1,)),
+        "wi_1": mdn(reduction_axis=(-2,), output_axis=(-1,)),
+        "wo": mdn(reduction_axis=(-2,), output_axis=(-1,)),
+    },
+    "shared_experts": {
+        "wi_0": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+        "wi_1": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+        "wo": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+    },
+}
+
+_DEEPSEEK4_ATTN_BASIC = {
+    "kv_norm": {"scale": None},
+    "o_a_proj": {"kernel": mdn(reduction_axis=(-2,), output_axis=(-1,))},
+    "o_b_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+    "q_norm": {"scale": None},
+    "sinks": None,
+    "wkv": {"kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1))},
+    "wq_a": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+    "wq_b": {"kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1))},
+}
+
+_DEEPSEEK4_ATTN_CSA = {
+    "csa_compressor": {
+        "gate_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+        "indexer": {
+            "gate_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+            "kv_norm": {"scale": None},
+            "kv_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+            "position_bias": mdn(reduction_axis=(0,), output_axis=(-1,)),
+            "q_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+            "weights_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+        },
+        "kv_norm": {"scale": None},
+        "kv_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+        "position_bias": mdn(reduction_axis=(0,), output_axis=(-1,)),
+    },
+    "kv_norm": {"scale": None},
+    "o_a_proj": {"kernel": mdn(reduction_axis=(-2,), output_axis=(-1,))},
+    "o_b_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+    "q_norm": {"scale": None},
+    "sinks": None,
+    "wkv": {"kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1))},
+    "wq_a": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+    "wq_b": {"kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1))},
+}
+
+_DEEPSEEK4_ATTN_HCA = {
+    "hca_compressor": {
+        "gate_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+        "kv_norm": {"scale": None},
+        "kv_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+        "position_bias": mdn(reduction_axis=(0,), output_axis=(-1,)),
+    },
+    "kv_norm": {"scale": None},
+    "o_a_proj": {"kernel": mdn(reduction_axis=(-2,), output_axis=(-1,))},
+    "o_b_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+    "q_norm": {"scale": None},
+    "sinks": None,
+    "wkv": {"kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1))},
+    "wq_a": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+    "wq_b": {"kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1))},
+}
+
+_DEEPSEEK4_LAYER_BASIC = {
+    "mhc_attention": _DEEPSEEK4_MHC_ATTENTION,
+    "mhc_mlp": _DEEPSEEK4_MHC_MLP,
+    "mlp": _DEEPSEEK4_MLP,
+    "post_self_attention_layer_norm": {"scale": None},
+    "pre_self_attention_layer_norm": {"scale": None},
+    "self_attention": _DEEPSEEK4_ATTN_BASIC,
+}
+
+_DEEPSEEK4_LAYER_CSA_PREFIX = {
+    "mhc_attention": _DEEPSEEK4_MHC_ATTENTION,
+    "mhc_mlp": _DEEPSEEK4_MHC_MLP,
+    "mlp": _DEEPSEEK4_MLP,
+    "post_self_attention_layer_norm": {"scale": None},
+    "pre_self_attention_layer_norm": {"scale": None},
+    "self_attention": _DEEPSEEK4_ATTN_CSA,
+}
+
+_DEEPSEEK4_LAYER_CSA_SCANNED = {
+    "mhc_attention": _DEEPSEEK4_MHC_ATTENTION,
+    "mhc_mlp": _DEEPSEEK4_MHC_MLP,
+    "mlp": _DEEPSEEK4_MLP_SCANNED,
+    "post_self_attention_layer_norm": {"scale": None},
+    "pre_self_attention_layer_norm": {"scale": None},
+    "self_attention": _DEEPSEEK4_ATTN_CSA,
+}
+
+_DEEPSEEK4_LAYER_HCA_SCANNED = {
+    "mhc_attention": _DEEPSEEK4_MHC_ATTENTION,
+    "mhc_mlp": _DEEPSEEK4_MHC_MLP,
+    "mlp": _DEEPSEEK4_MLP_SCANNED,
+    "post_self_attention_layer_norm": {"scale": None},
+    "pre_self_attention_layer_norm": {"scale": None},
+    "self_attention": _DEEPSEEK4_ATTN_HCA,
+}
+
+DEEPSEEK4_DIMENSION_NUMBER = {
+    "params": {
+        "decoder": {
+            "decoder_norm": {"scale": None},
+            "hc_head": {
+                "hc_base": None,
+                "hc_fn": mdn(reduction_axis=(0,), output_axis=(-1,)),
+                "hc_scale": None,
+            },
+            "layers_0": _DEEPSEEK4_LAYER_BASIC,
+            "layers_1": _DEEPSEEK4_LAYER_BASIC,
+            "layers_2": _DEEPSEEK4_LAYER_CSA_PREFIX,
+            "logits_dense": {"kernel": None},
+            "scanned_blocks": {
+                "layers_0": _DEEPSEEK4_LAYER_HCA_SCANNED,
+                "layers_1": _DEEPSEEK4_LAYER_CSA_SCANNED,
+            },
+        },
+        "token_embedder": {"embedding": None},
+    }
+}
+
+
 class MuonDimensionTest(parameterized.TestCase):
   """Unit tests for Muon dimension number generation.
 
@@ -229,6 +399,7 @@ class MuonDimensionTest(parameterized.TestCase):
   @parameterized.named_parameters(
       ("deepseek2-16b", "deepseek2-16b", DEEPSEEK2_DIMENSION_NUMBER),
       ("deepseek3-671b", "deepseek3-671b", DEEPSEEK3_DIMENSION_NUMBER),
+      ("deepseek4-284b", "deepseek4-284b", DEEPSEEK4_DIMENSION_NUMBER),
       ("kimi-k2-1t", "kimi-k2-1t", DEEPSEEK3_DIMENSION_NUMBER),
       ("llama2-7b", "llama2-7b", LLAMA2_DIMENSION_NUMBER),
       ("llama3-8b", "llama3-8b", LLAMA2_DIMENSION_NUMBER),
@@ -244,7 +415,10 @@ def test_model_integration(self, model_name, expected_output):
     Muon dimension numbers match the hardcoded reference.
     """
     actual_output = muon_utils.get_model_mdn(model_name, scan_layers=True, pure_nnx=False)
-    self.assertEqual(actual_output, expected_output)
+    if "params" in expected_output and "params" in actual_output:
+      self.assertEqual(actual_output["params"], expected_output["params"])
+    else:
+      self.assertEqual(actual_output, expected_output)
 
 
 class AdamWMaskTest(parameterized.TestCase):
@@ -621,6 +795,49 @@ def __init__(self, rngs: nnx.Rngs):
     # Check attention out: [0, -2] -> [-1]
     self.assertEqual(result.self_attention.out.kernel, mdn((0, -2), (-1,)))
 
+  def test_muon_ds4_ns_config(self):
+    """Verifies that muon optimizer configures Newton-Schulz parameters correctly based on model."""
+    model = MagicMock()
+    learning_rate_schedule = MagicMock()
+
+    # Case 1: DeepSeek4 Model (Auto-configures 10-step schedule)
+    argv_ds4 = [
+        "",
+        get_test_config_path(),
+        "run_name=test",
+        "opt_type=muon",
+        "model_name=deepseek4-284b",
+        "attention=dot_product",
+    ]
+    config_ds4 = pyconfig.initialize(argv_ds4)
+
+    with (
+        patch("maxtext.optimizers.optimizers.get_muon_weight_dimension_numbers") as mock_get_mdn,
+        patch("maxtext.optimizers.optimizers.muon") as mock_muon,
+    ):
+      mock_get_mdn.return_value = {}
+      optimizers.get_optimizer(config_ds4, learning_rate_schedule, model=model)
+      mock_muon.assert_called_once()
+      _, kwargs = mock_muon.call_args
+      self.assertEqual(kwargs["ns_steps"], 10)
+      self.assertEqual(len(kwargs["ns_coeffs"]), 10)
+      self.assertEqual(kwargs["ns_coeffs"][-1], (2.0, -1.5, 0.5))
+
+    # Case 2: Standard Model (Llama2) (Defaults to 5-step schedule)
+    argv_llama = ["", get_test_config_path(), "run_name=test", "opt_type=muon", "model_name=llama2-7b"]
+    config_llama = pyconfig.initialize(argv_llama)
+
+    with (
+        patch("maxtext.optimizers.optimizers.get_muon_weight_dimension_numbers") as mock_get_mdn,
+        patch("maxtext.optimizers.optimizers.muon") as mock_muon,
+    ):
+      mock_get_mdn.return_value = {}
+      optimizers.get_optimizer(config_llama, learning_rate_schedule, model=model)
+      mock_muon.assert_called_once()
+      _, kwargs = mock_muon.call_args
+      self.assertEqual(kwargs["ns_steps"], 5)
+      self.assertEqual(kwargs["ns_coeffs"], (3.4445, -4.7750, 2.0315))
+
 
 if __name__ == "__main__":
   unittest.main()