feat: implement Muon optimizer support for DeepSeek V4 architecture

Author: snehalv2002

src/maxtext/checkpoint_conversion/utils/param_mapping.py

@@ -3873,6 +3873,7 @@ def reshape_vision_attn_out(input_tensor, target_shape):
 
 
 def DEEPSEEKV4_MAXTEXT_TO_HF_PARAM_MAPPING(config, maxtext_config, scan_layers=False):
+  """Maps MaxText parameter keys to HuggingFace parameter keys for DeepSeek V4."""
   n_layers = config["num_hidden_layers"]
   num_experts = config.get("n_routed_experts", 8)
 
@@ -4027,6 +4028,8 @@ def get_hf_expert_keys(expert_subpath_template):
 
 
 def DEEPSEEKV4_MAXTEXT_TO_HF_PARAM_HOOK_FN(config, maxtext_config, scan_layers=False, saving_to_hf=False):
+  """Returns hook functions for transforming weights between MaxText and HuggingFace for DeepSeek V4."""
+
   def transpose(input_tensor, target_shape=None):
     return np.transpose(input_tensor)
 
@@ -4046,7 +4049,6 @@ def reshape_transpose_wq_b(input_tensor, target_shape=None):
     # HF: [n_heads * q_head_dim, kv_lora_rank]
     # MaxText: [kv_lora_rank, n_heads, q_head_dim]
     tensor = np.transpose(input_tensor)  # [kv_lora_rank, n_heads * q_head_dim]
-    n_heads = config["num_attention_heads"]
     return tensor.reshape(target_shape)
 
   def reshape_transpose_wkv(input_tensor, target_shape=None):
@@ -4132,7 +4134,7 @@ def mhc_split_scale_res(input_tensor, target_shape=None):
       mapping[key] = transpose
     elif "hc_head-hc_base" in key or "hc_head-hc_scale" in key:
       mapping[key] = identity
-    elif type(hf_key) == list:
+    elif isinstance(hf_key, list):
       mapping[key] = transpose if not saving_to_hf else transpose_stack
     elif "-kernel" in key or "-embedding" in key or "-sinks" in key:
       mapping[key] = transpose

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/deepseek_v4_vs_reference_test.py

@@ -1249,6 +1249,7 @@ def setUp(self):
     self.mesh = jax.sharding.Mesh(devices, ("tensor",))
 
   def _apply_param_mapping(self, mt_layer, pt_layer, l):
+    """Maps PT weights to MaxText layer."""
     from maxtext.checkpoint_conversion.utils.param_mapping import (
         DEEPSEEKV4_MAXTEXT_TO_HF_PARAM_MAPPING,
         DEEPSEEKV4_MAXTEXT_TO_HF_PARAM_HOOK_FN,

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."""