Merge pull request #33 from MotifTechnologies/feat/per-head-muon

feat: add Per-Head Muon (Kimi K3 §2.5) via head_keys [skip-build]

CLAUDE.md

@@ -87,6 +87,7 @@ For detailed implementation documentation (pipeline internals, distributed utili
 - **`construct_shard_mesh()` / `get_slices_of_dtensor()`** (`distributed/utils.py`) — Utilities for building shard meshes from DTensor placements and computing per-rank local slices. Handles both `Shard` and `_StridedShard` (PyTorch 2.10+).
 - **Newton-Schulz iteration** (`newton_schulz.py`) — `_zeropower_via_newtonschulz5()`: 5 quintic iterations in bfloat16 with pre-optimized coefficients for gradient orthogonalization. Uses Triton kernel `matmul_transpose_assign` for efficient X @ X.T.
 - **QK Clipping** (`qk_clip.py`) — Optional dynamic clipping of attention head projections when QK logits exceed a threshold. Configured via `q_indices`, `k_indices`, `head_dim`, `threshold`.
+- **Per-Head Muon** (`head_keys`) — Kimi K3 §2.5 variant: attention Q/K/V momentum matrices are partitioned along the head dimension (dim 0) and each head block is orthogonalized separately via the batched Newton-Schulz. `head_keys` maps name keys to per-head row strides (e.g. `{"q_proj": 128}`); adjusted LR uses the head-block shape. Works across all execution paths since NS always runs on the full gathered matrix.
 - **Fused AdamW** (`adamw.py`) — Uses PyTorch's `torch._fused_adamw_` for non-Muon parameters, grouping tensors by device/dtype and DTensor placement.
 
 ### Dependency Graph

docs/implementation.md

@@ -9,11 +9,12 @@ This document explains the internal architecture of the Muon optimizer for revie
 3. [Execution Paths](#execution-paths)
 4. [Parallel Pipeline (the core feature)](#parallel-pipeline)
 5. [MoE Expert Weight Support](#moe-expert-weight-support-expert_keys)
-6. [Distributed Utilities](#distributed-utilities)
-7. [Newton-Schulz Orthogonalization](#newton-schulz-orthogonalization)
-8. [QK Clipping](#qk-clipping)
-9. [AdamW for Non-Muon Parameters](#adamw-for-non-muon-parameters)
-10. [Source File Map](#source-file-map)
+6. [Per-Head Muon](#per-head-muon-head_keys)
+7. [Distributed Utilities](#distributed-utilities)
+8. [Newton-Schulz Orthogonalization](#newton-schulz-orthogonalization)
+9. [QK Clipping](#qk-clipping)
+10. [AdamW for Non-Muon Parameters](#adamw-for-non-muon-parameters)
+11. [Source File Map](#source-file-map)
 
 ---
 
@@ -158,7 +159,7 @@ class _muon_state:
 
 #### Stage 3: Compute
 
-The owning rank runs `_zeropower_via_newtonschulz5()` on the full gathered gradient. This is the most compute-intensive stage. Runs inline (no yield) since it is synchronous GPU work.
+The owning rank runs `_zeropower_via_newtonschulz5()` on the full gathered gradient. This is the most compute-intensive stage. Runs inline (no yield) since it is synchronous GPU work. Parameters with `_muon_state.head_dim` set (Per-Head Muon) instead dispatch to `zeropower_via_newtonschulz5_per_head()`, which orthogonalizes each attention head's block separately.
 
 #### Stages 4-5: Scatter
 
@@ -215,6 +216,26 @@ After expansion, the 2D params flow through the standard routing: DTensors with
 
 For EP/EFSDP background and torchtitan integration details, see [`docs/expert_parallel.md`](expert_parallel.md).
 
+## Per-Head Muon (`head_keys`)
+
+**Files:** `newton_schulz.py` — `zeropower_via_newtonschulz5_per_head()`, `core.py` — `get_head_dim()` / `resolve_head_dim()`
+
+Implements the per-head variant of Muon from the Kimi K3 tech report (§2.5). Full-matrix orthogonalization treats all attention heads as a single coupled block, so heads with larger momentum scale dominate the shared update while small-scale heads receive insufficiently normalized updates. Per-Head Muon partitions the momentum matrix of Q/K/V projections along the head dimension (dim 0) and orthogonalizes each head's block separately, equalizing the update scale across heads. It is also slightly cheaper: the Gram matrix per block is `head_dim x head_dim` instead of `min(M, N)^2`.
+
+### Configuration
+
+```python
+optim = Muon(params, head_keys={"q_proj": 128, "k_proj": 128, "v_proj": 128}, ...)
+```
+
+`head_keys` maps parameter-name keys (component-level matching, same rule as `expert_keys`) to the dim-0 row stride of one head. For MLA, pass each projection's per-head stride, e.g. `{"wq_b": qk_nope_head_dim + qk_rope_head_dim, "wkv_b": qk_nope_head_dim + v_head_dim}`. Only dim-0 (output-side) head layouts are supported — Q/K/V projections, not `o_proj`. Keys must be disjoint from `expert_keys`.
+
+### How It Works
+
+- **Momentum is unchanged** (elementwise); only the Newton-Schulz block granularity changes. A matching `[H * head_dim, N]` matrix is viewed as `[H, head_dim, N]` and run through the batched bmm-based Newton-Schulz (the same kernel as per-expert NS), which normalizes each block by its own Frobenius norm.
+- **The size-adjusted LR uses the head-block shape** `(head_dim, N)`, mirroring the per-expert path where each block is the NS unit.
+- **Sharding-agnostic**: all three execution paths (`base()`, `distributed_muon()`, `parallel()`) apply the head partition to the *full global* matrix — the parallel path gathers the full matrix on the owner rank before NS — so the result is identical for any DP/TP/HSDP placement, including shard boundaries that cut mid-head. No communication changes are needed; `_muon_state.head_dim` just switches the compute call.
+
 ## Distributed Utilities
 
 **File:** `distributed/utils.py`

test/conftest.py

@@ -4,7 +4,6 @@ import pytest
 import torch
 import torch.distributed as dist
 from packaging import version
-from transformers import AutoModelForCausalLM
 
 logger = logging.getLogger(__name__)
 logging.basicConfig(level=logging.INFO)
@@ -99,6 +98,10 @@ def inputs():
         - list[torch.Tensor]: A list of random gradients for each model parameter.
         - dict[int, torch.Tensor]: A dictionary mapping layer indices to random QK logits.
     """
+    # Lazy import so test files that don't use this fixture (e.g. the
+    # synthetic per-head tests) can run in images without transformers.
+    from transformers import AutoModelForCausalLM
+
     model_name = "Motif-Technologies/Motif-2.6B-4layer-random"
 
     torch.manual_seed(SEED)

test/run_test_per_head.sh

@@ -0,0 +1 @@
+torchrun --nproc-per-node=8 --local-ranks-filter=0 -m pytest test_muon_per_head.py

test/test_muon_per_head.py

@@ -0,0 +1,485 @@
+"""Tests for Per-Head Muon (Kimi K3 tech report, §2.5).
+
+Per-Head Muon partitions the momentum matrix of attention projections
+along the head dimension (dim 0) and orthogonalizes each head's block
+separately, instead of running Newton-Schulz on the full projection
+matrix.  These tests verify:
+
+1. Paper intent — per-head orthogonalization decouples heads (rescaling
+   one head's gradient does not change any head's update) and equalizes
+   the per-head update scale, unlike full-matrix Muon where heads with
+   larger momentum scale dominate the coupled update.
+2. Sharding correctness — the parallel (all2all) and distributed_muon
+   paths match the single-device base path bit-exactly under DP/TP/HSDP
+   placements, including shard boundaries that cut through the middle of
+   a head block.
+3. EP coexistence — expert_keys (per-expert NS) and head_keys (per-head
+   NS) work together in one optimizer instance.
+
+Run: torchrun --nproc-per-node=8 --local-ranks-filter=0 -m pytest test_muon_per_head.py
+"""
+
+import logging
+import math
+
+import pytest
+import torch
+import torch.distributed as dist
+from optimizer.core import adjust_lr_for_muon
+from optimizer.muon import Muon
+from optimizer.newton_schulz import set_ns_compile
+from torch.distributed.tensor import (DTensor, Replicate, Shard,
+                                      distribute_tensor)
+
+logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
+LR = 0.02
+WD = 0.01
+
+
+def _param_group(params, names, ns_steps=5, weight_decay=WD):
+    return [{
+        "params": params,
+        "names": names,
+        "use_muon": True,
+        "lr": LR,
+        "weight_decay": weight_decay,
+        "momentum": 0.95,
+        "nesterov": True,
+        "ns_steps": ns_steps,
+        "none_grad": False,
+    }]
+
+
+def _run_muon(full_params,
+              step_grads,
+              names,
+              head_keys=None,
+              expert_keys=None,
+              ns_steps=5,
+              **muon_kwargs):
+    """Run Muon on fresh plain-tensor params (single-device base path)."""
+    params = [torch.nn.Parameter(t.clone()) for t in full_params]
+    optim = Muon(params=_param_group(params, list(names), ns_steps=ns_steps),
+                 head_keys=head_keys,
+                 expert_keys=expert_keys,
+                 **muon_kwargs)
+    for grads in step_grads:
+        for p, g in zip(params, grads):
+            p.grad = g.clone()
+        optim.step()
+    return [p.data for p in params]
+
+
+def _run_muon_sharded(full_params,
+                      step_grads,
+                      names,
+                      mesh,
+                      placements,
+                      head_keys=None,
+                      expert_keys=None,
+                      ns_steps=5,
+                      **muon_kwargs):
+    """Run Muon on DTensor-sharded params, return gathered full tensors."""
+    params = []
+    for t in full_params:
+        dt = distribute_tensor(t.clone(), mesh, placements)
+        params.append(torch.nn.Parameter(dt))
+    optim = Muon(params=_param_group(params, list(names), ns_steps=ns_steps),
+                 head_keys=head_keys,
+                 expert_keys=expert_keys,
+                 **muon_kwargs)
+    for grads in step_grads:
+        for p, g in zip(params, grads):
+            p.grad = distribute_tensor(g.clone(), mesh, p.data.placements)
+        optim.step()
+    return [p.data.full_tensor() for p in params]
+
+
+def _make_attention_like_tensors(num_steps=2, seed=42):
+    """Attention-like params: shard boundaries deliberately cut mid-head
+    for q_proj (head_dim=40, 8-way row shard → 30 rows per rank), and
+    v_proj has tall head blocks (head_dim 96 > in_dim 48) to exercise the
+    transpose path of the batched Newton-Schulz.
+    """
+    torch.manual_seed(seed)
+    specs = [
+        ("layers.0.self_attn.q_proj.weight", (240, 64)),  # 6 heads x 40
+        ("layers.0.self_attn.k_proj.weight", (128, 64)),  # 8 heads x 16
+        ("layers.0.self_attn.v_proj.weight", (192, 48)),  # 2 heads x 96
+        ("layers.0.mlp.up_proj.weight", (144, 96)),  # control: full-matrix
+    ]
+    names = [n for n, _ in specs]
+    fulls = [torch.randn(s, device="cuda") for _, s in specs]
+    step_grads = [[torch.randn(s, device="cuda") for _, s in specs]
+                  for _ in range(num_steps)]
+    return names, fulls, step_grads
+
+
+HEAD_KEYS = {"q_proj": 40, "k_proj": 16, "v_proj": 96}
+
+# Note: plain nested same-dim sharding ([Shard(0), Shard(0)]) is rejected
+# by construct_shard_mesh (ambiguous layout).  The real FSDP2 + colwise-TP
+# nested dim-0 layout (_StridedShard) is covered by
+# test_per_head_fsdp2_tp_colwise below.
+SHARDING_CASES = [
+    pytest.param(((8, ), ("dp_shard", ), [Shard(0)]), id="fsdp"),
+    pytest.param(((8, ), ("tp", ), [Shard(1)]), id="tp_row"),
+    pytest.param(
+        ((2, 4),
+         ("dp_replicate", "dp_shard"), [Replicate(), Shard(0)]),
+        id="hsdp"),
+    pytest.param(((2, 4), ("dp_shard", "tp"), [Shard(0), Shard(1)]),
+                 id="fsdp+tp_row"),
+]
+
+# ======================================================================
+# Sharding correctness (TP / DP / HSDP)
+# ======================================================================
+
+
+@pytest.mark.parametrize("sharding", SHARDING_CASES)
+@pytest.mark.parametrize("use_distributed_muon", [False, True],
+                         ids=["parallel", "distributed"])
+def test_per_head_sharded_matches_base(init_dist, sharding,
+                                       use_distributed_muon):
+    """Sharded per-head Muon must match the single-device base path
+    bit-exactly.  Runs 2 steps so the second step exercises head_dim
+    propagation through the _parallel_cache cached-state path.
+    """
+    mesh_shape, mesh_dim_names, placements = sharding
+    set_ns_compile(False)
+    try:
+        names, fulls, step_grads = _make_attention_like_tensors()
+        base_result = _run_muon(fulls, step_grads, names, head_keys=HEAD_KEYS)
+
+        mesh = dist.init_device_mesh("cuda",
+                                     mesh_shape,
+                                     mesh_dim_names=mesh_dim_names)
+        sharded_result = _run_muon_sharded(
+            fulls,
+            step_grads,
+            names,
+            mesh,
+            placements,
+            head_keys=HEAD_KEYS,
+            chunk_size=2,
+            warmup_step=1,
+            use_distributed_muon=use_distributed_muon)
+
+        for name, b, s in zip(names, base_result, sharded_result):
+            torch.testing.assert_close(s,
+                                       b,
+                                       atol=0,
+                                       rtol=0,
+                                       msg=lambda m, name=name: f"{name}: {m}")
+    finally:
+        set_ns_compile(True)
+
+
+@pytest.mark.parametrize("sharding", [SHARDING_CASES[0], SHARDING_CASES[-1]])
+def test_per_head_sharded_matches_base_compiled(init_dist, sharding):
+    """Same parity check with torch.compile + CUDA graphs enabled for the
+    batched per-head Newton-Schulz (tolerances match test_muon.py's
+    compiled runs)."""
+    mesh_shape, mesh_dim_names, placements = sharding
+    set_ns_compile(True)
+
+    names, fulls, step_grads = _make_attention_like_tensors()
+    base_result = _run_muon(fulls, step_grads, names, head_keys=HEAD_KEYS)
+
+    mesh = dist.init_device_mesh("cuda",
+                                 mesh_shape,
+                                 mesh_dim_names=mesh_dim_names)
+    sharded_result = _run_muon_sharded(fulls,
+                                       step_grads,
+                                       names,
+                                       mesh,
+                                       placements,
+                                       head_keys=HEAD_KEYS,
+                                       chunk_size=2,
+                                       warmup_step=1)
+
+    for name, b, s in zip(names, base_result, sharded_result):
+        torch.testing.assert_close(s,
+                                   b,
+                                   atol=1e-5,
+                                   rtol=1e-2,
+                                   msg=lambda m, name=name: f"{name}: {m}")
+
+
+# ======================================================================
+# Paper-intent verification (Kimi K3 §2.5)
+# ======================================================================
+
+
+def _single_param_update(name, head_keys, grad, ns_steps=10):
+    """One Muon step on a zero-init param with wd=0.
+
+    With p0 = 0 and weight_decay = 0, the post-step param is exactly
+    ``-adjusted_lr * u``, i.e. proportional to the orthogonalized update.
+    """
+    p = torch.nn.Parameter(torch.zeros_like(grad))
+    optim = Muon(params=_param_group([p], [name],
+                                     ns_steps=ns_steps,
+                                     weight_decay=0.0),
+                 head_keys=head_keys)
+    p.grad = grad.clone()
+    optim.step()
+    return p.data
+
+
+def test_per_head_decouples_heads(init_dist):
+    """K3 §2.5: full-matrix orthogonalization treats all heads as a single
+    coupled block, so heads with larger momentum scale dominate the shared
+    update; per-head orthogonalization removes that coupling.
+
+    Rescaling one head's gradient by 1024 (a power of two — exact in
+    fp32/bf16, and each head block is normalized by its own Frobenius
+    norm) must leave every per-head update bit-identical, while the
+    full-matrix update of the *other* heads changes materially.
+    """
+    set_ns_compile(False)
+    try:
+        torch.manual_seed(7)
+        H, d, N = 8, 32, 128
+        scaled_head = 3
+        name = "layers.0.self_attn.q_proj.weight"
+
+        g = torch.randn(H * d, N, device="cuda")
+        g_scaled = g.clone()
+        g_scaled[scaled_head * d:(scaled_head + 1) * d] *= 1024.0
+
+        u1 = _single_param_update(name, {"q_proj": d}, g).view(H, d, N)
+        u2 = _single_param_update(name, {"q_proj": d}, g_scaled).view(H, d, N)
+        torch.testing.assert_close(
+            u2,
+            u1,
+            atol=0,
+            rtol=0,
+            msg="per-head updates must be invariant to per-head rescaling")
+
+        f1 = _single_param_update(name, None, g).view(H, d, N)
+        f2 = _single_param_update(name, None, g_scaled).view(H, d, N)
+        other = [h for h in range(H) if h != scaled_head]
+        rel = ((f2[other] - f1[other]).norm() / f1[other].norm()).item()
+        # The per-head diff above is exactly zero; any clearly nonzero
+        # coupling on the full-matrix side proves the contrast (measured
+        # ~0.4 for this configuration).
+        assert rel > 0.1, (
+            f"full-matrix NS should couple heads, got rel diff {rel:.3f}")
+        logger.info("decoupling: full-matrix rel diff on other heads = %.3f",
+                    rel)
+    finally:
+        set_ns_compile(True)
+
+
+def test_per_head_equalizes_update_scale(init_dist):
+    """K3 §2.5: with imbalanced per-head momentum scales, full-matrix NS
+    leaves small-scale heads insufficiently normalized, while per-head NS
+    gives every head a near semi-orthogonal update block (Frobenius norm
+    ≈ sqrt(head_dim) for all heads).
+
+    Uses a tall full matrix (H*d > N, the K3 regime) — for a wide matrix
+    the polar factor has orthonormal rows and the imbalance would be
+    invisible.
+    """
+    set_ns_compile(False)
+    try:
+        torch.manual_seed(11)
+        H, d, N = 16, 32, 256  # H*d = 512 > N = 256
+        name = "layers.0.self_attn.q_proj.weight"
+        head_scales = torch.logspace(-2, 2, H, device="cuda")
+        g = torch.randn(H * d, N, device="cuda")
+        g = (g.view(H, d, N) * head_scales.view(H, 1, 1)).reshape(H * d, N)
+
+        # --- per-head ---
+        adj_ph = adjust_lr_for_muon(LR, (d, N))
+        u_ph = (-_single_param_update(name, {"q_proj": d}, g) / adj_ph).view(
+            H, d, N)
+        norms_ph = u_ph.norm(dim=(-2, -1))
+        ratio_ph = (norms_ph.max() / norms_ph.min()).item()
+
+        # Every head block ~ semi-orthogonal: ||U_h||_F ≈ sqrt(d), and
+        # U_h @ U_h^T ≈ I.
+        assert ratio_ph < 1.2, (
+            f"per-head update scale not equalized: max/min = {ratio_ph:.3f}")
+        torch.testing.assert_close(norms_ph,
+                                   torch.full_like(norms_ph, math.sqrt(d)),
+                                   rtol=0.15,
+                                   atol=0)
+        gram = torch.bmm(u_ph, u_ph.transpose(-2, -1))
+        eye = torch.eye(d, device="cuda").expand(H, d, d)
+        max_dev = (gram - eye).abs().max().item()
+        assert max_dev < 0.2, (
+            f"per-head update blocks not row-orthogonal: max dev {max_dev}")
+
+        # --- full-matrix (control) ---
+        adj_full = adjust_lr_for_muon(LR, (H * d, N))
+        u_full = (-_single_param_update(name, None, g) / adj_full).view(
+            H, d, N)
+        norms_full = u_full.norm(dim=(-2, -1))
+        ratio_full = (norms_full.max() / norms_full.min()).item()
+        assert ratio_full > 3.0, (
+            f"expected full-matrix NS to under-normalize dominated heads, "
+            f"got max/min = {ratio_full:.3f}")
+        logger.info("update-scale max/min ratio: per-head=%.3f, full=%.1f",
+                    ratio_ph, ratio_full)
+    finally:
+        set_ns_compile(True)
+
+
+def test_per_head_fsdp2_tp_colwise(init_dist):
+    """Real FSDP2 + colwise-TP layout — the production layout for Q/K/V.
+
+    ColwiseParallel gives Shard(0) on the tp mesh dim, then fully_shard
+    adds a strided shard on dim 0 over the dp mesh dim (_StridedShard):
+    nested dim-0 sharding where the per-rank pieces interleave head
+    blocks.  Per-head NS must still match the single-device base path
+    bit-exactly because the pipeline gathers the full global matrix
+    before orthogonalization.
+    """
+    import torch.nn as nn
+    from torch.distributed.fsdp import fully_shard
+    from torch.distributed.tensor.parallel import (ColwiseParallel,
+                                                   parallelize_module)
+
+    set_ns_compile(False)
+    try:
+        torch.manual_seed(21)
+
+        class Blk(nn.Module):
+
+            def __init__(self):
+                super().__init__()
+                # weight shapes: (240, 64), (128, 64), (192, 48), (144, 96)
+                self.q_proj = nn.Linear(64, 240, bias=False)
+                self.k_proj = nn.Linear(64, 128, bias=False)
+                self.v_proj = nn.Linear(48, 192, bias=False)
+                self.up_proj = nn.Linear(96, 144, bias=False)
+
+        model = Blk().cuda()
+        names = [n for n, _ in model.named_parameters()]
+        fulls = [p.data.clone() for p in model.parameters()]
+        step_grads = [[torch.randn_like(t) for t in fulls] for _ in range(2)]
+
+        base_result = _run_muon(fulls, step_grads, names, head_keys=HEAD_KEYS)
+
+        mesh = dist.init_device_mesh("cuda", (2, 4),
+                                     mesh_dim_names=("dp_shard", "tp"))
+        parallelize_module(
+            model, mesh["tp"], {
+                "q_proj": ColwiseParallel(),
+                "k_proj": ColwiseParallel(),
+                "v_proj": ColwiseParallel(),
+                "up_proj": ColwiseParallel(),
+            })
+        fully_shard(model, mesh=mesh["dp_shard"])
+
+        params = list(model.parameters())
+        logger.info("FSDP2+TP placements: %s", params[0].data.placements)
+        optim = Muon(params=_param_group(params, names),
+                     head_keys=HEAD_KEYS,
+                     chunk_size=2,
+                     warmup_step=1)
+        for grads in step_grads:
+            for p, g in zip(params, grads):
+                full = DTensor.from_local(g.clone(),
+                                          device_mesh=p.data.device_mesh,
+                                          placements=[Replicate()] *
+                                          p.data.device_mesh.ndim)
+                p.grad = full.redistribute(device_mesh=p.data.device_mesh,
+                                           placements=p.data.placements)
+            optim.step()
+
+        for name, b, p in zip(names, base_result, model.parameters()):
+            torch.testing.assert_close(p.data.full_tensor(),
+                                       b,
+                                       atol=0,
+                                       rtol=0,
+                                       msg=lambda m, name=name: f"{name}: {m}")
+    finally:
+        set_ns_compile(True)
+
+
+# ======================================================================
+# EP coexistence and validation
+# ======================================================================
+
+
+def test_per_head_with_expert_keys(init_dist):
+    """head_keys (attention, per-head NS) and expert_keys (MoE, per-expert
+    NS) coexist in one optimizer.  Expert params take the batched-expert
+    path (which pre-launches the first A2A gather), attention params take
+    the parallel per-head path.  Sharded result must match the
+    single-device baseline.
+    """
+    world_size = dist.get_world_size()
+    set_ns_compile(False)
+    try:
+        torch.manual_seed(13)
+        E = world_size  # one expert per rank under Shard(0) — EP-style
+        specs = [
+            ("layers.0.moe.experts.w1", (E, 48, 32)),
+            ("layers.0.self_attn.q_proj.weight", (160, 64)),  # 5 heads x 32
+            ("layers.0.mlp.up_proj.weight", (96, 64)),
+        ]
+        names = [n for n, _ in specs]
+        fulls = [torch.randn(s, device="cuda") for _, s in specs]
+        step_grads = [[torch.randn_like(t) for t in fulls] for _ in range(2)]
+        head_keys = {"q_proj": 32}
+        expert_keys = ["experts"]
+
+        base_result = _run_muon(fulls,
+                                step_grads,
+                                names,
+                                head_keys=head_keys,
+                                expert_keys=expert_keys)
+
+        mesh = dist.init_device_mesh("cuda", (world_size, ),
+                                     mesh_dim_names=("dp", ))
+        sharded_result = _run_muon_sharded(fulls,
+                                           step_grads,
+                                           names,
+                                           mesh, [Shard(0)],
+                                           head_keys=head_keys,
+                                           expert_keys=expert_keys,
+                                           chunk_size=1,
+                                           warmup_step=1)
+
+        for name, b, s in zip(names, base_result, sharded_result):
+            torch.testing.assert_close(s,
+                                       b,
+                                       atol=0,
+                                       rtol=0,
+                                       msg=lambda m, name=name: f"{name}: {m}")
+    finally:
+        set_ns_compile(True)
+
+
+def test_head_keys_validation(init_dist):
+    """Misconfigured head_keys fail loudly, before touching comms."""
+    # dim-0 not divisible by head_dim → ValueError at step time.
+    set_ns_compile(False)
+    try:
+        p = torch.nn.Parameter(torch.randn(30, 16, device="cuda"))
+        p.grad = torch.randn_like(p)
+        optim = Muon(params=_param_group([p],
+                                         ["layers.0.self_attn.q_proj.weight"]),
+                     head_keys={"q_proj": 7})
+        with pytest.raises(ValueError, match="not divisible"):
+            optim.step()
+    finally:
+        set_ns_compile(True)
+
+    # head_dim must be a positive int.
+    p = torch.nn.Parameter(torch.randn(8, 8, device="cuda"))
+    with pytest.raises(ValueError, match="positive int"):
+        Muon(params=_param_group([p], ["w"]), head_keys={"q_proj": 0})
+
+    # head_keys and expert_keys must be disjoint.
+    with pytest.raises(ValueError, match="disjoint"):
+        Muon(params=_param_group([p], ["w"]),
+             head_keys={"experts": 4},
+             expert_keys=["experts"])

torch-ext/optimizer/core.py

@@ -29,6 +29,9 @@ class _muon_state:
     rank_numels: dict[int, int]  # local_rank -> numel
     name: str
     qk_clip_state: torch.Tensor | None = None
+    # Per-Head Muon: dim-0 row stride of one attention head.  When set,
+    # Newton-Schulz runs per head block instead of on the full matrix.
+    head_dim: int | None = None
 
 
 def _batch_momentum(
@@ -157,6 +160,34 @@ def is_expert_param(name, expert_keys):
     return any(_match_key(parts, key) for key in expert_keys)
 
 
+def get_head_dim(name, head_keys):
+    """Return the per-head dim-0 row stride for a param matching head_keys.
+
+    ``head_keys`` maps name keys (component-level matching, same rule as
+    ``expert_keys``) to the number of dim-0 rows each attention head
+    occupies.  Returns None if the param matches no key.
+    """
+    if not head_keys:
+        return None
+    parts = normalize_fqn(name).split(".")
+    for key, head_dim in head_keys.items():
+        if _match_key(parts, key):
+            return head_dim
+    return None
+
+
+def resolve_head_dim(name, shape, head_keys):
+    """``get_head_dim`` plus dim-0 divisibility validation."""
+    head_dim = get_head_dim(name, head_keys)
+    if head_dim is None:
+        return None
+    if shape[0] % head_dim != 0:
+        raise ValueError(
+            f"Param {name}: dim-0 size {shape[0]} is not divisible by "
+            f"head_dim={head_dim} from head_keys.")
+    return head_dim
+
+
 def default_is_muon(name, x, expert_keys=None):
     normalized = normalize_fqn(name)
     parts = normalized.split(".")

torch-ext/optimizer/distributed/utils.py

@@ -205,7 +205,11 @@ def construct_shard_mesh(
     else:
         shard_meshes = [sorted_mesh]
     shard_placements = sorted_placements[num_rep:]
-    assert len(shard_placements) == len(set(shard_placements))
+    assert len(shard_placements) == len(set(shard_placements)), (
+        f"Duplicate shard placements {shard_placements}: plain nested "
+        f"same-dim sharding (e.g. [Shard(0), Shard(0)]) is ambiguous and "
+        f"unsupported. Real FSDP2 + TP layouts use _StridedShard for the "
+        f"nested dim, which is supported.")
 
     # -- Step 4: Create/retrieve ProcessGroup for current rank's sub-mesh. --
     # Each rank only creates the group it belongs to, using

torch-ext/optimizer/muon.py

@@ -11,14 +11,17 @@ from torch.profiler import record_function
 from .adamw import _placement_cache, _tensor_cache, step_adamw
 from .async_utils import run_pipeline
 from .core import (_muon_state, adjust_lr_for_muon, batch_pre_ortho,
-                   get_default_muon_param_groups, is_expert_param, update_p)
+                   get_default_muon_param_groups, is_expert_param,
+                   resolve_head_dim, update_p)
 from .cpu_offload import CPUOffloadPool
 from .distributed.utils import (_is_shard, construct_shard_mesh,
                                 get_slices_of_dtensor)
 from .newton_schulz import (COMM_DTYPE, DEFAULT_CHUNK_SIZE_RATIO,
                             _zeropower_via_newtonschulz5,
+                            _zeropower_via_newtonschulz5_per_head,
                             zeropower_via_newtonschulz5,
-                            zeropower_via_newtonschulz5_batched)
+                            zeropower_via_newtonschulz5_batched,
+                            zeropower_via_newtonschulz5_per_head)
 from .pipeline import muon_chunk_pipeline, prelaunch_first_gather
 from .qk_clip import compute_scales, get_qk_clip_info, qk_clip
 
@@ -190,6 +193,22 @@ class Muon(torch.optim.Optimizer):
                      ``expert_keys=["experts"]`` matches any param whose name
                      contains "experts".  3D+ params not matched by any key
                      will raise an error.
+        head_keys: Per-Head Muon (Kimi K3 tech report, §2.5).  Mapping from
+                   parameter-name keys (component-level matching, same rule
+                   as ``expert_keys``) to the dim-0 row stride of one
+                   attention head.  A matching ``[H * head_dim, in]`` param
+                   has its momentum matrix partitioned along dim 0 into
+                   ``H`` head blocks that are orthogonalized independently,
+                   equalizing the update scale across heads.  The
+                   size-adjusted LR uses the head-block shape.
+                   Example (MHA/GQA): ``head_keys={"q_proj": 128,
+                   "k_proj": 128, "v_proj": 128}``.
+                   For MLA, pass each projection's per-head row stride,
+                   e.g. ``{"wq_b": qk_nope_head_dim + qk_rope_head_dim,
+                   "wkv_b": qk_nope_head_dim + v_head_dim}``.
+                   Only dim-0 (output-side) head layouts are supported —
+                   Q/K/V projections, not o_proj.  Keys must be disjoint
+                   from ``expert_keys``.
     """
 
     def __init__(self,
@@ -207,7 +226,8 @@ class Muon(torch.optim.Optimizer):
                  warmup_step=5,
                  chunk_size=-1,
                  use_distributed_muon=False,
-                 expert_keys=None):
+                 expert_keys=None,
+                 head_keys=None):
         defaults = dict(
             lr=lr,
             weight_decay=weight_decay,
@@ -241,6 +261,16 @@ class Muon(torch.optim.Optimizer):
         self.chunk_size = chunk_size
         self.use_distributed_muon = use_distributed_muon
         self.expert_keys = expert_keys
+        if head_keys:
+            for key, hd in head_keys.items():
+                if not isinstance(hd, int) or hd <= 0:
+                    raise ValueError(f"head_keys[{key!r}] must be a positive "
+                                     f"int (got {hd!r}).")
+                if expert_keys and key in expert_keys:
+                    raise ValueError(
+                        f"Key {key!r} appears in both expert_keys and "
+                        f"head_keys; they must be disjoint.")
+        self.head_keys = head_keys
         self.cpu_offload = False
         self.manual_offload = False
         self._cpu_offload_pool: CPUOffloadPool | None = None
@@ -318,6 +348,7 @@ class Muon(torch.optim.Optimizer):
             worker_rank = shard_mesh_flattened[round_robin].item() % num_ranks
             round_robin = (round_robin + 1) % len(shard_mesh_flattened)
             qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits)
+            head_dim = resolve_head_dim(n, p.shape, self.head_keys)
 
             # Precompute per-rank indices and numels for all-to-all.
             rank_indices: dict[int, tuple] = {}
@@ -342,6 +373,7 @@ class Muon(torch.optim.Optimizer):
                 rank_numels=rank_numels,
                 name=n,
                 qk_clip_state=qk_clip_state,
+                head_dim=head_dim,
             )
 
         return param_to_state, ordered_params
@@ -353,10 +385,15 @@ class Muon(torch.optim.Optimizer):
             if g is None:
                 continue
 
-            u = zeropower_via_newtonschulz5(g.to(COMM_DTYPE),
-                                            steps=group["ns_steps"])
-
-            adjusted_lr = adjust_lr_for_muon(lr, p.shape)
+            head_dim = resolve_head_dim(n, p.shape, self.head_keys)
+            if head_dim is not None:
+                u = zeropower_via_newtonschulz5_per_head(
+                    g.to(COMM_DTYPE), head_dim, steps=group["ns_steps"])
+                adjusted_lr = adjust_lr_for_muon(lr, (head_dim, p.shape[1]))
+            else:
+                u = zeropower_via_newtonschulz5(g.to(COMM_DTYPE),
+                                                steps=group["ns_steps"])
+                adjusted_lr = adjust_lr_for_muon(lr, p.shape)
             update_p(p, u, lr, adjusted_lr, weight_decay)
 
             qk_clip_state = get_qk_clip_info(self.clip_config, n, qk_logits)
@@ -386,6 +423,20 @@ class Muon(torch.optim.Optimizer):
             # Momentum is already applied by _step_muon before this method.
             ns_steps = group["ns_steps"]
 
+            def _ns_full(n, g_full):
+                """Full-matrix or per-head NS depending on head_keys."""
+                head_dim = resolve_head_dim(n, g_full.shape, self.head_keys)
+                if head_dim is not None:
+                    return _zeropower_via_newtonschulz5_per_head(
+                        g_full, head_dim, ns_steps)
+                return _zeropower_via_newtonschulz5(g_full, steps=ns_steps)
+
+            def _adjusted_lr(n, shape):
+                head_dim = resolve_head_dim(n, shape, self.head_keys)
+                if head_dim is not None:
+                    return adjust_lr_for_muon(lr, (head_dim, shape[1]))
+                return adjust_lr_for_muon(lr, shape)
+
             # Separate plain tensors (no communication) from DTensors.
             plain_names, plain_params = [], []
             dtensor_names, dtensor_params = [], []
@@ -401,9 +452,8 @@ class Muon(torch.optim.Optimizer):
 
             # Process plain tensors per-param (no communication).
             for n, p in zip(plain_names, plain_params):
-                u = _zeropower_via_newtonschulz5(p.grad.to(COMM_DTYPE),
-                                                 steps=ns_steps)
-                adjusted_lr = adjust_lr_for_muon(lr, p.shape)
+                u = _ns_full(n, p.grad.to(COMM_DTYPE))
+                adjusted_lr = _adjusted_lr(n, p.shape)
                 update_p(p, u, lr, adjusted_lr, weight_decay)
 
                 qk_clip_state = get_qk_clip_info(self.clip_config, n,
@@ -461,11 +511,10 @@ class Muon(torch.optim.Optimizer):
                 for n, p in zip(single_names, single_params):
                     with record_function("distributed_muon::newton_schulz"):
                         g_full = p.grad.full_tensor().to(COMM_DTYPE)
-                        u_full = _zeropower_via_newtonschulz5(g_full,
-                                                              steps=ns_steps)
+                        u_full = _ns_full(n, g_full)
                         del g_full
                     with record_function("distributed_muon::update"):
-                        adjusted_lr = adjust_lr_for_muon(lr, p.shape)
+                        adjusted_lr = _adjusted_lr(n, p.shape)
                         p._local_tensor.mul_(1 - lr * weight_decay)
                         local_indices = get_slices_of_dtensor(
                             p, rank, shard_mesh, shard_placements)
@@ -551,12 +600,11 @@ class Muon(torch.optim.Optimizer):
                             g_full[indices] = shard.reshape(
                                 g_full[indices].shape)
 
-                        u_full = _zeropower_via_newtonschulz5(g_full,
-                                                              steps=ns_steps)
+                        u_full = _ns_full(n, g_full)
                         del g_full
 
                     with record_function("distributed_muon::update"):
-                        adjusted_lr = adjust_lr_for_muon(lr, p.shape)
+                        adjusted_lr = _adjusted_lr(n, p.shape)
                         p._local_tensor.mul_(1 - lr * weight_decay)
                         local_indices = get_slices_of_dtensor(
                             p, rank, shard_mesh, shard_placements)
@@ -650,6 +698,7 @@ class Muon(torch.optim.Optimizer):
                     name=n,
                     qk_clip_state=get_qk_clip_info(self.clip_config, n,
                                                    qk_logits),
+                    head_dim=cached_state.head_dim,
                 )
 
         return ordered_params, param_to_state, rank, chunk_size

torch-ext/optimizer/newton_schulz.py

@@ -197,6 +197,27 @@ def _zeropower_via_newtonschulz5_batched(G, steps):
     return X
 
 
+@torch.no_grad()
+def _zeropower_via_newtonschulz5_per_head(G, head_dim, steps):
+    """Per-head polar factor for attention projection matrices.
+
+    Implements Per-Head Muon (Kimi K3 tech report, §2.5): instead of
+    orthogonalizing the full ``[H * head_dim, N]`` projection matrix as a
+    single coupled block, the matrix is viewed as ``H`` independent
+    ``[head_dim, N]`` head blocks (dim-0 partition) and each block is
+    orthogonalized separately via the batched Newton-Schulz path.  Each
+    block is normalized by its own Frobenius norm, so heads with large
+    momentum scale no longer dominate the shared update direction and
+    the update scale is equalized across heads.
+    """
+    assert len(G.shape) == 2
+    M, N = G.shape
+    assert M % head_dim == 0, (
+        f"dim-0 size {M} is not divisible by head_dim {head_dim}")
+    X = G.reshape(M // head_dim, head_dim, N)
+    return _zeropower_via_newtonschulz5_batched(X, steps).reshape(M, N)
+
+
 _ns_per_shape: dict[tuple[int, ...], callable] = {}
 _use_compile = True
 
@@ -235,3 +256,17 @@ def zeropower_via_newtonschulz5_batched(G, steps=5):
             })
     torch.compiler.cudagraph_mark_step_begin()
     return _ns_per_shape[key](G, steps).clone()
+
+
+def zeropower_via_newtonschulz5_per_head(G, head_dim, steps=5):
+    """Compile-cached per-head Newton-Schulz (Per-Head Muon).
+
+    See ``_zeropower_via_newtonschulz5_per_head``.  Dispatches to the
+    compile-cached batched kernel keyed by the ``[H, head_dim, N]``
+    view shape.
+    """
+    M, N = G.shape
+    assert M % head_dim == 0, (
+        f"dim-0 size {M} is not divisible by head_dim {head_dim}")
+    X = G.reshape(M // head_dim, head_dim, N)
+    return zeropower_via_newtonschulz5_batched(X, steps=steps).reshape(M, N)

torch-ext/optimizer/pipeline.py

@@ -7,7 +7,8 @@ from torch.distributed.tensor import DTensor
 from torch.profiler import record_function
 
 from .core import _muon_state, adjust_lr_for_muon
-from .newton_schulz import COMM_DTYPE, zeropower_via_newtonschulz5
+from .newton_schulz import (COMM_DTYPE, zeropower_via_newtonschulz5,
+                            zeropower_via_newtonschulz5_per_head)
 from .qk_clip import compute_scales
 
 logger = logging.getLogger(__name__)
@@ -143,15 +144,27 @@ def _compute_ns(
     owned_params: list[DTensor],
     gathered_grads: dict[int, torch.Tensor | None],
     ns_steps: int,
+    param_to_state: dict[int, _muon_state],
 ) -> dict[int, torch.Tensor | None]:
     """Run Newton-Schulz orthogonalization on owned parameters.
 
+    Parameters whose ``_muon_state.head_dim`` is set are orthogonalized
+    per attention-head block (Per-Head Muon) instead of as one full
+    matrix.  The gathered grad is always the full global matrix, so the
+    head partition is well-defined regardless of how the parameter is
+    sharded across ranks.
+
     Returns:
         computed_us: ``{id(p): orthogonalized_update}`` for owned params.
     """
     computed_us: dict[int, torch.Tensor | None] = {}
     for p in owned_params:
-        u = zeropower_via_newtonschulz5(gathered_grads[id(p)], ns_steps)
+        head_dim = param_to_state[id(p)].head_dim
+        if head_dim is not None:
+            u = zeropower_via_newtonschulz5_per_head(gathered_grads[id(p)],
+                                                     head_dim, ns_steps)
+        else:
+            u = zeropower_via_newtonschulz5(gathered_grads[id(p)], ns_steps)
         gathered_grads[id(p)] = None  # free gathered grad
         computed_us[id(p)] = u
     return computed_us
@@ -305,9 +318,15 @@ def _update_params(
 
     # Group params by adjusted_lr so _foreach_add_ can use a single
     # alpha per group (preserves float32 precision for alpha scaling).
+    # Per-head params scale by the head-block shape: each block is the
+    # Newton-Schulz unit, mirroring the per-expert path.
     lr_groups: dict[float, tuple[list, list]] = {}
     for p in params:
-        adjusted_lr = adjust_lr_for_muon(lr, p.shape)
+        head_dim = param_to_state[id(p)].head_dim
+        if head_dim is not None:
+            adjusted_lr = adjust_lr_for_muon(lr, (head_dim, p.shape[1]))
+        else:
+            adjusted_lr = adjust_lr_for_muon(lr, p.shape)
         if adjusted_lr not in lr_groups:
             lr_groups[adjusted_lr] = ([], [])
         lr_groups[adjusted_lr][0].append(p._local_tensor)
@@ -443,7 +462,8 @@ def muon_chunk_pipeline(
 
     # Stage 3: Newton-Schulz orthogonalization.
     with record_function("muon::newton_schulz"):
-        computed_us = _compute_ns(owned_params, gathered_grads, ns_steps)
+        computed_us = _compute_ns(owned_params, gathered_grads, ns_steps,
+                                  param_to_state)
         gathered_grads.clear()
 
     # Stages 4-5: launch async scatter.