From af40f4cba2cd81335588858dd8020174c1c73b74 Mon Sep 17 00:00:00 2001 From: hallerite Date: Sun, 30 Aug 2026 19:54:33 +0000 Subject: [PATCH] fix: route muon-indivisible params to adamw instead of crashing the step dion's Muon distributes orthogonalization by sharding a parameter's leading dim across the distributed mesh and asserts divisibility ('Shard dimension 0 size 72 is not divisible by world size 16'). Models with small projections whose leading dim does not divide the mesh (e.g. attention-gate weights of shape (72, hidden)) crashed on the first optimizer step. Extend the muon_enabled predicate to route such parameters to the existing adamw group, like the other muon-unsuitable parameters (1D, embeddings, lm_head). The check uses the same mesh the default muon group is constructed with; expert groups on a dedicated EP mesh are unaffected in practice since fused expert dims divide cleanly. Fixes #3437. Co-Authored-By: Claude Fable 5 --- src/prime_rl/trainer/optim/__init__.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/src/prime_rl/trainer/optim/__init__.py b/src/prime_rl/trainer/optim/__init__.py index f19b9b07c9..c7f638da95 100644 --- a/src/prime_rl/trainer/optim/__init__.py +++ b/src/prime_rl/trainer/optim/__init__.py @@ -122,6 +122,11 @@ def _create_muon_optimizer( parallel_dims: ParallelDims, lr: float | None = None, ) -> Optimizer: + if parallel_dims.dp_shard_enabled or parallel_dims.cp_enabled: + muon_mesh_size = parallel_dims.get_mesh("dp_shard_cp").size() + else: + muon_mesh_size = parallel_dims.world_mesh.size() + def muon_enabled(n, p): if p.ndim < 2: return False @@ -129,6 +134,12 @@ def muon_enabled(n, p): return False if "embed_tokens" in n: return False + # Muon distributes orthogonalization by sharding a parameter's leading dim + # across the mesh; dion asserts divisibility. Route indivisible params (e.g. + # small attention-gate projections) to adamw like the other muon-unsuitable + # params rather than crashing the step. + if p.shape[0] % muon_mesh_size != 0: + return False return True muon_params = []