Fix muon moe dimension numbers

[parsley9877]

Description

This PR updates maxtext.utils.muon_utils.transform_logic to assign MuonDimensionNumbers to MoE expert weights across (Qwen 3, Qwen 3-next, and GPT-OSS) and ensures all bias variants are correctly excluded from Muon orthogonalization.

Problem

1. Missing MoE block names: transform_logic previously checked if "MoeBlock_0" in path: to detect expert weights. However, newer MoE architectures use different module names:
   • routed_experts (Qwen3-next in qwen3.py)
   • moe_block (Qwen3 MoE in qwen3.py / qwen3_custom.py)
   • GptOssMlp (GPT-OSS in gpt_oss.py)
     Because these module names did not contain "MoeBlock_0", their expert weight tensors fell through to the default mapping mdn((0,), (-1,)) instead of mdn((-2,), (-1,)), leading to incorrect orthogonalization across the expert axis.
2. Prefused MoE kernels: When prefuse_moe_weights=True, the fused gate+up kernel is named wi, which was missing from the recognized expert weight kernels ("wi_0", "wi_1", "wo").
3. Bias variants: transform_logic checked segment == "bias", which missed compound bias names such as wi_0_bias, wo_bias (used in GPT-OSS with mlp_bias=True), and mlp_bias.

Solution

• Expanded MoE block container matching in transform_logic to recognize ("MoeBlock_0", "routed_experts", "moe_block", "GptOssMlp").
• Added "wi" to the expert weight kernel tuple ("wi", "wi_0", "wi_1", "wo") returning mdn((-2,), (-1,)).
• Added "bias" to the substring exclusion tuple so all bias variations (wi_0_bias, wo_bias, etc.) return None and fall back to AdamW as intended.

FIXES: #4798

Tests

Added unit test coverage in tests/unit/muon_utils_test.py:

• test_bias_variant_is_excluded: Verifies wi_0_bias, wo_bias, and mlp_bias return None.
• test_qwen3_next_moe_routed_experts: Verifies routed_experts expert kernels (wi_0, wi_1, wo) return mdn((-2,), (-1,)).
• test_qwen3_moe_block: Verifies moe_block expert kernels return mdn((-2,), (-1,)).
• test_gpt_oss_mlp_moe: Verifies GptOssMlp expert kernels return mdn((-2,), (-1,)).
• test_moe_prefused_wi_uses_last_two_axes: Verifies prefused wi kernels return mdn((-2,), (-1,)).
• test_moe_gate_falls_through_to_standard: Verifies gate routing kernels fall through to standard mapping.

Ran test suite locally:

PYTHONPATH=src python3 -m unittest tests/unit/muon_utils_test.py

Checklist

Before submitting this PR, please make sure (put X in square brackets):

• I have performed a self-review of my code. For an optional AI review, add the gemini-review label.
• I have necessary comments in my code, particularly in hard-to-understand areas.
• I have run end-to-end tests and provided workload links above if applicable.
• I have made or will make corresponding changes to the doc if needed, including adding new documentation pages to the relevant Table of Contents (toctree directive) as explained in our documentation.

[gemini-code-assist]

Code Review

This pull request updates the Muon optimizer utilities to support additional Mixture of Experts (MoE) block variants, including 'routed_experts', 'moe_block', and 'GptOssMlp', as well as prefused 'wi' weights. It also refactors the bias exclusion logic to simplify the path check and adds comprehensive unit tests to verify these changes. There are no review comments, and I have no additional feedback to provide.

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

📢 Thoughts on this report? Let us know!

[github-actions]

🤖 Hi @RissyRan, I've received your request, and I'm working on it now! You can track my progress in the logs for more details.

[github-actions]

## 📋 Review Summary

This pull request successfully resolves issue #4798 by ensuring correct Muon dimension numbers are assigned to Mixture-of-Experts (MoE) models and that all bias variants are correctly excluded from orthogonalization. The implementation is clean, robust, and accompanied by comprehensive unit test coverage verifying all proposed edge cases.

🔍 General Feedback

• Comprehensive Testing: Excellent job adding detailed unit tests covering all target architectures (Qwen3, Qwen3-next, GPT-OSS, and prefused/gate configurations).
• Robust Substring Matching: The group exclusion checks for bias variant substrings ("bias" in segment) successfully prevents any bias variant from being wrongly subjected to Muon orthogonalization.
• Pipeline Parallelism Alignment (Future Scope): In src/maxtext/utils/pipeline_utils.py (around line 80), is_moe_block_0 = "MoeBlock_0" in path_keys is used to identify MoE blocks. Since you have introduced alternative MoE block names ("routed_experts", "moe_block", "GptOssMlp"), those names may also need to be recognized in pipeline_utils.py in a separate follow-up PR to ensure they are handled correctly during scanned pipeline parallelism execution.

[github-actions]

🟢 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,))

[RissyRan]

LGTM at high level, and trying to understand a little bit on tests.

FYI @shuningjin who may have most context of Muon conv

[RissyRan]

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?