[DO NOT MERGE] fix: route muon-indivisible params to adamw (topology-dependent dion assert)

[hallerite]

Fixes #3437: the first Muon step crashes training a large custom-impl MoE checkpoint on a 2-train-node topology:

dion/muon.py:457 muon_update_batch_async
AssertionError: Shard dimension 0 size 72 is not divisible by world size 16.

Root cause (corrected after review — this is not a #3411 regression, and nothing changed in how weights reach the optimizer): dion's distributed Muon shards each parameter's leading dim across the mesh and asserts divisibility. The checkpoint has heterogeneous attention gating — 12 full-attention layers with g_proj (48, hidden) and 36 sliding-attention layers with g_proj (72, hidden). The muon mesh is dp_shard_cp: 72 divides a 1-node world (8) but not a 2-node world with cp=4 (16) — so the same config steps fine on one topology and asserts on another. Dion pin and the muon param-grouping are unchanged for months; the failure is a divisibility accident of topology × checkpoint.

Fix: extend the muon_enabled predicate — parameters whose dim0 is not divisible by the muon mesh size go to the existing adamw group, exactly like the other muon-unsuitable parameters (1D, embeddings, lm_head). That makes routing topology-robust; for the affected model it's 36 small gate matrices (~8M params), negligible for optimization.

Notes:

• The check uses the mesh the default muon group is built on (dp_shard_cp, else world). The expert group's dedicated EP mesh can differ; fused expert dims divide cleanly in practice — a per-group check is the follow-up if that ever changes.
• Alternatives considered: padding/remainder handling inside dion (external package) — heavier, and arguably the better long-term home; this PR keeps prime-rl safe either way.

🤖 Generated with Claude Code

[S1ro1]

This would silently introduce different optimization based on dp_shard_cp size and also route various 2D params silently to adamw even though they should've been optimized by Muon. Should fix in DION library