[faresobeid]

Two fixes to how Muon handles parameters.

• Depthwise nn.Conv1d weights go to AdamW. These are the gated delta net and Mamba short convolutions and the Qwen3.8 PLE convolution, shape [C, 1, K]. Muon treats a 3D tensor as a batch of matrices, so each channel's 1xK filter was orthogonalized on its own. That only normalizes each channel's taps to the same update RMS, whatever its gradient. They are now matched by module type.
• GroupedLinear weights get one Newton-Schulz per group. DeepSeek-V4 o_a_proj is [8192, 4096] = 8 independent [1024, 4096] matrices (block-diagonal). It was orthogonalized as one matrix and now gets matrix_partitions, like the packed qkv and gate_up weights.

Affected parameters (meta-device check with the cached configs):

model	parameter	before	after
Qwen3.5 (dense and MoE)	linear_attn.conv1d.weight [6144,1,4]	Muon	AdamW
Qwen3.8-Flash-Next	linear_attn.conv1d.weight, ple.conv1d.weight [10240,1,4]	Muon	AdamW
Nemotron-H	mamba.conv1d.weight [6144,1,4]	Muon	AdamW
DeepSeek-V4	self_attn.o_a_proj.weight [8192,4096]	1 NS	8 x [1024,4096] NS

Nothing else changes.

Not changed, open for an A/B:

• LoRA A/B matrices (including expert LoRA) go to Muon, with no published recipe for this.
• If the vision tower is unfrozen, its Conv3d patch_embed and its pos_embed (an nn.Embedding that is not named embed_tokens) go to Muon.

Checks

• 1 node of 8xH200, offload = "none", optimizer step only:
  • DeepSeek-V4-Flash, 2 layers: o_a_proj now runs as (1024, 4096) Newton-Schulz calls and no (8192, 4096) call remains. Step 0.196 s (main 0.195 s).
  • Qwen3.5-35B-A3B, 4 layers: runs, 0.04 s.
• tests/unit/orchestrator + tests/unit/test_configs.py -m "not gpu": 350 passed. ruff check, ruff format --check.

---

Note

Medium Risk
Changes which optimizer and orthogonalization path apply to specific trainable weights on several model families, so training dynamics for those parameters will differ from before.

Overview
Corrects Muon parameter routing so two weight shapes are no longer treated as single large matrices.

Depthwise nn.Conv1d weights (e.g. linear-attention, Mamba, PLE convs) are excluded from Muon and trained with AdamW. Muon was orthogonalizing each channel’s 1×K taps separately, which is wrong for these per-channel filters.

GroupedLinear weights (e.g. DeepSeek-V4 o_a_proj) now get matrix_partitions so Newton–Schulz runs once per output group—matching how packed fusion weights are handled—instead of one orthogonalization over the full stacked block.

^{Reviewed by Cursor Bugbot for commit 64d9806. Bugbot is set up for automated code reviews on this repo. Configure here.}