Skip to content

Fix muon moe dimension numbers #4843

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
parsley9877 wants to merge 2 commits into
base: main
Choose a base branch
from
+42 −4
Open

Fix muon moe dimension numbers #4843

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
4feb29f
Exclude bias from Muon orthogonalization
parsley9877 Aug 12, 2026
d70e2e0
Support Qwen3, GPT-OSS MoE blocks in Muon
parsley9877 Aug 12, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions src/maxtext/utils/muon_utils.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -93,17 +93,17 @@ def transform_logic(path: Tuple[str, ...]) -> Optional[mdn]:
"tid2eid",
)
)
or segment == "bias"
or (segment.endswith("bias") and segment != "position_bias")
for segment in path
):
return None

# 2 Special weights
# 2.1 Special weights: MoE, [0, L, -2, -1]
# L (optional) stands for layer when scan_layers=True
if "MoeBlock_0" in path:
if _is_path_contain_any(("MoeBlock_0", "routed_experts", "moe_block", "GptOssMlp"), path):
# exclude gate
if _is_path_contain_any(("wi_0", "wi_1", "wo"), path):
if _is_path_contain_any(("wi", "wi_0", "wi_1", "wo"), path):
return mdn((-2,), (-1,))
Comment on lines +104 to 107

@github-actions github-actions Bot Aug 11, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🟢 Low: Extracting the MoE block names and expert kernel names into variables improves code readability and makes it easier to extend with other MoE architectures in the future.

Suggested change
if _is_path_contain_any(("MoeBlock_0", "routed_experts", "moe_block", "GptOssMlp"), path):
# exclude gate
if _is_path_contain_any(("wi_0", "wi_1", "wo"), path):
if _is_path_contain_any(("wi", "wi_0", "wi_1", "wo"), path):
return mdn((-2,), (-1,))
moe_blocks = ("MoeBlock_0", "routed_experts", "moe_block", "GptOssMlp")
if _is_path_contain_any(moe_blocks, path):
# exclude gate
expert_kernels = ("wi", "wi_0", "wi_1", "wo")
if _is_path_contain_any(expert_kernels, path):
return mdn((-2,), (-1,))


# 2.2 Special weights: Self attention
Expand Down
40 changes: 39 additions & 1 deletion tests/unit/muon_utils_test.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,14 @@ def test_scale_is_excluded(self):
def test_bias_is_excluded(self):
self.assertIsNone(muon_utils.transform_logic(("decoder", "dense", "bias")))

def test_bias_variant_is_excluded(self):

@RissyRan RissyRan Aug 12, 2026

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Could you help explain a little bit why those bias to be excluded? or add a comment there?

Also, for a model/config, are you consider all bias? for instance, wondering why GptOss only excludes wo_bias?

@parsley9877 parsley9877 Aug 12, 2026

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I am under the impression that all kind of biases need to be excluded. These unittests are just some examples of bias tensors.

I think all bias needs to be excluded. GPTOss is just an example to test the code.

# Bias names (e.g., in MoE with mlp_bias=True) are excluded across all expert projections.
self.assertIsNone(muon_utils.transform_logic(("decoder", "moe_block", "wi_0_bias")))
self.assertIsNone(muon_utils.transform_logic(("decoder", "GptOssMlp", "wi_0_bias")))
self.assertIsNone(muon_utils.transform_logic(("decoder", "GptOssMlp", "wi_1_bias")))
self.assertIsNone(muon_utils.transform_logic(("decoder", "GptOssMlp", "wo_bias")))
self.assertIsNone(muon_utils.transform_logic(("decoder", "mlp", "mlp_bias")))

def test_embedding_is_excluded(self):
self.assertIsNone(muon_utils.transform_logic(("token_embedder", "embedding")))

Expand All @@ -69,9 +77,27 @@ def test_moe_wi_1_uses_last_two_axes(self):
def test_moe_wo_uses_last_two_axes(self):
self.assertEqual(muon_utils.transform_logic(("decoder", "MoeBlock_0", "wo")), mdn((-2,), (-1,)))

def test_moe_prefused_wi_uses_last_two_axes(self):
self.assertEqual(muon_utils.transform_logic(("decoder", "MoeBlock_0", "wi")), mdn((-2,), (-1,)))
self.assertEqual(muon_utils.transform_logic(("decoder", "routed_experts", "wi")), mdn((-2,), (-1,)))

def test_qwen3_next_moe_routed_experts(self):
self.assertEqual(muon_utils.transform_logic(("decoder", "mlp", "routed_experts", "wi_0")), mdn((-2,), (-1,)))
self.assertEqual(muon_utils.transform_logic(("decoder", "mlp", "routed_experts", "wi_1")), mdn((-2,), (-1,)))
self.assertEqual(muon_utils.transform_logic(("decoder", "mlp", "routed_experts", "wo")), mdn((-2,), (-1,)))

def test_qwen3_moe_block(self):
self.assertEqual(muon_utils.transform_logic(("decoder", "moe_block", "wi_0")), mdn((-2,), (-1,)))
self.assertEqual(muon_utils.transform_logic(("decoder", "moe_block", "wo")), mdn((-2,), (-1,)))

def test_gpt_oss_mlp_moe(self):
self.assertEqual(muon_utils.transform_logic(("decoder", "GptOssMlp", "wi_0")), mdn((-2,), (-1,)))
self.assertEqual(muon_utils.transform_logic(("decoder", "GptOssMlp", "wo")), mdn((-2,), (-1,)))

def test_moe_gate_falls_through_to_standard(self):
# 'gate' is inside MoeBlock_0 but not one of (wi_0, wi_1, wo) → standard.
# 'gate' is inside MoeBlock_0 but not one of (wi, wi_0, wi_1, wo) → standard.
self.assertEqual(muon_utils.transform_logic(("decoder", "MoeBlock_0", "gate", "kernel")), mdn((0,), (-1,)))
self.assertEqual(muon_utils.transform_logic(("decoder", "routed_experts", "gate", "kernel")), mdn((0,), (-1,)))

# --- 2.2 Self-attention ---
def test_self_attention_out_projection(self):
Expand Down Expand Up @@ -119,6 +145,18 @@ def test_deepseek_v4_self_attention_grouped_projection(self):
# and output on out_features_per_group (-1)
self.assertEqual(muon_utils.transform_logic(("decoder", "self_attention", "o_a_proj")), mdn((-2,), (-1,)))

def test_deepseek_v4_position_bias_is_standard_weight(self):
# position_bias in DeepSeek V4 compressed attention is a 2D weight matrix (not a 1D bias vector),
# so it receives standard Muon dimension numbers mdn((0,), (-1,)).
self.assertEqual(
muon_utils.transform_logic(("decoder", "self_attention", "csa_compressor", "position_bias")),
mdn((0,), (-1,)),
)
self.assertEqual(
muon_utils.transform_logic(("decoder", "self_attention", "hca_compressor", "position_bias")),
mdn((0,), (-1,)),
)


class TestGetTransformTree(unittest.TestCase):
"""Tests for get_transform_tree: recursive dict walk that applies transform_logic."""
Expand Down
Loading