[muon] Per-head Muon for linear-attention layers: read head geometry from the owning module

deepspeed/runtime/engine.py

@@ -613,6 +613,9 @@ def __init__(self,
         del autoep_replacement_sources
         if self.autotp_size() > 1:
             self._configure_tensor_parallel(model, self.tensor_parallel_config())
+            # Head counts were recorded against the whole model; the parameters are shards now.
+            from deepspeed import resolve_per_head_muon_after_sharding
+            resolve_per_head_muon_after_sharding(model)
         see_memory_usage("DeepSpeed Engine: After args sanity test", force=self.memory_breakdown())
         if mpu is not None:
             if self.elasticity_enabled():

deepspeed/runtime/zero/muon/original_muon.py

@@ -33,6 +33,19 @@
 from deepspeed.accelerator import get_accelerator
 
 
+def ns_compute_dtype(ns_method: str = "gram") -> torch.dtype:
+    """The dtype a Newton-Schulz iteration runs in, by method.
+
+    `gram` uses fp16 for better precision than bf16, `standard` uses bf16, and either falls
+    back to fp32 where the accelerator does not support its choice. Exported so that anything
+    reasoning about NS precision -- test tolerances in particular -- reads it from here rather
+    than restating it, which would let the two drift apart silently.
+    """
+    if ns_method == "gram":
+        return torch.float16 if get_accelerator().is_fp16_supported() else torch.float32
+    return torch.bfloat16 if get_accelerator().is_bf16_supported() else torch.float32
+
+
 @compiler.compile()
 def zeropower_via_newtonschulz5(G, steps: int):
     """
@@ -46,8 +59,7 @@ def zeropower_via_newtonschulz5(G, steps: int):
     """
     assert G.ndim >= 2  # batched Muon implementation by @scottjmaddox, and put into practice in the record by @YouJiacheng
     a, b, c = (3.4445, -4.7750, 2.0315)
-    # Use bf16 when hardware supports it; fp32 otherwise
-    compute_dtype = torch.bfloat16 if get_accelerator().is_bf16_supported() else torch.float32
+    compute_dtype = ns_compute_dtype("standard")
     X = G.to(compute_dtype)
     if G.size(-2) > G.size(-1):
         X = X.mT
@@ -86,8 +98,7 @@ def zeropower_via_gram_newtonschulz(G, steps: int):
     """
     assert G.ndim >= 2
     a, b, c = (3.4445, -4.7750, 2.0315)
-    # Use fp16 for better precision than bf16 when hardware supports it; fp32 otherwise
-    compute_dtype = torch.float16 if get_accelerator().is_fp16_supported() else torch.float32
+    compute_dtype = ns_compute_dtype("gram")
     X = G.to(compute_dtype)
     if G.size(-2) > G.size(-1):
         X = X.mT
@@ -142,11 +153,47 @@ def zeropower_via_gram_newtonschulz(G, steps: int):
 NS_METHODS = {"standard", "gram"}
 
 
+def _per_head_orthogonalize(update, num_heads, ns_steps, ns_method):
+    """Newton-Schulz per attention head, then fold the head dim back."""
+    if update.ndim != 2:
+        raise ValueError(f"Per-head Muon expects a 2D attention projection, got shape {tuple(update.shape)}.")
+
+    out_features, in_features = update.shape
+    if num_heads < 1 or out_features % num_heads != 0:
+        raise ValueError(f"Per-head Muon needs the output dim to split evenly across heads, but "
+                         f"{out_features} is not divisible by num_heads={num_heads}.")
+
+    head_dim = out_features // num_heads
+    ns_fn = zeropower_via_gram_newtonschulz if ns_method == "gram" else zeropower_via_newtonschulz5
+    # Scale per head block, matching what the full-matrix path does for the whole matrix.
+    scale = max(1, head_dim / in_features)**0.5
+    per_head = ns_fn(update.view(num_heads, head_dim, in_features), steps=ns_steps) * scale
+
+    return per_head.reshape(out_features, in_features)
+
+
 @compiler.compile()
-def muon_update(grad, momentum, beta=0.95, ns_steps=5, nesterov=True, ns_method="gram", is_expert_group=False):
+def muon_update(grad,
+                momentum,
+                beta=0.95,
+                ns_steps=5,
+                nesterov=True,
+                ns_method="gram",
+                is_expert_group=False,
+                num_heads=None):
+    """Muon update, optionally orthogonalizing each attention head separately.
+
+    With ``num_heads`` set, the update for an attention projection of shape
+    ``[num_heads * head_dim, in_features]`` is viewed as ``[num_heads, head_dim, in_features]``
+    and Newton-Schulz runs on that batch, so each head is orthogonalized against itself instead
+    of sharing one update direction with every other head. Both NS kernels are already batched,
+    so this is the same path the expert-group branch takes.
+    """
     orig_dtype = grad.dtype
     momentum.lerp_(grad, 1 - beta)
     update = grad.lerp_(momentum, beta) if nesterov else momentum
+    if num_heads is not None:
+        return _per_head_orthogonalize(update, num_heads, ns_steps, ns_method).to(orig_dtype)
     if is_expert_group:
         ns_fn = zeropower_via_gram_newtonschulz if ns_method == "gram" else zeropower_via_newtonschulz5
         scale = max(1, update.size(-2) / update.size(-1))**0.5
@@ -219,7 +266,8 @@ def step(self, closure=None):
                                          state["momentum_buffer"],
                                          beta=group["momentum"],
                                          ns_method=group.get("ns_method", "gram"),
-                                         is_expert_group=getattr(p, 'is_expert_group', False))
+                                         is_expert_group=getattr(p, 'is_expert_group', False),
+                                         num_heads=getattr(p, 'muon_num_heads', None))
                     p.mul_(1 - group["lr"] * group["weight_decay"])
                     p.add_(update.reshape(p.shape), alpha=-group["lr"])
                 dist.all_gather(params_pad[base_i:base_i + dist.get_world_size()],
@@ -257,7 +305,8 @@ def step(self, closure=None):
                                      state["momentum_buffer"],
                                      beta=group["momentum"],
                                      ns_method=group.get("ns_method", "gram"),
-                                     is_expert_group=getattr(p, 'is_expert_group', False))
+                                     is_expert_group=getattr(p, 'is_expert_group', False),
+                                     num_heads=getattr(p, 'muon_num_heads', None))
                 p.mul_(1 - group["lr"] * group["weight_decay"])
                 p.add_(update.reshape(p.shape), alpha=-group["lr"])
 
@@ -349,7 +398,8 @@ def step(self, closure=None):
                                              state["momentum_buffer"],
                                              beta=group["momentum"],
                                              ns_method=group.get("ns_method", "gram"),
-                                             is_expert_group=getattr(p, 'is_expert_group', False))
+                                             is_expert_group=getattr(p, 'is_expert_group', False),
+                                             num_heads=getattr(p, 'muon_num_heads', None))
                         p.mul_(1 - group["lr"] * group["weight_decay"])
                         p.add_(update.reshape(p.shape), alpha=-group["lr"])
                     dist.all_gather(params_pad[base_i:base_i + dist.get_world_size()],
@@ -421,7 +471,8 @@ def step(self, closure=None):
                                          state["momentum_buffer"],
                                          beta=group["momentum"],
                                          ns_method=group.get("ns_method", "gram"),
-                                         is_expert_group=getattr(p, 'is_expert_group', False))
+                                         is_expert_group=getattr(p, 'is_expert_group', False),
+                                         num_heads=getattr(p, 'muon_num_heads', None))
                     p.mul_(1 - group["lr"] * group["weight_decay"])
                     p.add_(update.reshape(p.shape), alpha=-group["lr"])
             else:

deepspeed/runtime/zero/stage3.py

@@ -1692,7 +1692,11 @@ def _apply_distributed_muon_update(self, communication_data_type: torch.dtype, b
                     param = params[base_i + rank]
                     g = param.grad
                     m = gathered_momentums_pad[base_i + rank]
-                    update = muon_update(g, m, beta=self.muon_beta, ns_method=getattr(self, 'muon_ns_method', 'gram'))
+                    update = muon_update(g,
+                                         m,
+                                         beta=self.muon_beta,
+                                         ns_method=getattr(self, 'muon_ns_method', 'gram'),
+                                         num_heads=getattr(param, 'muon_num_heads', None))
                     g.data.copy_(update, non_blocking=False)
                 grad_handle = dist.all_gather(grads_pad[base_i:base_i + world_sz],
                                               grads_pad[base_i + rank],

deepspeed/runtime/zero/stage_1_and_2.py

@@ -2178,7 +2178,8 @@ def get_flat_partition(self,
                                          buffer,
                                          self.optimizer.param_groups[param_group_idx]['momentum'],
                                          ns_method=ns_method,
-                                         is_expert_group=getattr(tensor, 'is_expert_group', False))
+                                         is_expert_group=getattr(tensor, 'is_expert_group', False),
+                                         num_heads=getattr(tensor, 'muon_num_heads', None))
             tensor = grad_accum
             num_elements = tensor.numel()
             buffer_idx += num_elements

docs/_pages/config-json.md

@@ -60,6 +60,56 @@ Muon supports the following params:
 | torch\_adam    | Use torch Adam/AdamW for non-Muon parameters instead of the DeepSpeed Adam backend.                                  | false     |
 | adam\_w\_mode | Use AdamW rather than Adam for non-Muon parameters.                                                                  | true      |
 | ns\_method     | Newton-Schulz orthogonalization method: `"gram"` for Gram NS (~2x faster on rectangular matrices), `"standard"` for the original iteration. Use `"standard"` to fall back if you encounter convergence issues. | `"gram"`  |
+| per\_head\_muon | Orthogonalize each attention head separately instead of the whole projection. See below. | false |
+
+#### Per-head Muon
+
+With `per_head_muon: true`, an attention projection shaped `[num_heads * head_dim, in_features]`
+is viewed as `[num_heads, head_dim, in_features]` and Newton-Schulz runs on that batch, so each
+head is orthogonalized against itself rather than sharing one update direction with every other
+head. This is the split described by Kimi K3 ("Per-Head Muon") and GLM-5 ("Muon Split"). Off by
+default; communication volume is unchanged.
+
+What is tagged, and what deliberately is not:
+
+| matrix | per-head | why |
+| --- | --- | --- |
+| `q_proj` / `query` / `wq` | yes | blocked by the query head count |
+| `k_proj` / `v_proj` / `key` / `value` / `wk` / `wv` | yes | blocked by the KV head count, which differs from the query count under GQA |
+| MLA `q_b_proj`, `kv_b_proj` | yes | the two up-projections, whose per-head widths are `qk_nope + qk_rope` and `qk_nope + v_head_dim` rather than `head_dim` |
+| `o_proj` and other output projections | no | the head structure is on the input dimension, so splitting dim 0 would cut across the wrong axis |
+| fused `qkv_proj` / `query_key_value` / `c_attn` / `wqkv` | no | the three sections do not share a head count under GQA |
+| MLA `q_a_proj`, `kv_a_proj_with_mqa` | no | down-projections mixing latent and rope components, with no head structure |
+| linear-attention `q_proj` / `k_proj` / `v_proj` | yes | blocked by the head count the attention module was built with, which for hybrids such as Kimi-K3 is not `num_attention_heads * head_dim` |
+| sparse-attention indexers | no | the indexer selects which keys attention will see; the split is defined on attention itself |
+
+**Where the geometry comes from.** The config first: head counts through `AutoTPMeta`, per-head
+widths from the fields the architecture defines. When no config geometry confirms, the module
+that owns the projection is asked for the counts it was built with, which is how hybrid models
+that keep their linear-attention geometry outside the top-level config fields are covered. Only
+modules that identify as attention are asked, so a module with a head count of its own that is
+not attention keeps the full-matrix path.
+
+**The shape confirms the name.** A leaf name is treated as a claim about the layout, never as
+proof of it. Every geometry the config makes plausible for that name is evaluated, and a
+parameter is tagged only when its rows equal `num_heads * width` exactly for one of them. Two
+geometries that confirm and agree on the head count are not a conflict; two that confirm and
+disagree are, and the parameter is skipped with a warning.
+
+**Tensor parallelism.** Column-parallel TP splits an attention projection on dim 0, which is
+the axis the heads are on, so a rank holds whole heads and the per-head width is unchanged. That
+makes the per-head split exact under TP: Newton-Schulz on a rank's heads is the same computation
+whether the other ranks' heads are present or not. The head *count* is not invariant, so with
+AutoTP the counts are re-resolved against the shards after partitioning; a shard whose rows are
+not a multiple of the per-head width does not hold whole heads and stays on the full-matrix path.
+A model that arrives already sharded by an external tensor-parallel implementation cannot be
+tagged at all, because the config then describes a width no parameter has.
+
+**The flag reports what it did.** Because it is an explicit opt-in, DeepSpeed raises at
+`deepspeed.initialize` if it is enabled and no attention projection could be tagged, rather than
+training on without it. Parameters that match an attention name but confirm no geometry are
+reported as a warning and stay on the full-matrix path, so a hybrid model still gets per-head on
+its recognized layers.
 
 By default, non-Muon parameters use `FusedAdam`. When optimizer state is offloaded to the CPU, DeepSpeed selects `DeepSpeedCPUAdam`. This is the same backend selection used by the Adam and AdamW optimizer types.