Skip to content

Commit 307ebb3

Browse files
snehalv2002snehalv2002
committed
feat: implement Muon optimizer support for DeepSeek V4 architecture
1 parent dfd8d29 commit 307ebb3

6 files changed

Lines changed: 289 additions & 5 deletions

File tree

‎src/maxtext/configs/base.yml‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -936,7 +936,7 @@ mu_dtype: "" # data type to store "mu" of AdamW tracking the first moment. Inher
936936
# Muon optimizer parameters
937937
# https://github.com/google-deepmind/optax/blob/main/optax/contrib/_muon.py
938938
# "mu_dtype", "adam_eps" are shared by AdamW
939-
# "nesterov", "ns_coeffs", "ns_steps", "weight_decay_mask", "adaptive" use default
939+
# "nesterov", "weight_decay_mask", "adaptive" use default
940940
muon_beta: 0.95 # Decay rate for the exponentially weighted average of grads.
941941
muon_weight_decay: 0 # Strength of the weight decay regularization. This is multiplied with the learning rate.
942942
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‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3535,6 +3535,7 @@ def calculate_global_batch_sizes(per_device_batch_size, expansion_factor, num_de
35353535

35363536
if self.opt_type == "muon" and self.decoder_block not in [
35373537
DecoderBlockType.DEEPSEEK,
3538+
DecoderBlockType.DEEPSEEK4,
35383539
DecoderBlockType.QWEN3,
35393540
DecoderBlockType.GEMMA3,
35403541
DecoderBlockType.LLAMA2,

‎src/maxtext/optimizers/optimizers.py‎

Lines changed: 12 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121

2222
import optax
2323
from optax.contrib._muon import muon
24+
from maxtext.configs.types import DecoderBlockType
2425
from maxtext.utils.muon_utils import get_muon_weight_dimension_numbers
2526

2627

@@ -197,16 +198,26 @@ def get_optimizer(config, learning_rate_schedule, model=None):
197198
muon_weight_dimension_numbers = get_muon_weight_dimension_numbers(model, config)
198199
else:
199200
raise ValueError("Please specify model to extract muon dimension number.")
201+
202+
if config.decoder_block == DecoderBlockType.DEEPSEEK4:
203+
ns_coeffs = [(3.4445, -4.7750, 2.0315)] * 8 + [(2.0, -1.5, 0.5)] * 2
204+
ns_steps = 10
205+
else:
206+
ns_coeffs = (3.4445, -4.7750, 2.0315)
207+
ns_steps = 5
208+
200209
muon_kwargs = {
201210
# Shared parameters: "nesterov" uses default
202211
"learning_rate": learning_rate_schedule,
203212
"eps": config.adam_eps,
204213
"mu_dtype": config.mu_dtype,
205-
# Muon-specific parameters: "ns_coeffs", "ns_steps", "weight_decay_mask", "adaptive" uses default
214+
# Muon-specific parameters: "weight_decay_mask", "adaptive" uses default
206215
"beta": config.muon_beta,
207216
"weight_decay": config.muon_weight_decay,
208217
"muon_weight_dimension_numbers": muon_weight_dimension_numbers,
209218
"consistent_rms": config.muon_consistent_rms,
219+
"ns_coeffs": ns_coeffs,
220+
"ns_steps": ns_steps,
210221
# AdamW-specific parameters
211222
"adam_b1": config.adam_b1,
212223
"adam_b2": config.adam_b2,

‎src/maxtext/utils/muon_utils.py‎

Lines changed: 38 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,32 @@ def transform_logic(path: Tuple[str, ...]) -> Optional[mdn]:
7070
"""
7171

7272
# 1 Exclude parameters not suitable for Muon (scalar, embeddings, unembedding)
73-
if _is_path_contain_any(("scale", "bias", "embedding", "logits_dense"), path):
73+
# "embedding": embedding
74+
# "logits_dense": output embedding
75+
# "tid2eid": lookup table in hash routing moe
76+
# "scale": scalar, common module
77+
# "sinks": scalar, attention sink
78+
# "bias": scalar, common module
79+
# "hc_base": scalar, in mhc head
80+
# "post_beta", "pre_beta", "res_beta": scalar, in mhc
81+
if any(
82+
any(
83+
x in segment
84+
for x in (
85+
"scale",
86+
"embedding",
87+
"logits_dense",
88+
"post_beta",
89+
"pre_beta",
90+
"res_beta",
91+
"hc_base",
92+
"sinks",
93+
"tid2eid",
94+
"bias",
95+
)
96+
)
97+
for segment in path
98+
):
7499
return None
75100

76101
# 2 Special weights
@@ -86,9 +111,12 @@ def transform_logic(path: Tuple[str, ...]) -> Optional[mdn]:
86111
# Attention output projection: [0, L, -2, -1]
87112
if "out" in path:
88113
return mdn((0, -2), (-1,))
114+
# Block-diagonal grouped linear layer: [n_groups, L, in_features_per_group, out_features_per_group]
115+
elif "o_a_proj" in path:
116+
return mdn((-2,), (-1,))
89117
# Attention qkv projection: [0, L, -2, -1]
90118
# MLA, exclude wq_a / wkv_a
91-
elif _is_path_contain_any(("query", "key", "value", "wq_b", "wkv_b"), path):
119+
elif _is_path_contain_any(("query", "key", "value", "wq_b", "wkv_b", "wkv"), path):
92120
return mdn((0,), (-2, -1))
93121

94122
# 3 Standard weights, [0, L, -1]
@@ -182,7 +210,15 @@ def get_model_mdn(model_name, scan_layers=True, verbose=False, pure_nnx=False):
182210
f"scan_layers={scan_layers}",
183211
"attention=dot_product",
184212
f"pure_nnx={pure_nnx}",
213+
"skip_jax_distributed_system=True",
185214
]
215+
if not pure_nnx:
216+
argv.extend(
217+
[
218+
"enable_nnx=False",
219+
"pure_nnx_decoder=False",
220+
]
221+
)
186222
config = pyconfig.initialize(argv)
187223
# Setup model
188224
devices_array = maxtext_utils.create_device_mesh(config)

‎tests/unit/muon_utils_test.py‎

Lines changed: 19 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -100,6 +100,25 @@ def test_self_attention_mla_wq_a_is_excluded_from_special(self):
100100
def test_standard_weight(self):
101101
self.assertEqual(muon_utils.transform_logic(("decoder", "mlp", "kernel")), mdn((0,), (-1,)))
102102

103+
# --- 4. DeepSeek V4 Specific ---
104+
def test_deepseek_v4_exclusions(self):
105+
# Scale, beta, base, sinks, and tid2eid nested segments should be excluded (None)
106+
self.assertIsNone(muon_utils.transform_logic(("decoder", "hc_head", "hc_scale")))
107+
self.assertIsNone(muon_utils.transform_logic(("decoder", "hc_head", "hc_base")))
108+
self.assertIsNone(muon_utils.transform_logic(("decoder", "layers", "layers_0", "mhc_attention", "res_beta")))
109+
self.assertIsNone(muon_utils.transform_logic(("decoder", "layers", "layers_0", "self_attention", "sinks")))
110+
self.assertIsNone(
111+
muon_utils.transform_logic(("decoder", "layers", "layers_0", "mlp", "MoeBlock_0", "gate", "tid2eid"))
112+
)
113+
self.assertIsNone(
114+
muon_utils.transform_logic(("decoder", "layers", "layers_0", "self_attention", "rotary_embedding", "inv_freq"))
115+
)
116+
117+
def test_deepseek_v4_self_attention_grouped_projection(self):
118+
# o_a_proj projects with reduction on in_features_per_group (-2)
119+
# and output on out_features_per_group (-1)
120+
self.assertEqual(muon_utils.transform_logic(("decoder", "self_attention", "o_a_proj")), mdn((-2,), (-1,)))
121+
103122

104123
class TestGetTransformTree(unittest.TestCase):
105124
"""Tests for get_transform_tree: recursive dict walk that applies transform_logic."""

‎tests/unit/optimizers_test.py‎

Lines changed: 218 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -219,6 +219,176 @@
219219
}
220220

221221

222+
# deepseek4 building blocks
223+
_DEEPSEEK4_MHC_ATTENTION = {
224+
"mhc_norm": {"scale": None},
225+
"post_alpha": mdn(reduction_axis=(0,), output_axis=(-1,)),
226+
"post_alpha_scale": None,
227+
"post_beta": None,
228+
"pre_alpha": mdn(reduction_axis=(0,), output_axis=(-1,)),
229+
"pre_alpha_scale": None,
230+
"pre_beta": None,
231+
"res_alpha": mdn(reduction_axis=(0,), output_axis=(-1,)),
232+
"res_alpha_scale": None,
233+
"res_beta": None,
234+
}
235+
236+
_DEEPSEEK4_MHC_MLP = {
237+
"mhc_norm": {"scale": None},
238+
"post_alpha": mdn(reduction_axis=(0,), output_axis=(-1,)),
239+
"post_alpha_scale": None,
240+
"post_beta": None,
241+
"pre_alpha": mdn(reduction_axis=(0,), output_axis=(-1,)),
242+
"pre_alpha_scale": None,
243+
"pre_beta": None,
244+
"res_alpha": mdn(reduction_axis=(0,), output_axis=(-1,)),
245+
"res_alpha_scale": None,
246+
"res_beta": None,
247+
}
248+
249+
_DEEPSEEK4_MLP = {
250+
"MoeBlock_0": {
251+
"gate": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
252+
"wi_0": mdn(reduction_axis=(-2,), output_axis=(-1,)),
253+
"wi_1": mdn(reduction_axis=(-2,), output_axis=(-1,)),
254+
"wo": mdn(reduction_axis=(-2,), output_axis=(-1,)),
255+
},
256+
"shared_experts": {
257+
"wi_0": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
258+
"wi_1": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
259+
"wo": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
260+
},
261+
}
262+
263+
_DEEPSEEK4_MLP_SCANNED = {
264+
"MoeBlock_0": {
265+
"gate": {
266+
"bias": None,
267+
"kernel": mdn(reduction_axis=(0,), output_axis=(-1,)),
268+
},
269+
"wi_0": mdn(reduction_axis=(-2,), output_axis=(-1,)),
270+
"wi_1": mdn(reduction_axis=(-2,), output_axis=(-1,)),
271+
"wo": mdn(reduction_axis=(-2,), output_axis=(-1,)),
272+
},
273+
"shared_experts": {
274+
"wi_0": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
275+
"wi_1": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
276+
"wo": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
277+
},
278+
}
279+
280+
_DEEPSEEK4_ATTN_BASIC = {
281+
"kv_norm": {"scale": None},
282+
"o_a_proj": {"kernel": mdn(reduction_axis=(-2,), output_axis=(-1,))},
283+
"o_b_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
284+
"q_norm": {"scale": None},
285+
"sinks": None,
286+
"wkv": {"kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1))},
287+
"wq_a": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
288+
"wq_b": {"kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1))},
289+
}
290+
291+
_DEEPSEEK4_ATTN_CSA = {
292+
"csa_compressor": {
293+
"gate_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
294+
"indexer": {
295+
"gate_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
296+
"kv_norm": {"scale": None},
297+
"kv_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
298+
"position_bias": mdn(reduction_axis=(0,), output_axis=(-1,)),
299+
"q_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
300+
"weights_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
301+
},
302+
"kv_norm": {"scale": None},
303+
"kv_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
304+
"position_bias": mdn(reduction_axis=(0,), output_axis=(-1,)),
305+
},
306+
"kv_norm": {"scale": None},
307+
"o_a_proj": {"kernel": mdn(reduction_axis=(-2,), output_axis=(-1,))},
308+
"o_b_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
309+
"q_norm": {"scale": None},
310+
"sinks": None,
311+
"wkv": {"kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1))},
312+
"wq_a": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
313+
"wq_b": {"kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1))},
314+
}
315+
316+
_DEEPSEEK4_ATTN_HCA = {
317+
"hca_compressor": {
318+
"gate_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
319+
"kv_norm": {"scale": None},
320+
"kv_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
321+
"position_bias": mdn(reduction_axis=(0,), output_axis=(-1,)),
322+
},
323+
"kv_norm": {"scale": None},
324+
"o_a_proj": {"kernel": mdn(reduction_axis=(-2,), output_axis=(-1,))},
325+
"o_b_proj": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
326+
"q_norm": {"scale": None},
327+
"sinks": None,
328+
"wkv": {"kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1))},
329+
"wq_a": {"kernel": mdn(reduction_axis=(0,), output_axis=(-1,))},
330+
"wq_b": {"kernel": mdn(reduction_axis=(0,), output_axis=(-2, -1))},
331+
}
332+
333+
_DEEPSEEK4_LAYER_BASIC = {
334+
"mhc_attention": _DEEPSEEK4_MHC_ATTENTION,
335+
"mhc_mlp": _DEEPSEEK4_MHC_MLP,
336+
"mlp": _DEEPSEEK4_MLP,
337+
"post_self_attention_layer_norm": {"scale": None},
338+
"pre_self_attention_layer_norm": {"scale": None},
339+
"self_attention": _DEEPSEEK4_ATTN_BASIC,
340+
}
341+
342+
_DEEPSEEK4_LAYER_CSA_PREFIX = {
343+
"mhc_attention": _DEEPSEEK4_MHC_ATTENTION,
344+
"mhc_mlp": _DEEPSEEK4_MHC_MLP,
345+
"mlp": _DEEPSEEK4_MLP,
346+
"post_self_attention_layer_norm": {"scale": None},
347+
"pre_self_attention_layer_norm": {"scale": None},
348+
"self_attention": _DEEPSEEK4_ATTN_CSA,
349+
}
350+
351+
_DEEPSEEK4_LAYER_CSA_SCANNED = {
352+
"mhc_attention": _DEEPSEEK4_MHC_ATTENTION,
353+
"mhc_mlp": _DEEPSEEK4_MHC_MLP,
354+
"mlp": _DEEPSEEK4_MLP_SCANNED,
355+
"post_self_attention_layer_norm": {"scale": None},
356+
"pre_self_attention_layer_norm": {"scale": None},
357+
"self_attention": _DEEPSEEK4_ATTN_CSA,
358+
}
359+
360+
_DEEPSEEK4_LAYER_HCA_SCANNED = {
361+
"mhc_attention": _DEEPSEEK4_MHC_ATTENTION,
362+
"mhc_mlp": _DEEPSEEK4_MHC_MLP,
363+
"mlp": _DEEPSEEK4_MLP_SCANNED,
364+
"post_self_attention_layer_norm": {"scale": None},
365+
"pre_self_attention_layer_norm": {"scale": None},
366+
"self_attention": _DEEPSEEK4_ATTN_HCA,
367+
}
368+
369+
DEEPSEEK4_DIMENSION_NUMBER = {
370+
"params": {
371+
"decoder": {
372+
"decoder_norm": {"scale": None},
373+
"hc_head": {
374+
"hc_base": None,
375+
"hc_fn": mdn(reduction_axis=(0,), output_axis=(-1,)),
376+
"hc_scale": None,
377+
},
378+
"layers_0": _DEEPSEEK4_LAYER_BASIC,
379+
"layers_1": _DEEPSEEK4_LAYER_BASIC,
380+
"layers_2": _DEEPSEEK4_LAYER_CSA_PREFIX,
381+
"logits_dense": {"kernel": None},
382+
"scanned_blocks": {
383+
"layers_0": _DEEPSEEK4_LAYER_HCA_SCANNED,
384+
"layers_1": _DEEPSEEK4_LAYER_CSA_SCANNED,
385+
},
386+
},
387+
"token_embedder": {"embedding": None},
388+
}
389+
}
390+
391+
222392
class MuonDimensionTest(parameterized.TestCase):
223393
"""Unit tests for Muon dimension number generation.
224394
@@ -229,6 +399,7 @@ class MuonDimensionTest(parameterized.TestCase):
229399
@parameterized.named_parameters(
230400
("deepseek2-16b", "deepseek2-16b", DEEPSEEK2_DIMENSION_NUMBER),
231401
("deepseek3-671b", "deepseek3-671b", DEEPSEEK3_DIMENSION_NUMBER),
402+
("deepseek4-284b", "deepseek4-284b", DEEPSEEK4_DIMENSION_NUMBER),
232403
("kimi-k2-1t", "kimi-k2-1t", DEEPSEEK3_DIMENSION_NUMBER),
233404
("llama2-7b", "llama2-7b", LLAMA2_DIMENSION_NUMBER),
234405
("llama3-8b", "llama3-8b", LLAMA2_DIMENSION_NUMBER),
@@ -244,7 +415,10 @@ def test_model_integration(self, model_name, expected_output):
244415
Muon dimension numbers match the hardcoded reference.
245416
"""
246417
actual_output = muon_utils.get_model_mdn(model_name, scan_layers=True, pure_nnx=False)
247-
self.assertEqual(actual_output, expected_output)
418+
if "params" in expected_output and "params" in actual_output:
419+
self.assertEqual(actual_output["params"], expected_output["params"])
420+
else:
421+
self.assertEqual(actual_output, expected_output)
248422

249423

250424
class AdamWMaskTest(parameterized.TestCase):
@@ -621,6 +795,49 @@ def __init__(self, rngs: nnx.Rngs):
621795
# Check attention out: [0, -2] -> [-1]
622796
self.assertEqual(result.self_attention.out.kernel, mdn((0, -2), (-1,)))
623797

798+
def test_muon_newton_schulz_config(self):
799+
"""Verifies that muon optimizer configures Newton-Schulz parameters correctly based on model."""
800+
model = MagicMock()
801+
learning_rate_schedule = MagicMock()
802+
803+
# Case 1: DeepSeek4 Model (Auto-configures 10-step schedule)
804+
argv_ds4 = [
805+
"",
806+
get_test_config_path(),
807+
"run_name=test",
808+
"opt_type=muon",
809+
"model_name=deepseek4-284b",
810+
"attention=dot_product",
811+
]
812+
config_ds4 = pyconfig.initialize(argv_ds4)
813+
814+
with (
815+
patch("maxtext.optimizers.optimizers.get_muon_weight_dimension_numbers") as mock_get_mdn,
816+
patch("maxtext.optimizers.optimizers.muon") as mock_muon,
817+
):
818+
mock_get_mdn.return_value = {}
819+
optimizers.get_optimizer(config_ds4, learning_rate_schedule, model=model)
820+
mock_muon.assert_called_once()
821+
_, kwargs = mock_muon.call_args
822+
self.assertEqual(kwargs["ns_steps"], 10)
823+
self.assertEqual(len(kwargs["ns_coeffs"]), 10)
824+
self.assertEqual(kwargs["ns_coeffs"][-1], (2.0, -1.5, 0.5))
825+
826+
# Case 2: Standard Model (Llama2) (Defaults to 5-step schedule)
827+
argv_llama = ["", get_test_config_path(), "run_name=test", "opt_type=muon", "model_name=llama2-7b"]
828+
config_llama = pyconfig.initialize(argv_llama)
829+
830+
with (
831+
patch("maxtext.optimizers.optimizers.get_muon_weight_dimension_numbers") as mock_get_mdn,
832+
patch("maxtext.optimizers.optimizers.muon") as mock_muon,
833+
):
834+
mock_get_mdn.return_value = {}
835+
optimizers.get_optimizer(config_llama, learning_rate_schedule, model=model)
836+
mock_muon.assert_called_once()
837+
_, kwargs = mock_muon.call_args
838+
self.assertEqual(kwargs["ns_steps"], 5)
839+
self.assertEqual(kwargs["ns_coeffs"], (3.4445, -4.7750, 2.0315))
840+
624841

625842
if __name__ == "__main__":
626843
unittest.main()

0 commit comments

Comments
 (0)