add muon support for deepseek v4

Author: snehalv2002

src/maxtext/configs/base.yml

@@ -893,10 +893,11 @@ 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).
+ds4_ns: false # Whether to use the DeepSeek-V4 hybrid Newton-Schulz iteration parameters (10 steps with custom coefficients).
 
 # Stack trace parameters
 collect_stack_trace: False

src/maxtext/configs/types.py

@@ -1500,6 +1500,10 @@ class Muon(BaseModel):
       None,
       description="If None, apply width scaling to updates. If float, apply consistent rms scaling (recommend 0.2).",
   )
+  ds4_ns: bool = Field(
+      False,
+      description="Whether to use the DeepSeek-V4 hybrid Newton-Schulz iteration parameters (10 steps with custom coefficients).",
+  )
 
 
 class PositionalEmbedding(BaseModel):
@@ -3033,6 +3037,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.DEEPSEEK_V4,
         DecoderBlockType.QWEN3,
         DecoderBlockType.GEMMA3,
         DecoderBlockType.LLAMA2,

src/maxtext/optimizers/optimizers.py

@@ -197,16 +197,26 @@ 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.")
+    
+    if config.ds4_ns:
+      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,10 @@ 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", "bias", "embedding", "logits_dense", "beta", "base", "sinks", "tid2eid", "inv_freq"))
+      for segment in path
+  ):
     return None
 
   # 2 Special weights
@@ -86,6 +89,9 @@ 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):
@@ -179,6 +185,7 @@ 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",
   ]
   config = pyconfig.initialize(argv)
   # Setup model

tests/unit/muon_utils_test.py

@@ -100,6 +100,20 @@ 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 (Grouped linear) 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,188 @@
 }
 
 
+# deepseek_v4 building blocks
+_DEEPSEEK_V4_MHC_ATTENTION = {
+    "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,
+}
+
+_DEEPSEEK_V4_MHC_MLP = _DEEPSEEK_V4_MHC_ATTENTION
+
+_DEEPSEEK_V4_MLP_STD = {
+    "MoeBlock_0": {
+        "gate": {
+            "e_score_correction_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": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+        "wo": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+    },
+}
+
+_DEEPSEEK_V4_MLP_PRE = {
+    "MoeBlock_0": {
+        "gate": {
+            "kernel": mdn(reduction_axis=(0,), output_axis=(-1,)),
+            "tid2eid": None,
+        },
+        "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": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+        "wo": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+    },
+}
+
+_DEEPSEEK_V4_ATTN_INDEXER = {
+    "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": None,
+            "q_b_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": None,
+    },
+    "kv_norm": {"scale": None},
+    "kv_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+    "o_a_proj": {"kernel": mdn(reduction_axis=(-2,), output_axis=(-1,))},
+    "o_b_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+    "q_a_norm": {"scale": None},
+    "q_a_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+    "q_b_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+    "sinks": None,
+}
+
+_DEEPSEEK_V4_ATTN_COMPRESSOR = {
+    "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": None,
+    },
+    "kv_norm": {"scale": None},
+    "kv_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+    "o_a_proj": {"kernel": mdn(reduction_axis=(-2,), output_axis=(-1,))},
+    "o_b_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+    "q_a_norm": {"scale": None},
+    "q_a_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+    "q_b_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+    "sinks": None,
+}
+
+_DEEPSEEK_V4_ATTN_BASIC = {
+    "kv_norm": {"scale": None},
+    "kv_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+    "o_a_proj": {"kernel": mdn(reduction_axis=(-2,), output_axis=(-1,))},
+    "o_b_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+    "q_a_norm": {"scale": None},
+    "q_a_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+    "q_b_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
+    "sinks": None,
+}
+
+_DEEPSEEK_V4_LAYER_INDEXER = {
+    "mhc_attention": _DEEPSEEK_V4_MHC_ATTENTION,
+    "mhc_mlp": _DEEPSEEK_V4_MHC_MLP,
+    "mlp": _DEEPSEEK_V4_MLP_STD,
+    "post_self_attention_layer_norm": {"scale": None},
+    "pre_self_attention_layer_norm": {"scale": None},
+    "self_attention": _DEEPSEEK_V4_ATTN_INDEXER,
+}
+
+_DEEPSEEK_V4_LAYER_COMPRESSOR = {
+    "mhc_attention": _DEEPSEEK_V4_MHC_ATTENTION,
+    "mhc_mlp": _DEEPSEEK_V4_MHC_MLP,
+    "mlp": _DEEPSEEK_V4_MLP_STD,
+    "post_self_attention_layer_norm": {"scale": None},
+    "pre_self_attention_layer_norm": {"scale": None},
+    "self_attention": _DEEPSEEK_V4_ATTN_COMPRESSOR,
+}
+
+_DEEPSEEK_V4_LAYER_BASIC_STD = {
+    "mhc_attention": _DEEPSEEK_V4_MHC_ATTENTION,
+    "mhc_mlp": _DEEPSEEK_V4_MHC_MLP,
+    "mlp": _DEEPSEEK_V4_MLP_STD,
+    "post_self_attention_layer_norm": {"scale": None},
+    "pre_self_attention_layer_norm": {"scale": None},
+    "self_attention": _DEEPSEEK_V4_ATTN_BASIC,
+}
+
+_DEEPSEEK_V4_LAYER_BASIC_PRE = {
+    "mhc_attention": _DEEPSEEK_V4_MHC_ATTENTION,
+    "mhc_mlp": _DEEPSEEK_V4_MHC_MLP,
+    "mlp": _DEEPSEEK_V4_MLP_PRE,
+    "post_self_attention_layer_norm": {"scale": None},
+    "pre_self_attention_layer_norm": {"scale": None},
+    "self_attention": _DEEPSEEK_V4_ATTN_BASIC,
+}
+
+_DEEPSEEK_V4_LAYER_INDEXER_PRE = {
+    "mhc_attention": _DEEPSEEK_V4_MHC_ATTENTION,
+    "mhc_mlp": _DEEPSEEK_V4_MHC_MLP,
+    "mlp": _DEEPSEEK_V4_MLP_PRE,
+    "post_self_attention_layer_norm": {"scale": None},
+    "pre_self_attention_layer_norm": {"scale": None},
+    "self_attention": _DEEPSEEK_V4_ATTN_INDEXER,
+}
+
+_DEEPSEEK_V4_LAYER_COMPRESSOR_PRE = {
+    "mhc_attention": _DEEPSEEK_V4_MHC_ATTENTION,
+    "mhc_mlp": _DEEPSEEK_V4_MHC_MLP,
+    "mlp": _DEEPSEEK_V4_MLP_PRE,
+    "post_self_attention_layer_norm": {"scale": None},
+    "pre_self_attention_layer_norm": {"scale": None},
+    "self_attention": _DEEPSEEK_V4_ATTN_COMPRESSOR,
+}
+
+DEEPSEEK_V4_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": {
+                "layers_0": _DEEPSEEK_V4_LAYER_INDEXER,
+                "layers_1": _DEEPSEEK_V4_LAYER_COMPRESSOR,
+            },
+            "logits_dense": {"kernel": None},
+            "post_layers": {
+                "layers_0": _DEEPSEEK_V4_LAYER_BASIC_STD,
+            },
+            "pre_layers": {
+                "layers_0": _DEEPSEEK_V4_LAYER_BASIC_PRE,
+                "layers_1": _DEEPSEEK_V4_LAYER_INDEXER_PRE,
+                "layers_2": _DEEPSEEK_V4_LAYER_COMPRESSOR_PRE,
+            },
+        },
+        "token_embedder": {"embedding": None},
+    }
+}
+
+
 class MuonDimensionTest(parameterized.TestCase):
   """Unit tests for Muon dimension number generation.
 
@@ -229,6 +411,7 @@ class MuonDimensionTest(parameterized.TestCase):
   @parameterized.named_parameters(
       ("deepseek2-16b", "deepseek2-16b", DEEPSEEK2_DIMENSION_NUMBER),
       ("deepseek3-671b", "deepseek3-671b", DEEPSEEK3_DIMENSION_NUMBER),
+      ("deepseek_v4-tiny", "deepseek_v4-tiny", DEEPSEEK_V4_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 +427,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 +807,38 @@ def __init__(self, rngs: nnx.Rngs):
     # Check attention out: [0, -2] -> [-1]
     self.assertEqual(result.self_attention.out.kernel.value, mdn((0, -2), (-1,)))
 
+  def test_muon_ds4_ns_config(self):
+    """Verifies that muon optimizer configures Newton-Schulz parameters correctly based on ds4_ns."""
+    model = MagicMock()
+    learning_rate_schedule = MagicMock()
+
+    # Case 1: ds4_ns = True
+    argv_true = ["", get_test_config_path(), "run_name=test", "opt_type=muon", "ds4_ns=true"]
+    config_true = pyconfig.initialize(argv_true)
+
+    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_true, 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: ds4_ns = False (Default)
+    argv_false = ["", get_test_config_path(), "run_name=test", "opt_type=muon", "ds4_ns=false"]
+    config_false = pyconfig.initialize(argv_false)
+
+    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_false, 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()