From 15e3bf4f5770e7f75357a1cdb52c86b9b724cf46 Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Tue, 1 Sep 2026 19:33:15 +0800 Subject: [PATCH 01/11] [muon] Add per-head Newton-Schulz to muon_update Full-matrix orthogonalization treats every attention head as one coupled block, so heads with larger momentum dominate the shared update direction while smaller-scale heads get insufficiently normalized updates. Kimi K3 (arXiv:2607.24653 5 2.5) and GLM-5 Muon Split (arXiv:2602.15763) both orthogonalize per head instead. With num_heads set, the update for a [num_heads * head_dim, in_features] projection is viewed as [num_heads, head_dim, in_features] and Newton-Schulz runs on that batch, with the existing max(1, m/n)**0.5 scaling applied per head block. Both NS kernels are already batch-safe, so this reuses the path the expert-group branch takes. Kernel only; the metadata plumbing and config surface for #8367 follow separately. Signed-off-by: alanhuangyoo --- deepspeed/runtime/zero/muon/original_muon.py | 38 +++++- tests/unit/runtime/zero/test_per_head_muon.py | 125 ++++++++++++++++++ 2 files changed, 162 insertions(+), 1 deletion(-) create mode 100644 tests/unit/runtime/zero/test_per_head_muon.py diff --git a/deepspeed/runtime/zero/muon/original_muon.py b/deepspeed/runtime/zero/muon/original_muon.py index 1cbc46392410..31968cc2d083 100644 --- a/deepspeed/runtime/zero/muon/original_muon.py +++ b/deepspeed/runtime/zero/muon/original_muon.py @@ -143,10 +143,46 @@ def zeropower_via_gram_newtonschulz(G, steps: int): @compiler.compile() -def muon_update(grad, momentum, beta=0.95, ns_steps=5, nesterov=True, ns_method="gram", is_expert_group=False): +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) + + +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 diff --git a/tests/unit/runtime/zero/test_per_head_muon.py b/tests/unit/runtime/zero/test_per_head_muon.py new file mode 100644 index 000000000000..c2940819d3e3 --- /dev/null +++ b/tests/unit/runtime/zero/test_per_head_muon.py @@ -0,0 +1,125 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""Per-head Muon: Newton-Schulz on each attention head instead of the whole projection. + +Full-matrix orthogonalization treats every head as one coupled block, so heads with larger +momentum dominate the shared update direction. Kimi K3 (arXiv:2607.24653 §2.5) and GLM-5 +"Muon Split" (arXiv:2602.15763) both orthogonalize per head instead. See #8367. + +CPU-only: these pin the arithmetic, not the accelerator path. +""" + +import pytest +import torch + +from deepspeed.accelerator import get_accelerator +from deepspeed.runtime.zero.muon.original_muon import ( + muon_update, + zeropower_via_gram_newtonschulz, + zeropower_via_newtonschulz5, +) + + +def _ns_tolerance(ns_method): + """A few ulps of whatever dtype the kernel iterates in. + + `gram` runs in fp16 and `newtonschulz5` in bf16 (fp32 where unsupported), and the iteration + amplifies rounding, so batched and unbatched NS agree to a handful of ulps rather than + bitwise. Deriving the bound from the dtype keeps it honest instead of tuned to pass. + """ + if ns_method == "gram": + dtype = torch.float16 if get_accelerator().is_fp16_supported() else torch.float32 + else: + dtype = torch.bfloat16 if get_accelerator().is_bf16_supported() else torch.float32 + eps = torch.finfo(dtype).eps + return dict(rtol=8 * eps, atol=8 * eps) + + +def _norm_rtol(ns_method): + """Scale agreement: a couple of ulps of the compute dtype, and never looser than 1%.""" + if ns_method == "gram": + dtype = torch.float16 if get_accelerator().is_fp16_supported() else torch.float32 + else: + dtype = torch.bfloat16 if get_accelerator().is_bf16_supported() else torch.float32 + return max(1e-2, 2 * torch.finfo(dtype).eps) + + +def _update_only(grad, momentum, beta=0.95, nesterov=True): + """The pre-orthogonalization update muon_update forms, without mutating the caller's tensors.""" + grad, momentum = grad.clone(), momentum.clone() + momentum.lerp_(grad, 1 - beta) + return grad.lerp_(momentum, beta) if nesterov else momentum + + +@pytest.mark.parametrize("ns_method", ["gram", "newtonschulz5"]) +@pytest.mark.parametrize("num_heads,head_dim,in_features", [(4, 8, 32), (2, 16, 32), (8, 4, 64)]) +def test_per_head_matches_orthogonalizing_each_head_alone(ns_method, num_heads, head_dim, in_features): + """The batched path must equal running NS on each head block on its own.""" + torch.manual_seed(0) + out_features = num_heads * head_dim + grad = torch.randn(out_features, in_features) + momentum = torch.randn(out_features, in_features) + + got = muon_update(grad.clone(), momentum.clone(), ns_method=ns_method, num_heads=num_heads) + + update = _update_only(grad, momentum) + ns_fn = zeropower_via_gram_newtonschulz if ns_method == "gram" else zeropower_via_newtonschulz5 + scale = max(1, head_dim / in_features)**0.5 + expected = torch.cat([ns_fn(update[h * head_dim:(h + 1) * head_dim], steps=5) * scale + for h in range(num_heads)]).to(got.dtype) + + assert got.shape == (out_features, in_features) + torch.testing.assert_close(got, expected, **_ns_tolerance(ns_method)) + # Elementwise agreement is ulp-limited, so also pin the overall scale. + torch.testing.assert_close(got.norm(), expected.norm(), rtol=_norm_rtol(ns_method), atol=0.0) + + +@pytest.mark.parametrize("ns_method", ["gram", "newtonschulz5"]) +def test_single_head_reproduces_the_full_matrix_path(ns_method): + """num_heads=1 is the whole projection, so it has to agree with the existing behaviour.""" + torch.manual_seed(0) + grad = torch.randn(16, 32) + momentum = torch.randn(16, 32) + + per_head = muon_update(grad.clone(), momentum.clone(), ns_method=ns_method, num_heads=1) + full = muon_update(grad.clone(), momentum.clone(), ns_method=ns_method) + + # Batched and unbatched NS take the same arithmetic path but not bit-identically in the + # half-precision compute dtype, so compare at that granularity. + torch.testing.assert_close(per_head, full, **_ns_tolerance(ns_method)) + torch.testing.assert_close(per_head.norm(), full.norm(), rtol=_norm_rtol(ns_method), atol=0.0) + + +def test_per_head_differs_from_full_matrix_when_heads_are_unbalanced(): + """The point of the change: one loud head must stop setting the direction for the quiet ones. + + Without this the test would pass even if num_heads were ignored. + """ + torch.manual_seed(0) + num_heads, head_dim, in_features = 4, 8, 32 + grad = torch.randn(num_heads * head_dim, in_features) + grad[:head_dim] *= 100.0 # one head with a far larger gradient scale + momentum = torch.zeros_like(grad) + + per_head = muon_update(grad.clone(), momentum.clone(), num_heads=num_heads) + full = muon_update(grad.clone(), momentum.clone()) + + quiet = slice(head_dim, None) + assert not torch.allclose(per_head[quiet], full[quiet], rtol=1e-2, atol=1e-2) + # Every head should come out with a comparable update scale. + norms = torch.stack([per_head[h * head_dim:(h + 1) * head_dim].norm() for h in range(num_heads)]) + assert norms.max() / norms.min() < 1.5 + + +def test_rejects_shapes_that_do_not_split_into_heads(): + grad = torch.randn(15, 32) + momentum = torch.zeros_like(grad) + + with pytest.raises(ValueError, match="not divisible by num_heads"): + muon_update(grad.clone(), momentum.clone(), num_heads=4) + + conv_like = torch.randn(4, 4, 3, 3) + with pytest.raises(ValueError, match="expects a 2D attention projection"): + muon_update(conv_like.clone(), torch.zeros_like(conv_like), num_heads=4) From 4644933fe07b6dcb4c7055d425ba1c11288820bb Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Tue, 1 Sep 2026 22:11:53 +0800 Subject: [PATCH 02/11] [muon] Tag attention projections with their head count and wire it through Adds the metadata and config half of #8367 on top of the kernel. set_optimizer_flags now also tags muon_num_heads next to use_muon, gated on an opt-in optimizer.params.per_head_muon. Head structure comes from the model config rather than AutoTP, so it does not require AutoTP to be enabled: q/o projections are blocked by num_attention_heads, k/v by num_key_value_heads, which differ under GQA. Deliberately conservative about what it claims to recognize. A fused QKV matrix is left on the full-matrix path - its three sections split separately, and under GQA they do not even share a head count - and any projection whose output dim does not divide by the head count is skipped with a warning rather than reshaped on a guess. All six muon_update call sites pass the tag through. Each one already operates on a whole parameter rather than a flat shard: the ZeRO-1/2 path views the momentum back to tensor.size() and asserts ndim > 1, and the ZeRO-3 path takes param.grad directly. Signed-off-by: alanhuangyoo --- deepspeed/__init__.py | 43 ++++++++ deepspeed/runtime/zero/muon/original_muon.py | 12 ++- deepspeed/runtime/zero/stage3.py | 6 +- deepspeed/runtime/zero/stage_1_and_2.py | 3 +- .../zero/test_per_head_muon_tagging.py | 98 +++++++++++++++++++ 5 files changed, 156 insertions(+), 6 deletions(-) create mode 100644 tests/unit/runtime/zero/test_per_head_muon_tagging.py diff --git a/deepspeed/__init__.py b/deepspeed/__init__.py index 26bfc7c77df6..f7dbfb28dfcd 100755 --- a/deepspeed/__init__.py +++ b/deepspeed/__init__.py @@ -80,15 +80,58 @@ def _parse_version(version_str): # Set to torch's distributed package or deepspeed.comm based inside DeepSpeedEngine init dist = None +# Attention projections Muon can orthogonalize per head, and which head count each one is +# blocked by. Q and the output projection are blocked by the query heads; K and V by the +# key/value heads, which differ under GQA/MQA. +_QUERY_HEAD_PROJECTIONS = ("q_proj", "query", "o_proj", "out_proj", "dense") +_KV_HEAD_PROJECTIONS = ("k_proj", "key", "v_proj", "value") +# A single matrix holding Q, K and V. Splitting it per head means splitting each of the three +# sections separately, and under GQA they do not even have the same head count, so leave it to +# the full-matrix path rather than guessing the layout. +_FUSED_QKV_PROJECTIONS = ("qkv_proj", "query_key_value", "c_attn", "in_proj_qkv", "Wqkv") + + +def _attention_head_count(param_name: str, model: torch.nn.Module): + """Heads this projection splits into, or None to leave it on the full-matrix path.""" + config = getattr(model, "config", None) + if config is None: + return None + + text_config = getattr(config, "text_config", config) + num_attention_heads = getattr(text_config, "num_attention_heads", None) + if num_attention_heads is None: + return None + num_kv_heads = getattr(text_config, "num_key_value_heads", None) or num_attention_heads + + leaf = param_name.lower() + if any(k.lower() in leaf for k in _FUSED_QKV_PROJECTIONS): + return None + if any(k.lower() in leaf for k in _KV_HEAD_PROJECTIONS): + return num_kv_heads + if any(k.lower() in leaf for k in _QUERY_HEAD_PROJECTIONS): + return num_attention_heads + + return None + def set_optimizer_flags(config_class: DeepSpeedConfig, model: torch.nn.Module) -> None: if config_class.optimizer_name == MUON_OPTIMIZER: + per_head = bool((config_class.optimizer_params or {}).get("per_head_muon", False)) for name, p in model.named_parameters(): if p.ndim >= 2 and not any(keyword in name.lower() for keyword in ("embed", "lm_head")): setattr(p, "use_muon", True) else: setattr(p, "use_muon", False) + num_heads = _attention_head_count(name, model) if (per_head and p.use_muon) else None + # Only tag when the parameter actually splits evenly; a projection whose output dim + # does not divide by the head count is not the layout we think it is. + if num_heads is not None and (p.ndim != 2 or p.shape[0] % num_heads != 0): + logger.warning(f"per_head_muon: skipping {name} with shape {tuple(p.shape)}, which does not " + f"split into {num_heads} heads. It keeps the full-matrix update.") + num_heads = None + setattr(p, "muon_num_heads", num_heads) + def initialize( args: Any = None, diff --git a/deepspeed/runtime/zero/muon/original_muon.py b/deepspeed/runtime/zero/muon/original_muon.py index 31968cc2d083..901666a5f074 100644 --- a/deepspeed/runtime/zero/muon/original_muon.py +++ b/deepspeed/runtime/zero/muon/original_muon.py @@ -255,7 +255,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()], @@ -293,7 +294,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"]) @@ -385,7 +387,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()], @@ -457,7 +460,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: diff --git a/deepspeed/runtime/zero/stage3.py b/deepspeed/runtime/zero/stage3.py index 0afd8c1b9f89..1fe03674666d 100644 --- a/deepspeed/runtime/zero/stage3.py +++ b/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], diff --git a/deepspeed/runtime/zero/stage_1_and_2.py b/deepspeed/runtime/zero/stage_1_and_2.py index f05a53867c93..52cd23191c69 100644 --- a/deepspeed/runtime/zero/stage_1_and_2.py +++ b/deepspeed/runtime/zero/stage_1_and_2.py @@ -2168,7 +2168,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 diff --git a/tests/unit/runtime/zero/test_per_head_muon_tagging.py b/tests/unit/runtime/zero/test_per_head_muon_tagging.py new file mode 100644 index 000000000000..c19788aaa4b0 --- /dev/null +++ b/tests/unit/runtime/zero/test_per_head_muon_tagging.py @@ -0,0 +1,98 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""Which parameters per-head Muon tags, and with how many heads. See #8367. + +`set_optimizer_flags` already tags `use_muon` per parameter; head structure rides along the +same way so it does not depend on AutoTP being enabled. CPU-only. +""" + +from types import SimpleNamespace + +import pytest +import torch + +import deepspeed +from deepspeed.runtime.config import MUON_OPTIMIZER + + +class _Attn(torch.nn.Module): + + def __init__(self, hidden=64, q_heads=8, kv_heads=2, head_dim=8, fused=False): + super().__init__() + self.q_proj = torch.nn.Linear(hidden, q_heads * head_dim, bias=False) + self.k_proj = torch.nn.Linear(hidden, kv_heads * head_dim, bias=False) + self.v_proj = torch.nn.Linear(hidden, kv_heads * head_dim, bias=False) + self.o_proj = torch.nn.Linear(q_heads * head_dim, hidden, bias=False) + self.mlp = torch.nn.Linear(hidden, hidden, bias=False) + self.embed_tokens = torch.nn.Embedding(16, hidden) + if fused: + self.qkv_proj = torch.nn.Linear(hidden, (q_heads + 2 * kv_heads) * head_dim, bias=False) + self.config = SimpleNamespace(num_attention_heads=q_heads, num_key_value_heads=kv_heads) + + +def _flags(model, per_head=True): + cfg = SimpleNamespace(optimizer_name=MUON_OPTIMIZER, + optimizer_params={"per_head_muon": per_head} if per_head else {}) + deepspeed.set_optimizer_flags(cfg, model) + return {name: getattr(p, "muon_num_heads", "MISSING") for name, p in model.named_parameters()} + + +def test_query_and_output_projections_use_the_query_head_count(): + tags = _flags(_Attn(q_heads=8, kv_heads=2)) + + assert tags["q_proj.weight"] == 8 + assert tags["o_proj.weight"] == 8 + + +def test_kv_projections_use_the_kv_head_count_under_gqa(): + """K/V have fewer heads than Q under GQA, and splitting them by the query count would be wrong.""" + tags = _flags(_Attn(q_heads=8, kv_heads=2)) + + assert tags["k_proj.weight"] == 2 + assert tags["v_proj.weight"] == 2 + + +def test_non_attention_parameters_are_left_on_the_full_matrix_path(): + tags = _flags(_Attn()) + + assert tags["mlp.weight"] is None + assert tags["embed_tokens.weight"] is None + + +def test_fused_qkv_is_skipped(): + """One matrix holding Q, K and V does not split into uniform heads under GQA.""" + tags = _flags(_Attn(fused=True)) + + assert tags["qkv_proj.weight"] is None + + +def test_opt_in_is_required(): + tags = _flags(_Attn(), per_head=False) + + assert all(v is None for v in tags.values()), tags + + +def test_shape_that_does_not_divide_is_skipped(): + """A projection whose output dim is not a multiple of the head count is not that layout.""" + model = _Attn(q_heads=8, kv_heads=2) + model.q_proj = torch.nn.Linear(64, 63, bias=False) # 63 % 8 != 0 + + assert _flags(model)["q_proj.weight"] is None + + +def test_use_muon_tagging_is_unchanged(): + model = _Attn() + _flags(model) + + assert model.q_proj.weight.use_muon is True + assert model.embed_tokens.weight.use_muon is False + + +@pytest.mark.parametrize("q_heads,kv_heads", [(8, 8), (8, 1), (12, 4)]) +def test_head_counts_track_the_config(q_heads, kv_heads): + tags = _flags(_Attn(q_heads=q_heads, kv_heads=kv_heads, hidden=64, head_dim=8)) + + assert tags["q_proj.weight"] == q_heads + assert tags["k_proj.weight"] == kv_heads From 11876c644ce511263ec16c99940193c615d42b89 Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Wed, 2 Sep 2026 01:53:10 +0800 Subject: [PATCH 03/11] [muon] Only tag projections whose output dim is head-blocked Two mistakes in the previous commit's tagging, both of which produced a wrong update rather than an error. o_proj / out_proj were tagged with the query head count, but their head structure is on the input dimension ([hidden, num_heads * head_dim]) while the split is on dim 0. With the usual hidden == num_heads * head_dim they still divide evenly, so the matrix was silently cut across the wrong axis. Q/K/V only now. 'dense' was matched anywhere in the parameter path, which also names MLP matrices - intermediate.dense, output.dense, dense_h_to_4h, dense_4h_to_h - so a matrix with no head structure at all was split by the head count. Matching is now on the leaf module name against explicit Q/K/V names, and the shape has to confirm the layout: dim 0 divisible by the head count, and equal to num_heads * head_dim wherever the config states head_dim. Regression tests cover both; against the previous logic all five of the names they pin come back tagged. Signed-off-by: alanhuangyoo --- deepspeed/__init__.py | 71 ++++++++++++------- .../zero/test_per_head_muon_tagging.py | 36 +++++++++- 2 files changed, 78 insertions(+), 29 deletions(-) diff --git a/deepspeed/__init__.py b/deepspeed/__init__.py index f7dbfb28dfcd..fe73b07a4f5e 100755 --- a/deepspeed/__init__.py +++ b/deepspeed/__init__.py @@ -80,19 +80,32 @@ def _parse_version(version_str): # Set to torch's distributed package or deepspeed.comm based inside DeepSpeedEngine init dist = None -# Attention projections Muon can orthogonalize per head, and which head count each one is -# blocked by. Q and the output projection are blocked by the query heads; K and V by the -# key/value heads, which differ under GQA/MQA. -_QUERY_HEAD_PROJECTIONS = ("q_proj", "query", "o_proj", "out_proj", "dense") -_KV_HEAD_PROJECTIONS = ("k_proj", "key", "v_proj", "value") -# A single matrix holding Q, K and V. Splitting it per head means splitting each of the three -# sections separately, and under GQA they do not even have the same head count, so leave it to -# the full-matrix path rather than guessing the layout. -_FUSED_QKV_PROJECTIONS = ("qkv_proj", "query_key_value", "c_attn", "in_proj_qkv", "Wqkv") - - -def _attention_head_count(param_name: str, model: torch.nn.Module): - """Heads this projection splits into, or None to leave it on the full-matrix path.""" +# Only projections whose *output* dimension is blocked by heads, because the per-head split is on +# dim 0 of the weight: Q/K/V are `[num_heads * head_dim, hidden]`. The output projection is +# deliberately absent - its head structure is on the input dimension +# (`[hidden, num_heads * head_dim]`), so splitting dim 0 would cut across the wrong axis, and in +# the usual hidden == num_heads * head_dim case it still divides evenly, so it would be silently +# wrong rather than an error. +_QUERY_HEAD_LEAVES = ("q_proj", "query", "wq") +_KV_HEAD_LEAVES = ("k_proj", "key", "wk", "v_proj", "value", "wv") +# A single matrix holding Q, K and V. Its three sections split separately, and under GQA they do +# not even share a head count, so leave it on the full-matrix path rather than guessing. +_FUSED_QKV_LEAVES = ("qkv_proj", "query_key_value", "c_attn", "in_proj_qkv", "wqkv") + + +def _leaf_module_name(param_name: str) -> str: + """`model.layers.0.self_attn.q_proj.weight` -> `q_proj`. + + Matching the leaf rather than the whole path keeps generic names from matching by accident: + `dense` appears in both `attention.output.dense` and `intermediate.dense`, and an MLP matrix + tagged with a head count would be split on a dimension that has no heads in it. + """ + parts = param_name.split(".") + return parts[-2].lower() if len(parts) >= 2 else parts[-1].lower() + + +def _attention_head_count(param_name: str, param: torch.Tensor, model: torch.nn.Module): + """Heads this projection splits into on dim 0, or None to leave it on the full-matrix path.""" config = getattr(model, "config", None) if config is None: return None @@ -103,15 +116,25 @@ def _attention_head_count(param_name: str, model: torch.nn.Module): return None num_kv_heads = getattr(text_config, "num_key_value_heads", None) or num_attention_heads - leaf = param_name.lower() - if any(k.lower() in leaf for k in _FUSED_QKV_PROJECTIONS): + leaf = _leaf_module_name(param_name) + if any(leaf.startswith(k) for k in _FUSED_QKV_LEAVES): + return None + if any(leaf.startswith(k) for k in _KV_HEAD_LEAVES): + num_heads = num_kv_heads + elif any(leaf.startswith(k) for k in _QUERY_HEAD_LEAVES): + num_heads = num_attention_heads + else: + return None + + # The name says attention; only the shape can confirm the layout. dim 0 has to be the + # head-blocked one, and where the config states head_dim it has to agree. + if param.ndim != 2 or param.shape[0] % num_heads != 0: + return None + config_head_dim = getattr(text_config, "head_dim", None) + if config_head_dim is not None and param.shape[0] != num_heads * config_head_dim: return None - if any(k.lower() in leaf for k in _KV_HEAD_PROJECTIONS): - return num_kv_heads - if any(k.lower() in leaf for k in _QUERY_HEAD_PROJECTIONS): - return num_attention_heads - return None + return num_heads def set_optimizer_flags(config_class: DeepSpeedConfig, model: torch.nn.Module) -> None: @@ -123,13 +146,7 @@ def set_optimizer_flags(config_class: DeepSpeedConfig, model: torch.nn.Module) - else: setattr(p, "use_muon", False) - num_heads = _attention_head_count(name, model) if (per_head and p.use_muon) else None - # Only tag when the parameter actually splits evenly; a projection whose output dim - # does not divide by the head count is not the layout we think it is. - if num_heads is not None and (p.ndim != 2 or p.shape[0] % num_heads != 0): - logger.warning(f"per_head_muon: skipping {name} with shape {tuple(p.shape)}, which does not " - f"split into {num_heads} heads. It keeps the full-matrix update.") - num_heads = None + num_heads = _attention_head_count(name, p, model) if (per_head and p.use_muon) else None setattr(p, "muon_num_heads", num_heads) diff --git a/tests/unit/runtime/zero/test_per_head_muon_tagging.py b/tests/unit/runtime/zero/test_per_head_muon_tagging.py index c19788aaa4b0..1b64319c9c5f 100644 --- a/tests/unit/runtime/zero/test_per_head_muon_tagging.py +++ b/tests/unit/runtime/zero/test_per_head_muon_tagging.py @@ -39,11 +39,43 @@ def _flags(model, per_head=True): return {name: getattr(p, "muon_num_heads", "MISSING") for name, p in model.named_parameters()} -def test_query_and_output_projections_use_the_query_head_count(): +def test_query_projection_uses_the_query_head_count(): tags = _flags(_Attn(q_heads=8, kv_heads=2)) assert tags["q_proj.weight"] == 8 - assert tags["o_proj.weight"] == 8 + + +def test_output_projection_is_left_alone(): + """o_proj is `[hidden, num_heads * head_dim]` - its heads are on the input axis. + + The split is on dim 0, so tagging it would cut across the wrong axis, and with the usual + hidden == num_heads * head_dim it still divides evenly, i.e. silently wrong rather than an + error. Regression test: it was tagged in the first version of this. + """ + tags = _flags(_Attn(q_heads=8, kv_heads=2)) + + assert tags["o_proj.weight"] is None + + +@pytest.mark.parametrize("mlp_name", [ + "intermediate.dense.weight", + "output.dense.weight", + "mlp.dense_h_to_4h.weight", + "mlp.dense_4h_to_h.weight", +]) +def test_mlp_matrices_named_dense_are_not_treated_as_attention(mlp_name): + """`dense` names an MLP matrix as often as an attention one. + + Matching it anywhere in the path tagged `intermediate.dense` and `dense_h_to_4h` with a head + count, splitting a matrix that has no head structure. Regression test: the first version of + this matched on the full path and did exactly that. + """ + from deepspeed import _attention_head_count + + model = _Attn(q_heads=8, kv_heads=2) + weight = torch.zeros(4 * 64, 64) + + assert _attention_head_count(f"encoder.layer.0.{mlp_name}", weight, model) is None def test_kv_projections_use_the_kv_head_count_under_gqa(): From e9584d3f7b5dd770180db78f4633e7ecfa3d932a Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Wed, 2 Sep 2026 02:44:28 +0800 Subject: [PATCH 04/11] [muon] Cover the tagging against real model architectures The synthetic module the other cases use has the leaf names I chose, which is circular for a change whose whole job is recognizing real ones. These build actual HF configs instead. llama / qwen2 / mistral (split QKV, GQA): q_proj tagged with the query head count, k_proj and v_proj with the kv count, o_proj and the MLP projections left alone. gpt_neox / falcon (fused QKV): nothing tagged. These name their output projection 'dense' and their MLP matrices 'dense_h_to_4h' / 'dense_4h_to_h', which is exactly what the previous substring matching got wrong - all three came back tagged. Signed-off-by: alanhuangyoo --- .../zero/test_per_head_muon_tagging.py | 47 +++++++++++++++++++ 1 file changed, 47 insertions(+) diff --git a/tests/unit/runtime/zero/test_per_head_muon_tagging.py b/tests/unit/runtime/zero/test_per_head_muon_tagging.py index 1b64319c9c5f..7548ece621cc 100644 --- a/tests/unit/runtime/zero/test_per_head_muon_tagging.py +++ b/tests/unit/runtime/zero/test_per_head_muon_tagging.py @@ -14,6 +14,7 @@ import torch import deepspeed +from deepspeed import _attention_head_count from deepspeed.runtime.config import MUON_OPTIMIZER @@ -128,3 +129,49 @@ def test_head_counts_track_the_config(q_heads, kv_heads): assert tags["q_proj.weight"] == q_heads assert tags["k_proj.weight"] == kv_heads + + +@pytest.mark.parametrize("arch", ["llama", "qwen2", "mistral"]) +def test_split_qkv_architectures_tag_only_qkv(arch): + """Real HF configs rather than a stand-in, so the leaf names are the ones models actually use.""" + transformers = pytest.importorskip("transformers") + cfg_cls = { + "llama": transformers.LlamaConfig, + "qwen2": transformers.Qwen2Config, + "mistral": transformers.MistralConfig, + }[arch] + cfg = cfg_cls(hidden_size=64, + num_attention_heads=8, + num_key_value_heads=2, + num_hidden_layers=1, + intermediate_size=128, + vocab_size=32) + model = transformers.AutoModelForCausalLM.from_config(cfg) + + tags = {n.split(".")[-2]: _attention_head_count(n, p, model) for n, p in model.named_parameters() if p.ndim == 2} + + assert tags["q_proj"] == 8 + assert tags["k_proj"] == 2, "GQA: k/v are blocked by num_key_value_heads, not the query count" + assert tags["v_proj"] == 2 + assert tags["o_proj"] is None, "o_proj's heads are on the input axis" + for mlp_leaf in ("gate_proj", "up_proj", "down_proj"): + assert tags[mlp_leaf] is None, f"{mlp_leaf} has no head structure" + + +@pytest.mark.parametrize("arch", ["gpt_neox", "falcon"]) +def test_fused_qkv_architectures_tag_nothing(arch): + """These name their MLP matrices `dense_h_to_4h` / `dense_4h_to_h` and their output proj `dense`. + + Matching `dense` anywhere in the path tagged all three; this pins that none of them are. + """ + transformers = pytest.importorskip("transformers") + cfg_cls = {"gpt_neox": transformers.GPTNeoXConfig, "falcon": transformers.FalconConfig}[arch] + kwargs = dict(hidden_size=64, num_attention_heads=8, num_hidden_layers=1, vocab_size=32) + if arch == "gpt_neox": + kwargs["intermediate_size"] = 128 + model = transformers.AutoModelForCausalLM.from_config(cfg_cls(**kwargs)) + + tags = {n: _attention_head_count(n, p, model) for n, p in model.named_parameters() if p.ndim == 2} + + assert all(v is None for v in tags.values()), \ + {k: v for k, v in tags.items() if v is not None} From 870ee004d3c603af4aab41fd36c54dde750934e8 Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Thu, 3 Sep 2026 14:21:41 +0800 Subject: [PATCH 05/11] [muon] Support MLA up-projections and read head counts from AutoTPMeta Addresses the review on #8384. MLA blocks its two up-projections by head instead of using q/k/v: q_b_proj is [num_heads * (qk_nope + qk_rope), rank] and kv_b_proj is [num_heads * (qk_nope + v_head_dim), rank]. That is the split GLM-5's Muon Split applies, and neither width is head_dim, so the previous shape check rejected both. The down-projections (q_a_proj, kv_a_proj_with_mqa) mix latent and rope components and stay on the full-matrix path. Verified against the real tensor shapes of the two models named in #8367: inference-optimization/GLM-5.2-0.8B-A0.8B tags q_b_proj (4096, 512) and kv_b_proj (5120, 128) with 16 heads, matching 16*(192+64) and 16*(192+128), and leaves the down-projections, o_proj and the MLP alone. inference-optimization/Kimi-K3-0.40B is kimi_linear rather than MLA - q_proj is [256, 1024] against 8 heads of 74 - so the shape check keeps it off the per-head path. Head counts now come from AutoTPMeta.from_model_config, the repo's stated single source of truth, which descends into text_config and probes the several spellings models use (num_heads, n_head, attention_heads) instead of one hardcoded name. Adds end-to-end training over ZeRO 1/2/3 at world size 2: tags survive deepspeed.initialize into the ZeRO call sites, per_head_muon stays off by default, and both paths train. Signed-off-by: alanhuangyoo --- deepspeed/__init__.py | 72 +++++++--- tests/unit/ops/muon/test_per_head_muon_e2e.py | 123 ++++++++++++++++++ .../zero/test_per_head_muon_tagging.py | 73 +++++++++++ 3 files changed, 247 insertions(+), 21 deletions(-) create mode 100644 tests/unit/ops/muon/test_per_head_muon_e2e.py diff --git a/deepspeed/__init__.py b/deepspeed/__init__.py index fe73b07a4f5e..044f1bd8e13c 100755 --- a/deepspeed/__init__.py +++ b/deepspeed/__init__.py @@ -80,17 +80,22 @@ def _parse_version(version_str): # Set to torch's distributed package or deepspeed.comm based inside DeepSpeedEngine init dist = None -# Only projections whose *output* dimension is blocked by heads, because the per-head split is on -# dim 0 of the weight: Q/K/V are `[num_heads * head_dim, hidden]`. The output projection is -# deliberately absent - its head structure is on the input dimension -# (`[hidden, num_heads * head_dim]`), so splitting dim 0 would cut across the wrong axis, and in -# the usual hidden == num_heads * head_dim case it still divides evenly, so it would be silently -# wrong rather than an error. +# Projections whose *output* dimension is blocked by heads, because the per-head split is on +# dim 0 of the weight. Standard attention blocks Q/K/V as `[num_heads * head_dim, hidden]`. MLA +# blocks its two up-projections instead: `q_b_proj` is `[num_heads * (qk_nope + qk_rope), rank]` +# and `kv_b_proj` is `[num_heads * (qk_nope + v_head_dim), rank]`, which is the split GLM-5's +# "Muon Split" applies. The output projection is deliberately absent everywhere: its head +# structure is on the input dimension, so splitting dim 0 would cut across the wrong axis, and +# with the usual hidden == num_heads * head_dim it still divides evenly, i.e. silently wrong. _QUERY_HEAD_LEAVES = ("q_proj", "query", "wq") _KV_HEAD_LEAVES = ("k_proj", "key", "wk", "v_proj", "value", "wv") -# A single matrix holding Q, K and V. Its three sections split separately, and under GQA they do -# not even share a head count, so leave it on the full-matrix path rather than guessing. +# MLA up-projections, keyed by which per-head width the config gives them. +_MLA_Q_LEAVES = ("q_b_proj", ) +_MLA_KV_LEAVES = ("kv_b_proj", ) +# A single matrix holding Q, K and V, or an MLA down-projection that mixes latent and rope +# components. Neither splits into uniform heads, so leave them on the full-matrix path. _FUSED_QKV_LEAVES = ("qkv_proj", "query_key_value", "c_attn", "in_proj_qkv", "wqkv") +_MLA_DOWN_LEAVES = ("q_a_proj", "kv_a_proj", "kv_a_proj_with_mqa") def _leaf_module_name(param_name: str) -> str: @@ -104,34 +109,59 @@ def _leaf_module_name(param_name: str) -> str: return parts[-2].lower() if len(parts) >= 2 else parts[-1].lower() +def _mla_head_width(text_config, leaf: str): + """Per-head output width of an MLA up-projection, or None if this is not one.""" + qk_nope = getattr(text_config, "qk_nope_head_dim", None) + qk_rope = getattr(text_config, "qk_rope_head_dim", None) + v_head_dim = getattr(text_config, "v_head_dim", None) + if any(leaf.startswith(k) for k in _MLA_Q_LEAVES): + if qk_nope is not None and qk_rope is not None: + return qk_nope + qk_rope + elif any(leaf.startswith(k) for k in _MLA_KV_LEAVES): + if qk_nope is not None and v_head_dim is not None: + return qk_nope + v_head_dim + return None + + def _attention_head_count(param_name: str, param: torch.Tensor, model: torch.nn.Module): """Heads this projection splits into on dim 0, or None to leave it on the full-matrix path.""" - config = getattr(model, "config", None) - if config is None: + model_config = getattr(model, "config", None) + if model_config is None: return None - text_config = getattr(config, "text_config", config) - num_attention_heads = getattr(text_config, "num_attention_heads", None) + # AutoTPMeta.from_model_config is the repo's single source of truth for these counts: it + # descends into text_config and probes the several spellings models use (num_heads, n_head, + # attention_heads, ...) rather than assuming one attribute name. + from .module_inject.tp_shard import AutoTPMeta + + meta = AutoTPMeta.from_model_config(model_config) + num_attention_heads = meta.num_attention_heads if num_attention_heads is None: return None - num_kv_heads = getattr(text_config, "num_key_value_heads", None) or num_attention_heads + num_kv_heads = meta.num_kv_heads or num_attention_heads + text_config = getattr(model_config, "text_config", model_config) leaf = _leaf_module_name(param_name) - if any(leaf.startswith(k) for k in _FUSED_QKV_LEAVES): + if any(leaf.startswith(k) for k in _FUSED_QKV_LEAVES) or any(leaf.startswith(k) for k in _MLA_DOWN_LEAVES): return None - if any(leaf.startswith(k) for k in _KV_HEAD_LEAVES): - num_heads = num_kv_heads + + mla_width = _mla_head_width(text_config, leaf) + if mla_width is not None: + # MLA blocks both up-projections by the query head count; the two differ in per-head + # width, not in how many heads they carry. + num_heads, head_width = num_attention_heads, mla_width + elif any(leaf.startswith(k) for k in _KV_HEAD_LEAVES): + num_heads, head_width = num_kv_heads, getattr(text_config, "head_dim", None) elif any(leaf.startswith(k) for k in _QUERY_HEAD_LEAVES): - num_heads = num_attention_heads + num_heads, head_width = num_attention_heads, getattr(text_config, "head_dim", None) else: return None # The name says attention; only the shape can confirm the layout. dim 0 has to be the - # head-blocked one, and where the config states head_dim it has to agree. - if param.ndim != 2 or param.shape[0] % num_heads != 0: + # head-blocked one, and where a per-head width is known it has to agree exactly. + if param.ndim != 2 or num_heads is None or num_heads < 1 or param.shape[0] % num_heads != 0: return None - config_head_dim = getattr(text_config, "head_dim", None) - if config_head_dim is not None and param.shape[0] != num_heads * config_head_dim: + if head_width is not None and param.shape[0] != num_heads * head_width: return None return num_heads diff --git a/tests/unit/ops/muon/test_per_head_muon_e2e.py b/tests/unit/ops/muon/test_per_head_muon_e2e.py new file mode 100644 index 000000000000..c7239d0e7afa --- /dev/null +++ b/tests/unit/ops/muon/test_per_head_muon_e2e.py @@ -0,0 +1,123 @@ +# Copyright (c) Microsoft Corporation. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""End-to-end training with per-head Muon, across ZeRO stages and world size > 1. + +The unit tests next to this pin the arithmetic and the tagging. These run the whole path: +`deepspeed.initialize` tags the parameters, the ZeRO call sites carry the tag into +`muon_update`, and a real training loop takes steps with it. See #8367. +""" + +from types import SimpleNamespace + +import pytest +import torch + +import deepspeed +from unit.common import DistributedTest + + +class AttentionModel(torch.nn.Module): + """Small GQA-shaped model: split QKV, 8 query heads over 2 kv heads.""" + + def __init__(self, hidden_dim=64, q_heads=8, kv_heads=2, head_dim=8, nlayers=2): + super().__init__() + self.q_heads, self.kv_heads, self.head_dim = q_heads, kv_heads, head_dim + self.blocks = torch.nn.ModuleList() + for _ in range(nlayers): + self.blocks.append( + torch.nn.ModuleDict({ + "q_proj": torch.nn.Linear(hidden_dim, q_heads * head_dim, bias=False), + "k_proj": torch.nn.Linear(hidden_dim, kv_heads * head_dim, bias=False), + "v_proj": torch.nn.Linear(hidden_dim, kv_heads * head_dim, bias=False), + "o_proj": torch.nn.Linear(q_heads * head_dim, hidden_dim, bias=False), + "mlp": torch.nn.Linear(hidden_dim, hidden_dim, bias=False), + })) + self.cross_entropy_loss = torch.nn.CrossEntropyLoss() + self.config = SimpleNamespace(num_attention_heads=q_heads, + num_key_value_heads=kv_heads, + hidden_size=hidden_dim, + head_dim=head_dim) + + def forward(self, x, y): + for b in self.blocks: + q, k, v = b["q_proj"](x), b["k_proj"](x), b["v_proj"](x) + rep = self.q_heads // self.kv_heads + attn = q * k.repeat(1, rep) + v.repeat(1, rep) + x = x + b["mlp"](b["o_proj"](attn)) + return self.cross_entropy_loss(x, y) + + +def _config(zero_stage, per_head, lr=0.01): + return { + "train_batch_size": 4, + "optimizer": { + "type": "muon", + "params": { + "lr": lr, + "adam_lr": lr, + "per_head_muon": per_head + } + }, + "zero_optimization": { + "stage": zero_stage, + # Muon does not support reduce-scatter; the existing Muon suite disables it the + # same way (see TestMuonRejectsReduceScatter). + "reduce_scatter": False, + }, + "fp16": { + "enabled": False + }, + "bf16": { + "enabled": True + }, + } + + +def _train(model, config, steps=6, hidden_dim=64, seed=1234): + engine, *_ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config) + tags = {n: getattr(p, "muon_num_heads", "MISSING") for n, p in model.named_parameters()} + gen = torch.Generator().manual_seed(seed) + losses = [] + for _ in range(steps): + x = torch.randn(4, hidden_dim, generator=gen).to(engine.device).to(torch.bfloat16) + y = torch.randint(0, hidden_dim, (4, ), generator=gen).to(engine.device) + loss = engine(x, y) + engine.backward(loss) + engine.step() + losses.append(loss.item()) + return tags, losses + + +@pytest.mark.parametrize("zero_stage", [1, 2, 3]) +class TestPerHeadMuonEndToEnd(DistributedTest): + world_size = 2 + + def test_tags_reach_the_optimizer(self, zero_stage): + """Per-parameter tags have to survive `deepspeed.initialize` into the ZeRO call sites.""" + torch.manual_seed(1234) + tags, losses = _train(AttentionModel(), _config(zero_stage, per_head=True)) + + assert tags["blocks.0.q_proj.weight"] == 8 + assert tags["blocks.0.k_proj.weight"] == 2, "GQA: kv projections carry the kv head count" + assert tags["blocks.0.v_proj.weight"] == 2 + assert tags["blocks.0.o_proj.weight"] is None, "o_proj's heads are on the input axis" + assert tags["blocks.0.mlp.weight"] is None + assert all(torch.isfinite(torch.tensor(loss)) for loss in losses) + + def test_opt_in_is_off_by_default(self, zero_stage): + torch.manual_seed(1234) + tags, _ = _train(AttentionModel(), _config(zero_stage, per_head=False)) + + assert all(v is None for v in tags.values()), {k: v for k, v in tags.items() if v is not None} + + def test_training_makes_progress_either_way(self, zero_stage): + """Both paths have to train; this is the baseline delock asked for alongside per-head.""" + torch.manual_seed(1234) + _, full = _train(AttentionModel(), _config(zero_stage, per_head=False)) + torch.manual_seed(1234) + _, per_head = _train(AttentionModel(), _config(zero_stage, per_head=True)) + + assert full[-1] < full[0], f"baseline did not train: {full}" + assert per_head[-1] < per_head[0], f"per-head did not train: {per_head}" diff --git a/tests/unit/runtime/zero/test_per_head_muon_tagging.py b/tests/unit/runtime/zero/test_per_head_muon_tagging.py index 7548ece621cc..9c85989f76fc 100644 --- a/tests/unit/runtime/zero/test_per_head_muon_tagging.py +++ b/tests/unit/runtime/zero/test_per_head_muon_tagging.py @@ -175,3 +175,76 @@ def test_fused_qkv_architectures_tag_nothing(arch): assert all(v is None for v in tags.values()), \ {k: v for k, v in tags.items() if v is not None} + + +# Shapes and config values read off the real checkpoints delock pointed at in #8367: +# inference-optimization/GLM-5.2-0.8B-A0.8B and inference-optimization/Kimi-K3-0.40B. +def _glm52_mla_config(): + return SimpleNamespace(num_attention_heads=16, + num_key_value_heads=16, + hidden_size=2048, + head_dim=64, + q_lora_rank=512, + kv_lora_rank=128, + qk_nope_head_dim=192, + qk_rope_head_dim=64, + v_head_dim=128) + + +@pytest.mark.parametrize( + "leaf,shape,expected", + [ + ("q_b_proj", (4096, 512), 16), # 16 * (qk_nope 192 + qk_rope 64) + ("kv_b_proj", (5120, 128), 16), # 16 * (qk_nope 192 + v_head_dim 128) + ("q_a_proj", (512, 2048), None), # down-projection, no head structure + ("kv_a_proj_with_mqa", (192, 2048), None), # latent + rope, does not split into heads + ("o_proj", (2048, 2048), None), # heads on the input axis + ]) +def test_mla_tags_only_the_up_projections(leaf, shape, expected): + """MLA blocks its two up-projections by head, which is the split GLM-5's Muon Split applies. + + Their per-head width is not `head_dim`: q_b is qk_nope + qk_rope and kv_b is + qk_nope + v_head_dim, so a tagger that assumes `num_heads * head_dim` rejects both. + """ + model = SimpleNamespace(config=_glm52_mla_config()) + name = f"model.layers.0.self_attn.{leaf}.weight" + + assert _attention_head_count(name, torch.zeros(shape), model) == expected + + +def test_mla_head_width_is_checked_not_assumed(): + """A shape that does not equal num_heads * per-head width is not the layout we think it is.""" + model = SimpleNamespace(config=_glm52_mla_config()) + + ok = _attention_head_count("l.0.self_attn.q_b_proj.weight", torch.zeros(4096, 512), model) + wrong = _attention_head_count("l.0.self_attn.q_b_proj.weight", torch.zeros(4080, 512), model) + + assert ok == 16 + assert wrong is None + + +def test_linear_attention_named_like_standard_attention_is_rejected(): + """Kimi-K3-0.40B is `kimi_linear`, not MLA: q_proj is [256, 1024] with 8 heads of 74. + + The names match the standard-attention list, so only the shape check keeps it off the + per-head path. + """ + text = SimpleNamespace(num_attention_heads=8, + num_key_value_heads=8, + hidden_size=1024, + head_dim=74, + qk_nope_head_dim=64, + qk_rope_head_dim=32, + v_head_dim=64) + model = SimpleNamespace(config=SimpleNamespace(text_config=text)) + + for leaf in ("q_proj", "k_proj", "v_proj"): + name = f"model.layers.0.self_attn.{leaf}.weight" + assert _attention_head_count(name, torch.zeros(256, 1024), model) is None + + +def test_head_count_comes_from_the_shared_extractor(): + """Head counts are read through AutoTPMeta, so alternative config spellings work.""" + model = SimpleNamespace(config=SimpleNamespace(n_head=8, hidden_size=64, head_dim=8)) + + assert _attention_head_count("l.0.attn.q_proj.weight", torch.zeros(64, 64), model) == 8 From c6a7793bd4c73b9d0a46691f85eb7ed4cc9f071f Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Fri, 4 Sep 2026 12:31:35 +0800 Subject: [PATCH 06/11] [muon] Cover MLA models that have no q_lora_rank An MLA attention block only builds the q_a/q_b pair when q_lora_rank is set. Without it the query up-projection is a plain q_proj, and its per-head width is still qk_nope_head_dim + qk_rope_head_dim rather than head_dim: self.q_proj = nn.Linear(hidden_size, num_heads * self.qk_head_dim) if self.q_lora_rank is None else None DeepSeek-V2-Lite is in that shape. Measured by instantiating DeepseekV2Attention on the released deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct config under transformers 5.16.1: heads=16 head_dim=64 qk_nope=128 qk_rope=64 q_lora_rank=None q_proj.weight (3072, 2048) 16 * 192 kv_b_proj.weight (4096, 512) 16 * 256 kv_a_proj_with_mqa.weight (576, 2048) o_proj.weight (2048, 2048) head_dim is present and equal to 64, so the shape check compared 3072 against 16 * 64 = 1024 and put the model back on the full-matrix path. The fallback is the behaviour before this feature, so nothing was wrong, but a whole class of MLA checkpoints never reached the per-head split the feature exists for. q_proj now takes the MLA width when the config carries the MLA head dimensions, and falls through to head_dim when it does not, so ordinary attention models are unchanged. Signed-off-by: alanhuangyoo --- deepspeed/__init__.py | 12 +++-- .../zero/test_per_head_muon_tagging.py | 53 +++++++++++++++++++ 2 files changed, 62 insertions(+), 3 deletions(-) diff --git a/deepspeed/__init__.py b/deepspeed/__init__.py index 044f1bd8e13c..575658facdc1 100755 --- a/deepspeed/__init__.py +++ b/deepspeed/__init__.py @@ -89,8 +89,10 @@ def _parse_version(version_str): # with the usual hidden == num_heads * head_dim it still divides evenly, i.e. silently wrong. _QUERY_HEAD_LEAVES = ("q_proj", "query", "wq") _KV_HEAD_LEAVES = ("k_proj", "key", "wk", "v_proj", "value", "wv") -# MLA up-projections, keyed by which per-head width the config gives them. -_MLA_Q_LEAVES = ("q_b_proj", ) +# MLA up-projections, keyed by which per-head width the config gives them. Without a +# q_lora_rank there is no q_a/q_b pair and the query up-projection is a plain `q_proj` +# (DeepSeek-V2-Lite), still `qk_nope + qk_rope` wide per head rather than `head_dim`. +_MLA_Q_LEAVES = ("q_b_proj", "q_proj") _MLA_KV_LEAVES = ("kv_b_proj", ) # A single matrix holding Q, K and V, or an MLA down-projection that mixes latent and rope # components. Neither splits into uniform heads, so leave them on the full-matrix path. @@ -110,7 +112,11 @@ def _leaf_module_name(param_name: str) -> str: def _mla_head_width(text_config, leaf: str): - """Per-head output width of an MLA up-projection, or None if this is not one.""" + """Per-head output width of an MLA up-projection, or None if this is not one. + + A config without the MLA head dimensions returns None here, so `q_proj` on an ordinary + attention model falls through to the `head_dim` path below. + """ qk_nope = getattr(text_config, "qk_nope_head_dim", None) qk_rope = getattr(text_config, "qk_rope_head_dim", None) v_head_dim = getattr(text_config, "v_head_dim", None) diff --git a/tests/unit/runtime/zero/test_per_head_muon_tagging.py b/tests/unit/runtime/zero/test_per_head_muon_tagging.py index 9c85989f76fc..9701e716cbaf 100644 --- a/tests/unit/runtime/zero/test_per_head_muon_tagging.py +++ b/tests/unit/runtime/zero/test_per_head_muon_tagging.py @@ -223,6 +223,59 @@ def test_mla_head_width_is_checked_not_assumed(): assert wrong is None +# Shapes measured by instantiating DeepseekV2Attention on the released +# deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct config under transformers 5.16.1. +def _deepseek_v2_lite_mla_config(): + return SimpleNamespace(num_attention_heads=16, + num_key_value_heads=16, + hidden_size=2048, + head_dim=64, + q_lora_rank=None, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128) + + +@pytest.mark.parametrize( + "leaf,shape,expected", + [ + ("q_proj", (3072, 2048), 16), # 16 * (qk_nope 128 + qk_rope 64) + ("kv_b_proj", (4096, 512), 16), # 16 * (qk_nope 128 + v_head_dim 128) + ("kv_a_proj_with_mqa", (576, 2048), None), # kv_lora_rank + qk_rope, no head structure + ("o_proj", (2048, 2048), None), # heads on the input axis, and 2048 still divides by 16 + ]) +def test_mla_without_q_lora_rank_tags_the_plain_q_proj(leaf, shape, expected): + """Without a q_lora_rank there is no q_a/q_b pair; the up-projection is `q_proj` itself. + + Its per-head width stays `qk_nope + qk_rope`, so reading `head_dim` gives 16 * 64 = 1024 + against a real 3072 and drops the model off the per-head path. + """ + model = SimpleNamespace(config=_deepseek_v2_lite_mla_config()) + name = f"model.layers.0.self_attn.{leaf}.weight" + + assert _attention_head_count(name, torch.zeros(shape), model) == expected + + +def test_head_dim_alone_would_reject_the_mla_q_proj(): + """Guards the width source rather than the outcome: head_dim is present and wrong here.""" + config = _deepseek_v2_lite_mla_config() + + assert config.head_dim is not None + assert config.num_attention_heads * config.head_dim == 1024 + assert _attention_head_count("l.0.self_attn.q_proj.weight", torch.zeros(3072, 2048), + SimpleNamespace(config=config)) == 16 + + +def test_q_proj_on_a_non_mla_config_still_uses_head_dim(): + """The MLA width only applies where the config carries the MLA head dimensions.""" + config = SimpleNamespace(num_attention_heads=8, num_key_value_heads=8, hidden_size=512, head_dim=64) + model = SimpleNamespace(config=config) + + assert _attention_head_count("l.0.self_attn.q_proj.weight", torch.zeros(512, 512), model) == 8 + assert _attention_head_count("l.0.self_attn.q_proj.weight", torch.zeros(768, 512), model) is None + + def test_linear_attention_named_like_standard_attention_is_rejected(): """Kimi-K3-0.40B is `kimi_linear`, not MLA: q_proj is [256, 1024] with 8 heads of 74. From c3e93341d3954ab7a02297caab9722db4130bdf9 Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Fri, 4 Sep 2026 13:56:07 +0800 Subject: [PATCH 07/11] [muon] Keep muon_update compiled Adding the per-head helper above muon_update left @compiler.compile() attached to the new function instead of to muon_update, so muon_update lost the torch.compile it has on master. That was not intended and it applies to every Muon user, not only the per-head path. The decorated set now matches master again: zeropower_via_newtonschulz5, zeropower_via_gram_newtonschulz and muon_update. _per_head_orthogonalize is called from inside muon_update and does not need its own. Signed-off-by: alanhuangyoo --- deepspeed/runtime/zero/muon/original_muon.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/deepspeed/runtime/zero/muon/original_muon.py b/deepspeed/runtime/zero/muon/original_muon.py index 901666a5f074..71520a42b4bc 100644 --- a/deepspeed/runtime/zero/muon/original_muon.py +++ b/deepspeed/runtime/zero/muon/original_muon.py @@ -142,7 +142,6 @@ def zeropower_via_gram_newtonschulz(G, steps: int): NS_METHODS = {"standard", "gram"} -@compiler.compile() 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: @@ -162,6 +161,7 @@ def _per_head_orthogonalize(update, num_heads, ns_steps, ns_method): return per_head.reshape(out_features, in_features) +@compiler.compile() def muon_update(grad, momentum, beta=0.95, From 37740f6f2930e4ded08a956a8b30a58fda7faeba Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Sun, 6 Sep 2026 16:39:11 +0800 Subject: [PATCH 08/11] [muon] Resolve head geometry by candidate, and report what the flag did Addresses @delock's review. The dispatch was an ordered cascade, so which branch claimed a leaf depended on which config fields happened to exist rather than on the model. q_proj sat in both the MLA and the standard tables; Kimi-K3's KDA q_proj was claimed by the MLA branch on the strength of top-level qk_nope/qk_rope that belong to its two MLA layers, and survived only because 8 x 96 != 256. Safety by width mismatch, not by recognition. Now every geometry the config makes plausible for a leaf is collected and the shape confirms one: classify(leaf) -> kind geometry_candidates(kind, meta, cfg) -> [(heads, width, source)] confirm(param, candidates) -> (num_heads | None, reason) Candidates that confirm and agree on the head count are not a conflict, since the head count is the whole output. Candidates that confirm and disagree are, and the parameter is skipped with both named in the warning. This also removes a weaker rule that was hiding in the old code: with no head_dim the tagger fell back to divisibility alone, and rows % heads == 0 holds for matrices with no head structure at all -- the way o_proj slipped through in the first version. Every candidate now carries a width, derived as hidden_size // num_attention_heads where a config omits head_dim. An explicit opt-in that silently does nothing is its own failure. deepspeed.initialize now raises when per_head_muon is set and no projection could be tagged, listing the likely causes. The systemic one is tensor parallelism: the config describes the whole model while each rank holds a shard, so every width check fails and per-head is off model-wide while the user believes it is on. Non-systemic misses are aggregated per leaf as a warning so a hybrid model keeps per-head on its recognized layers, and an info line reports what was tagged. Also from the review: - AutoTPMeta.from_model_config is built once per model rather than per parameter. The model-taking _attention_head_count is kept for single-parameter callers. - ns_compute_dtype moves into original_muon.py, so the kernels and the test tolerances read the NS precision from one place instead of restating it. - per_head_muon is documented in config-json.md: what is tagged, what deliberately is not, the shape-confirms-the-name contract, and the reporting behaviour. - The three test files carried a Microsoft copyright line; they use the repo's two-line header now. Tests added: a Kimi-K3 hybrid fixture pinning that the KDA projections are declined by candidate confirmation rather than by branch order, a synthetic ambiguity case, agreement-is-not-ambiguity, the divisibility-alone rejection, head_dim derivation, and the initialize-time error. 52 unit, 9 end-to-end on 2 x H20. Signed-off-by: alanhuangyoo --- deepspeed/__init__.py | 207 +++++++++++++----- deepspeed/runtime/zero/muon/original_muon.py | 19 +- docs/_pages/config-json.md | 33 +++ tests/unit/ops/muon/test_per_head_muon_e2e.py | 1 - tests/unit/runtime/zero/test_per_head_muon.py | 15 +- .../zero/test_per_head_muon_tagging.py | 117 +++++++++- 6 files changed, 323 insertions(+), 69 deletions(-) diff --git a/deepspeed/__init__.py b/deepspeed/__init__.py index 575658facdc1..7b64f4c5dcf7 100755 --- a/deepspeed/__init__.py +++ b/deepspeed/__init__.py @@ -89,16 +89,38 @@ def _parse_version(version_str): # with the usual hidden == num_heads * head_dim it still divides evenly, i.e. silently wrong. _QUERY_HEAD_LEAVES = ("q_proj", "query", "wq") _KV_HEAD_LEAVES = ("k_proj", "key", "wk", "v_proj", "value", "wv") -# MLA up-projections, keyed by which per-head width the config gives them. Without a -# q_lora_rank there is no q_a/q_b pair and the query up-projection is a plain `q_proj` -# (DeepSeek-V2-Lite), still `qk_nope + qk_rope` wide per head rather than `head_dim`. -_MLA_Q_LEAVES = ("q_b_proj", "q_proj") +# MLA up-projections. Without a q_lora_rank there is no q_a/q_b pair and the query +# up-projection is a plain `q_proj` (DeepSeek-V2-Lite), so `q_proj` can be either kind and +# is resolved by shape rather than by which config fields happen to exist. +_MLA_Q_LEAVES = ("q_b_proj", ) _MLA_KV_LEAVES = ("kv_b_proj", ) # A single matrix holding Q, K and V, or an MLA down-projection that mixes latent and rope # components. Neither splits into uniform heads, so leave them on the full-matrix path. _FUSED_QKV_LEAVES = ("qkv_proj", "query_key_value", "c_attn", "in_proj_qkv", "wqkv") _MLA_DOWN_LEAVES = ("q_a_proj", "kv_a_proj", "kv_a_proj_with_mqa") +QUERY, KV, MLA_Q, MLA_KV, NOT_HEAD_BLOCKED = "query", "kv", "mla_q", "mla_kv", "not-head-blocked" + + +def _per_head_muon_meta(model: torch.nn.Module): + """The head-count reader and the config the widths come from, built once per model. + + AutoTPMeta.from_model_config is the repo's single source of truth for these counts: it + descends into text_config and probes the several spellings models use (num_heads, n_head, + attention_heads, ...) rather than assuming one attribute name. It is loop-invariant, so it + is built here rather than per parameter. + """ + model_config = getattr(model, "config", None) + if model_config is None: + return None, None + + from .module_inject.tp_shard import AutoTPMeta + + meta = AutoTPMeta.from_model_config(model_config) + if meta.num_attention_heads is None: + return None, None + return meta, getattr(model_config, "text_config", model_config) + def _leaf_module_name(param_name: str) -> str: """`model.layers.0.self_attn.q_proj.weight` -> `q_proj`. @@ -111,80 +133,165 @@ def _leaf_module_name(param_name: str) -> str: return parts[-2].lower() if len(parts) >= 2 else parts[-1].lower() -def _mla_head_width(text_config, leaf: str): - """Per-head output width of an MLA up-projection, or None if this is not one. +def _classify_leaf(leaf: str): + """Which kind of attention matrix this leaf name claims to be, or None if it claims none. - A config without the MLA head dimensions returns None here, so `q_proj` on an ordinary - attention model falls through to the `head_dim` path below. + A name is a claim, not a layout. What the leaf resolves to is decided later, by the shape. """ - qk_nope = getattr(text_config, "qk_nope_head_dim", None) - qk_rope = getattr(text_config, "qk_rope_head_dim", None) - v_head_dim = getattr(text_config, "v_head_dim", None) + if any(leaf.startswith(k) for k in _FUSED_QKV_LEAVES) or any(leaf.startswith(k) for k in _MLA_DOWN_LEAVES): + return NOT_HEAD_BLOCKED if any(leaf.startswith(k) for k in _MLA_Q_LEAVES): - if qk_nope is not None and qk_rope is not None: - return qk_nope + qk_rope - elif any(leaf.startswith(k) for k in _MLA_KV_LEAVES): - if qk_nope is not None and v_head_dim is not None: - return qk_nope + v_head_dim + return MLA_Q + if any(leaf.startswith(k) for k in _MLA_KV_LEAVES): + return MLA_KV + if any(leaf.startswith(k) for k in _KV_HEAD_LEAVES): + return KV + if any(leaf.startswith(k) for k in _QUERY_HEAD_LEAVES): + return QUERY return None -def _attention_head_count(param_name: str, param: torch.Tensor, model: torch.nn.Module): - """Heads this projection splits into on dim 0, or None to leave it on the full-matrix path.""" - model_config = getattr(model, "config", None) - if model_config is None: - return None +def _standard_head_dim(text_config, num_attention_heads: int): + """`head_dim`, or the value it is defined as when a config leaves it out.""" + head_dim = getattr(text_config, "head_dim", None) + if head_dim is not None: + return head_dim + hidden_size = getattr(text_config, "hidden_size", None) + if hidden_size is not None and num_attention_heads: + quotient, remainder = divmod(hidden_size, num_attention_heads) + if remainder == 0: + return quotient + return None - # AutoTPMeta.from_model_config is the repo's single source of truth for these counts: it - # descends into text_config and probes the several spellings models use (num_heads, n_head, - # attention_heads, ...) rather than assuming one attribute name. - from .module_inject.tp_shard import AutoTPMeta - meta = AutoTPMeta.from_model_config(model_config) +def _geometry_candidates(kind, meta, text_config): + """Every (heads, per-head width) this leaf could plausibly have, each labelled. + + Candidates are collected rather than chosen. `q_proj` is the case that matters: on an MLA + model without a q_lora_rank it is the query up-projection, and on an ordinary model it is + the standard query projection, and a config can carry the fields for both. Deciding by the + order the branches are written makes the outcome depend on which fields happen to exist; + collecting both and letting the shape confirm one makes it depend on the model. + """ num_attention_heads = meta.num_attention_heads - if num_attention_heads is None: - return None num_kv_heads = meta.num_kv_heads or num_attention_heads - text_config = getattr(model_config, "text_config", model_config) + head_dim = _standard_head_dim(text_config, num_attention_heads) + qk_nope = getattr(text_config, "qk_nope_head_dim", None) + qk_rope = getattr(text_config, "qk_rope_head_dim", None) + v_head_dim = getattr(text_config, "v_head_dim", None) - leaf = _leaf_module_name(param_name) - if any(leaf.startswith(k) for k in _FUSED_QKV_LEAVES) or any(leaf.startswith(k) for k in _MLA_DOWN_LEAVES): - return None + candidates = [] + if kind in (QUERY, MLA_Q) and qk_nope is not None and qk_rope is not None: + candidates.append((num_attention_heads, qk_nope + qk_rope, "mla-q")) + if kind == MLA_KV and qk_nope is not None and v_head_dim is not None: + candidates.append((num_attention_heads, qk_nope + v_head_dim, "mla-kv")) + if kind == QUERY and head_dim is not None: + candidates.append((num_attention_heads, head_dim, "head-dim")) + if kind == KV and head_dim is not None: + candidates.append((num_kv_heads, head_dim, "head-dim")) + return [c for c in candidates if c[0] and c[0] >= 1 and c[1] and c[1] >= 1] - mla_width = _mla_head_width(text_config, leaf) - if mla_width is not None: - # MLA blocks both up-projections by the query head count; the two differ in per-head - # width, not in how many heads they carry. - num_heads, head_width = num_attention_heads, mla_width - elif any(leaf.startswith(k) for k in _KV_HEAD_LEAVES): - num_heads, head_width = num_kv_heads, getattr(text_config, "head_dim", None) - elif any(leaf.startswith(k) for k in _QUERY_HEAD_LEAVES): - num_heads, head_width = num_attention_heads, getattr(text_config, "head_dim", None) - else: - return None - # The name says attention; only the shape can confirm the layout. dim 0 has to be the - # head-blocked one, and where a per-head width is known it has to agree exactly. - if param.ndim != 2 or num_heads is None or num_heads < 1 or param.shape[0] % num_heads != 0: - return None - if head_width is not None and param.shape[0] != num_heads * head_width: - return None +def _confirm(param: torch.Tensor, candidates): + """Resolve candidates against the shape. Returns (num_heads or None, reason). + + Only an exact `rows == heads * width` confirms a candidate. Several candidates can confirm + at once; that is only an ambiguity if they disagree on the head count, which is the whole + output, so agreeing candidates are not a conflict. + """ + if param.ndim != 2: + return None, "not-2d" + if not candidates: + return None, "no-candidate-geometry" + + rows = param.shape[0] + exact = [c for c in candidates if rows == c[0] * c[1]] + if not exact: + return None, "width-mismatch" + + head_counts = {c[0] for c in exact} + if len(head_counts) > 1: + return None, "ambiguous:" + "/".join(f"{c[2]}={c[0]}" for c in sorted(exact, key=lambda c: c[2])) + return exact[0][0], exact[0][2] + +def _attention_head_count(param_name: str, param: torch.Tensor, model: torch.nn.Module): + """Heads this projection splits into on dim 0, or None to leave it on the full-matrix path. + + Takes the model for callers holding one parameter. `set_optimizer_flags` builds the reader + once and calls `_resolve_attention_head_count` directly, since it is loop-invariant. + """ + meta, text_config = _per_head_muon_meta(model) + if meta is None: + return None + num_heads, _ = _resolve_attention_head_count(param_name, param, meta, text_config) return num_heads +def _resolve_attention_head_count(param_name: str, param: torch.Tensor, meta, text_config): + """As above, with the reason, so callers can report why a parameter was not tagged.""" + kind = _classify_leaf(_leaf_module_name(param_name)) + if kind is None: + return None, "not-attention" + if kind == NOT_HEAD_BLOCKED: + return None, NOT_HEAD_BLOCKED + return _confirm(param, _geometry_candidates(kind, meta, text_config)) + + +def _report_per_head_tagging(tagged: dict, skipped: dict) -> None: + """Report what an explicit `per_head_muon: true` actually did. + + An opt-in that silently does nothing is the failure this guards. The systemic case is + tensor parallelism: the config describes the whole model while each rank holds a shard, so + every projection fails its width check and the feature is off model-wide while the user + believes it is on. That is an error rather than a warning, because there is no partial + result to keep. + """ + if not tagged: + raise ValueError("per_head_muon is enabled but no attention projection could be tagged. Per-head " + "Newton-Schulz is therefore inactive for every parameter. Likely causes: tensor " + "parallelism is active, so each rank holds a shard whose width no longer matches " + "the config; the architecture's attention layout is not recognized; or the model " + "has no attention projections. Unset per_head_muon to train without it. " + f"Leaves examined: {dict(sorted(skipped.items())) or 'none'}") + + unrecognized = { + leaf: reason + for leaf, reason in skipped.items() if reason not in (NOT_HEAD_BLOCKED, "not-attention") + } + if unrecognized: + logger.warning( + "per_head_muon: %s matched an attention name but no candidate geometry confirmed them; " + "they stay on the full-matrix path", dict(sorted(unrecognized.items()))) + logger.info("per_head_muon: tagged %s", dict(sorted(tagged.items()))) + + def set_optimizer_flags(config_class: DeepSpeedConfig, model: torch.nn.Module) -> None: if config_class.optimizer_name == MUON_OPTIMIZER: per_head = bool((config_class.optimizer_params or {}).get("per_head_muon", False)) + meta, text_config = _per_head_muon_meta(model) if per_head else (None, None) + tagged: dict = {} + skipped: dict = {} + for name, p in model.named_parameters(): if p.ndim >= 2 and not any(keyword in name.lower() for keyword in ("embed", "lm_head")): setattr(p, "use_muon", True) else: setattr(p, "use_muon", False) - num_heads = _attention_head_count(name, p, model) if (per_head and p.use_muon) else None + num_heads = None + if per_head and p.use_muon and meta is not None: + num_heads, reason = _resolve_attention_head_count(name, p, meta, text_config) + leaf = _leaf_module_name(name) + if num_heads is not None: + tagged[leaf] = f"{num_heads} heads of {p.shape[0] // num_heads} ({reason})" + else: + skipped[leaf] = reason setattr(p, "muon_num_heads", num_heads) + if per_head: + _report_per_head_tagging(tagged, skipped) + def initialize( args: Any = None, diff --git a/deepspeed/runtime/zero/muon/original_muon.py b/deepspeed/runtime/zero/muon/original_muon.py index 71520a42b4bc..f5edc2a796e9 100644 --- a/deepspeed/runtime/zero/muon/original_muon.py +++ b/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 diff --git a/docs/_pages/config-json.md b/docs/_pages/config-json.md index 4fda25f47167..2c4f3850c08f 100644 --- a/docs/_pages/config-json.md +++ b/docs/_pages/config-json.md @@ -60,6 +60,39 @@ 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 | + +**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. + +**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. The most likely cause is tensor parallelism: the config describes the +whole model while each rank holds a shard, so every projection fails its width check. 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. diff --git a/tests/unit/ops/muon/test_per_head_muon_e2e.py b/tests/unit/ops/muon/test_per_head_muon_e2e.py index c7239d0e7afa..280a84d4a4cf 100644 --- a/tests/unit/ops/muon/test_per_head_muon_e2e.py +++ b/tests/unit/ops/muon/test_per_head_muon_e2e.py @@ -1,4 +1,3 @@ -# Copyright (c) Microsoft Corporation. # SPDX-License-Identifier: Apache-2.0 # DeepSpeed Team diff --git a/tests/unit/runtime/zero/test_per_head_muon.py b/tests/unit/runtime/zero/test_per_head_muon.py index c2940819d3e3..237004cf103e 100644 --- a/tests/unit/runtime/zero/test_per_head_muon.py +++ b/tests/unit/runtime/zero/test_per_head_muon.py @@ -1,4 +1,3 @@ -# Copyright (c) Microsoft Corporation. # SPDX-License-Identifier: Apache-2.0 # DeepSpeed Team @@ -14,10 +13,10 @@ import pytest import torch -from deepspeed.accelerator import get_accelerator from deepspeed.runtime.zero.muon.original_muon import ( muon_update, zeropower_via_gram_newtonschulz, + ns_compute_dtype, zeropower_via_newtonschulz5, ) @@ -29,21 +28,13 @@ def _ns_tolerance(ns_method): amplifies rounding, so batched and unbatched NS agree to a handful of ulps rather than bitwise. Deriving the bound from the dtype keeps it honest instead of tuned to pass. """ - if ns_method == "gram": - dtype = torch.float16 if get_accelerator().is_fp16_supported() else torch.float32 - else: - dtype = torch.bfloat16 if get_accelerator().is_bf16_supported() else torch.float32 - eps = torch.finfo(dtype).eps + eps = torch.finfo(ns_compute_dtype(ns_method)).eps return dict(rtol=8 * eps, atol=8 * eps) def _norm_rtol(ns_method): """Scale agreement: a couple of ulps of the compute dtype, and never looser than 1%.""" - if ns_method == "gram": - dtype = torch.float16 if get_accelerator().is_fp16_supported() else torch.float32 - else: - dtype = torch.bfloat16 if get_accelerator().is_bf16_supported() else torch.float32 - return max(1e-2, 2 * torch.finfo(dtype).eps) + return max(1e-2, 2 * torch.finfo(ns_compute_dtype(ns_method)).eps) def _update_only(grad, momentum, beta=0.95, nesterov=True): diff --git a/tests/unit/runtime/zero/test_per_head_muon_tagging.py b/tests/unit/runtime/zero/test_per_head_muon_tagging.py index 9701e716cbaf..f5ff1ea78200 100644 --- a/tests/unit/runtime/zero/test_per_head_muon_tagging.py +++ b/tests/unit/runtime/zero/test_per_head_muon_tagging.py @@ -1,4 +1,3 @@ -# Copyright (c) Microsoft Corporation. # SPDX-License-Identifier: Apache-2.0 # DeepSpeed Team @@ -30,7 +29,12 @@ def __init__(self, hidden=64, q_heads=8, kv_heads=2, head_dim=8, fused=False): self.embed_tokens = torch.nn.Embedding(16, hidden) if fused: self.qkv_proj = torch.nn.Linear(hidden, (q_heads + 2 * kv_heads) * head_dim, bias=False) - self.config = SimpleNamespace(num_attention_heads=q_heads, num_key_value_heads=kv_heads) + # Real configs carry the per-head width, either as head_dim or derivably from + # hidden_size. Without one the shape cannot confirm the name, and the tagger declines. + self.config = SimpleNamespace(num_attention_heads=q_heads, + num_key_value_heads=kv_heads, + hidden_size=hidden, + head_dim=head_dim) def _flags(model, per_head=True): @@ -301,3 +305,112 @@ def test_head_count_comes_from_the_shared_extractor(): model = SimpleNamespace(config=SimpleNamespace(n_head=8, hidden_size=64, head_dim=8)) assert _attention_head_count("l.0.attn.q_proj.weight", torch.zeros(64, 64), model) == 8 + + +# --- candidate resolution ------------------------------------------------------ +# +# `q_proj` can be either the standard query projection or, on an MLA model without a +# q_lora_rank, the query up-projection. Both candidates are evaluated and the shape decides, +# so the outcome does not depend on which config fields happen to be present. + + +def _kimi_k3_hybrid_config(): + """Kimi-K3-0.40B: linear-attention layers on a config that also carries MLA fields. + + `linear_attn_config` gives `num_heads: 8, head_dim: 32`, so its `q_proj` is (256, 1024). + The top-level MLA dimensions belong to the model's two MLA layers, and `head_dim` is 74. + Neither top-level geometry describes the KDA projection. + """ + return SimpleNamespace(num_attention_heads=8, + num_key_value_heads=8, + hidden_size=1024, + head_dim=74, + qk_nope_head_dim=64, + qk_rope_head_dim=32, + v_head_dim=64) + + +@pytest.mark.parametrize("leaf", ["q_proj", "k_proj", "v_proj"]) +def test_linear_attention_on_a_config_with_mla_leftovers_is_declined(leaf): + """No candidate confirms, so it stays on the full-matrix path. + + 8 x (qk_nope 64 + qk_rope 32) = 768 and 8 x head_dim 74 = 592, against 256 rows. This is + the case tracked in #8420; until the linear-attention geometry is read, declining is the + correct outcome and it must come from the shape rather than from branch ordering. + """ + model = SimpleNamespace(config=_kimi_k3_hybrid_config()) + + assert _attention_head_count(f"model.layers.0.self_attn.{leaf}.weight", torch.zeros(256, 1024), model) is None + + +def test_two_candidates_agreeing_on_the_head_count_are_not_ambiguous(): + """Ambiguity is about the answer, not the route. + + The output is a head count, so two candidates that confirm with the same count give the + same answer and there is nothing to be ambiguous about. + """ + config = SimpleNamespace(num_attention_heads=8, + num_key_value_heads=8, + hidden_size=512, + head_dim=96, + qk_nope_head_dim=64, + qk_rope_head_dim=32, + v_head_dim=64) + model = SimpleNamespace(config=config) + + # 8 x 96 = 768 by head_dim, and 8 x (64 + 32) = 768 by the MLA width. + assert _attention_head_count("l.0.self_attn.q_proj.weight", torch.zeros(768, 512), model) == 8 + + +def test_candidates_that_disagree_on_the_head_count_are_skipped(): + """A real ambiguity: both confirm the shape, and they give different answers.""" + from deepspeed import _confirm + + candidates = [(8, 96, "mla-q"), (12, 64, "head-dim")] + num_heads, reason = _confirm(torch.zeros(768, 512), candidates) + + assert num_heads is None + assert reason.startswith("ambiguous:") + assert "mla-q=8" in reason and "head-dim=12" in reason + + +def test_a_config_without_a_per_head_width_is_declined(): + """Divisibility alone is not confirmation. + + `rows % heads == 0` holds for matrices that are not head-blocked at all, which is how + o_proj used to slip through. Without a width there is nothing to confirm against. + """ + config = SimpleNamespace(num_attention_heads=8, num_key_value_heads=8) + model = SimpleNamespace(config=config) + + assert _attention_head_count("l.0.self_attn.q_proj.weight", torch.zeros(512, 512), model) is None + + +def test_head_dim_is_derived_when_the_config_omits_it(): + """Configs that leave head_dim out still define it as hidden_size // num_attention_heads.""" + config = SimpleNamespace(num_attention_heads=8, num_key_value_heads=8, hidden_size=512) + model = SimpleNamespace(config=config) + + assert _attention_head_count("l.0.self_attn.q_proj.weight", torch.zeros(512, 512), model) == 8 + + +def test_the_flag_errors_rather_than_silently_doing_nothing(): + """An explicit opt-in that tags nothing is the tensor-parallel failure mode. + + Under TP the config describes the whole model while each rank holds a shard, so every + projection fails its width check and per-head is off model-wide while the user believes it + is on. There is no partial result to keep, so this is an error. + """ + + class _NoAttention(torch.nn.Module): + + def __init__(self): + super().__init__() + self.mlp = torch.nn.Linear(64, 64, bias=False) + self.config = SimpleNamespace(num_attention_heads=8, num_key_value_heads=8, hidden_size=64, head_dim=8) + + with pytest.raises(ValueError, match="no attention projection could be tagged"): + _flags(_NoAttention()) + + # ...and with the flag off it is simply not asked for. + assert _flags(_NoAttention(), per_head=False)["mlp.weight"] is None From 943fd03b1cb8305fbde698d7e53f97c3e1bc9382 Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Sun, 6 Sep 2026 19:57:56 +0800 Subject: [PATCH 09/11] [muon] Re-resolve the head count against the shard under tensor parallelism set_optimizer_flags runs before the engine partitions the model, so the count it records is the model's, not the rank's. AutoTP replaces the parameter's .data in place, so the tag rides onto a column-parallel shard and nothing catches it: with tp=2 a tag of 8 heads lands on a shard holding 4 heads' worth of rows, out_features % num_heads still divides, and Newton-Schulz runs on half of each head. Measured on a Llama with 8 heads of 32 and autotp_size 2: every q/k/v tagged 8 heads of 16. The per-head width is what a column-parallel split leaves alone, so record it and re-derive the count from the shard. A shard whose rows are not a multiple of the width does not hold whole heads and is dropped. This is not only a repair. A column-parallel shard holds whole heads, so per-head Newton-Schulz on the shard is bit-identical to the corresponding blocks of per-head Newton-Schulz on the whole matrix - the split is along the same axis the batch is taken over. There is a test asserting equality, not closeness, for both kernels at tp=2 and tp=4, against the whole-matrix path, which differs by more than 10% on the same input. Signed-off-by: alanhuangyoo --- deepspeed/__init__.py | 57 ++++++- deepspeed/runtime/engine.py | 3 + docs/_pages/config-json.md | 16 +- .../muon/test_per_head_muon_under_sharding.py | 140 ++++++++++++++++++ .../test_per_head_muon_tensor_parallel.py | 106 +++++++++++++ 5 files changed, 313 insertions(+), 9 deletions(-) create mode 100644 tests/unit/ops/muon/test_per_head_muon_under_sharding.py create mode 100644 tests/unit/runtime/zero/test_per_head_muon_tensor_parallel.py diff --git a/deepspeed/__init__.py b/deepspeed/__init__.py index 7b64f4c5dcf7..e1dc66277df4 100755 --- a/deepspeed/__init__.py +++ b/deepspeed/__init__.py @@ -249,11 +249,13 @@ def _report_per_head_tagging(tagged: dict, skipped: dict) -> None: """ if not tagged: raise ValueError("per_head_muon is enabled but no attention projection could be tagged. Per-head " - "Newton-Schulz is therefore inactive for every parameter. Likely causes: tensor " - "parallelism is active, so each rank holds a shard whose width no longer matches " - "the config; the architecture's attention layout is not recognized; or the model " - "has no attention projections. Unset per_head_muon to train without it. " - f"Leaves examined: {dict(sorted(skipped.items())) or 'none'}") + "Newton-Schulz is therefore inactive for every parameter. Likely causes: the model " + "arrives already sharded by an external tensor-parallel implementation, so each " + "rank holds a shard whose width no longer matches the config; the architecture's " + "attention layout is not recognized; or the model has no attention projections. " + "AutoTP is not one of the causes - it partitions after this runs, and the counts " + "are re-resolved against the shards afterwards. Unset per_head_muon to train " + f"without it. Leaves examined: {dict(sorted(skipped.items())) or 'none'}") unrecognized = { leaf: reason @@ -288,11 +290,56 @@ def set_optimizer_flags(config_class: DeepSpeedConfig, model: torch.nn.Module) - else: skipped[leaf] = reason setattr(p, "muon_num_heads", num_heads) + # The width, not the count, is what survives a column-parallel split; see + # `resolve_per_head_muon_after_sharding`. + setattr(p, "muon_head_dim", p.shape[0] // num_heads if num_heads else None) if per_head: _report_per_head_tagging(tagged, skipped) +def resolve_per_head_muon_after_sharding(model: torch.nn.Module) -> None: + """Re-derive head counts from the shapes the parameters actually have. + + `set_optimizer_flags` runs before the engine partitions the model, so the count it records + is the model's, not the rank's. Column-parallel tensor parallelism splits an attention + projection on dim 0, which is the axis the heads are on, so after the split the per-head + width is unchanged and the head count is not. Nothing catches that on its own: with tp=2 a + tag of 8 heads lands on a shard holding 4 heads' worth of rows, `out_features % num_heads` + still divides, and Newton-Schulz runs on half of each head. + + Re-deriving the count from the width is not just a repair. A column-parallel shard holds + whole heads, so per-head Newton-Schulz on the shard is exactly the corresponding blocks of + per-head Newton-Schulz on the whole matrix - the split is along the same axis the batch is + taken over. A shard whose rows are not a multiple of the width does not hold whole heads, + and is dropped rather than guessed at. + """ + tagged, dropped = {}, {} + for name, p in model.named_parameters(): + head_dim = getattr(p, "muon_head_dim", None) + if head_dim is None: + continue + leaf = _leaf_module_name(name) + rows = p.shape[0] + if rows % head_dim: + setattr(p, "muon_num_heads", None) + dropped[leaf] = f"{rows} rows do not divide into heads of {head_dim}" + continue + setattr(p, "muon_num_heads", rows // head_dim) + tagged[leaf] = f"{rows // head_dim} heads of {head_dim}" + + if not tagged and not dropped: + return + if dropped: + logger.warning("per_head_muon: %s are sharded across head boundaries and stay on the full-matrix " + "path", dict(sorted(dropped.items()))) + if not tagged: + raise ValueError("per_head_muon is enabled but every tagged projection is sharded across head " + "boundaries, so per-head Newton-Schulz is inactive for all of them. Unset " + f"per_head_muon to train without it. Parameters examined: {dict(sorted(dropped.items()))}") + logger.info("per_head_muon: after sharding, tagged %s", dict(sorted(tagged.items()))) + + def initialize( args: Any = None, model: torch.nn.Module = None, diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 6d0857ae59b7..46ceeb162796 100644 --- a/deepspeed/runtime/engine.py +++ b/deepspeed/runtime/engine.py @@ -413,6 +413,9 @@ def __init__(self, self._configure_expert_parallel(model) 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(): diff --git a/docs/_pages/config-json.md b/docs/_pages/config-json.md index 2c4f3850c08f..d152d8d7515d 100644 --- a/docs/_pages/config-json.md +++ b/docs/_pages/config-json.md @@ -87,12 +87,20 @@ parameter is tagged only when its rows equal `num_heads * width` exactly for one 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. The most likely cause is tensor parallelism: the config describes the -whole model while each rank holds a shard, so every projection fails its width check. 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. +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. diff --git a/tests/unit/ops/muon/test_per_head_muon_under_sharding.py b/tests/unit/ops/muon/test_per_head_muon_under_sharding.py new file mode 100644 index 000000000000..7b2a57a92e1e --- /dev/null +++ b/tests/unit/ops/muon/test_per_head_muon_under_sharding.py @@ -0,0 +1,140 @@ +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""Why per-head Newton-Schulz survives a column-parallel split and the whole-matrix path does not. + +Column-parallel tensor parallelism splits an attention projection on dim 0, which is the axis +per-head Newton-Schulz batches over. Each rank therefore holds whole heads, and orthogonalizing +them is the same computation whether the other ranks' heads are present or not. The whole-matrix +path has no such property: orthogonalizing a block of rows is not the same as taking that block +out of the orthogonalization of all of them. See #8367. +""" + +import pytest +import torch + +import deepspeed +from deepspeed.accelerator import get_accelerator +from deepspeed.runtime.zero.muon.original_muon import ( + _per_head_orthogonalize, + zeropower_via_gram_newtonschulz, + zeropower_via_newtonschulz5, +) +from unit.common import DistributedTest + +HEADS, HEAD_DIM, HIDDEN, STEPS = 8, 32, 256, 5 + + +@pytest.fixture +def grad(): + torch.manual_seed(0) + return torch.randn(HEADS * HEAD_DIM, HIDDEN, device=get_accelerator().device_name()) + + +def _relative(a, b): + return ((a - b).norm() / b.norm()).item() + + +def _column_parallel_shards(tensor, tp): + rows = tensor.shape[0] // tp + return [tensor[r * rows:(r + 1) * rows].contiguous() for r in range(tp)] + + +@pytest.mark.parametrize("ns_method", ["gram", "standard"]) +@pytest.mark.parametrize("tp", [2, 4]) +def test_per_head_on_a_shard_is_the_shard_of_per_head(grad, ns_method, tp): + """Exactly equal, not close: the shards are the same matrices in the same batch.""" + whole = _per_head_orthogonalize(grad.clone(), HEADS, STEPS, ns_method) + sharded = torch.cat([ + _per_head_orthogonalize(shard.clone(), HEADS // tp, STEPS, ns_method) + for shard in _column_parallel_shards(grad, tp) + ]) + + assert torch.equal(sharded, whole) + + +@pytest.mark.parametrize("ns_method", ["gram", "standard"]) +def test_the_whole_matrix_path_does_not_survive_the_split(grad, ns_method): + """The comparison this is measured against, so "exact" above means something. + + Newton-Schulz on a block of rows is a different computation from the same block of + Newton-Schulz on every row, and the difference is not small. + """ + ns_fn = zeropower_via_gram_newtonschulz if ns_method == "gram" else zeropower_via_newtonschulz5 + whole = ns_fn(grad.clone(), steps=STEPS) + sharded = torch.cat([ns_fn(shard.clone(), steps=STEPS) for shard in _column_parallel_shards(grad, 2)]) + + assert _relative(sharded, whole) > 0.1 + + +def test_a_stale_head_count_splits_heads_in_half(grad): + """What the tag does under tp=2 if it is not re-resolved against the shard. + + 128 rows still divide by 8, so the divisibility check passes and each head is cut in two. + """ + whole = _per_head_orthogonalize(grad.clone(), HEADS, STEPS, "gram") + shard = _column_parallel_shards(grad, 2)[0] + + correct = _per_head_orthogonalize(shard.clone(), HEADS // 2, STEPS, "gram") + stale = _per_head_orthogonalize(shard.clone(), HEADS, STEPS, "gram") + + assert torch.equal(correct, whole[:shard.shape[0]]) + assert _relative(stale, whole[:shard.shape[0]]) > 0.1 + + +class TestPerHeadMuonUnderAutoTP(DistributedTest): + """The tags AutoTP leaves behind, through a real `deepspeed.initialize`. + + `set_optimizer_flags` runs before `_configure_tensor_parallel`, and AutoTP replaces the + parameter's `.data` in place, so the tag made against the whole model rides onto a shard. + """ + world_size = 2 + + def test_the_tags_describe_the_shard_and_not_the_model(self): + transformers = pytest.importorskip("transformers") + heads, head_dim = 8, 32 + config = transformers.LlamaConfig(hidden_size=heads * head_dim, + num_attention_heads=heads, + num_key_value_heads=heads, + num_hidden_layers=2, + intermediate_size=2 * heads * head_dim, + vocab_size=128) + model = transformers.AutoModelForCausalLM.from_config(config) + + engine, _, _, _ = deepspeed.initialize(model=model, + model_parameters=model.parameters(), + config={ + "train_micro_batch_size_per_gpu": 1, + "gradient_accumulation_steps": 1, + "bf16": { + "enabled": True + }, + "zero_optimization": { + "stage": 1 + }, + "tensor_parallel": { + "autotp_size": 2 + }, + "optimizer": { + "type": "Muon", + "params": { + "lr": 1e-3, + "per_head_muon": True + } + }, + }) + + tagged = {n: p for n, p in model.named_parameters() if getattr(p, "muon_num_heads", None)} + assert tagged, "AutoTP left nothing tagged" + for name, param in tagged.items(): + assert param.shape[0] // param.muon_num_heads == head_dim, \ + f"{name}: {param.shape[0]} rows over {param.muon_num_heads} heads is not {head_dim} wide" + assert param.muon_num_heads == heads // 2, \ + f"{name}: tp=2 leaves {heads // 2} heads on this rank, tagged {param.muon_num_heads}" + + # AutoTP asserts every rank in the TP group sees the same batch. + ids = torch.arange(8, device=engine.device).unsqueeze(0) % 128 + out = engine(input_ids=ids, labels=ids) + engine.backward(out.loss) + engine.step() + assert all(torch.isfinite(p).all() for p in model.parameters()) diff --git a/tests/unit/runtime/zero/test_per_head_muon_tensor_parallel.py b/tests/unit/runtime/zero/test_per_head_muon_tensor_parallel.py new file mode 100644 index 000000000000..cb12d608bade --- /dev/null +++ b/tests/unit/runtime/zero/test_per_head_muon_tensor_parallel.py @@ -0,0 +1,106 @@ +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""Per-head Muon under tensor parallelism. See #8367. + +`set_optimizer_flags` runs before the engine partitions the model, so the head count it records +counts the model's heads. Column-parallel TP then splits attention projections on dim 0, which +is the axis the heads are on: the per-head width survives the split and the count does not. +Nothing catches that on its own, because the stale count usually still divides the shard. +""" + +import pytest +import torch + +from deepspeed import resolve_per_head_muon_after_sharding + + +class _Attn(torch.nn.Module): + """8 heads of 32, as the whole model sees it.""" + + def __init__(self, hidden=256, heads=8, head_dim=32): + super().__init__() + self.q_proj = torch.nn.Linear(hidden, heads * head_dim, bias=False) + self.k_proj = torch.nn.Linear(hidden, heads * head_dim, bias=False) + self.mlp = torch.nn.Linear(hidden, hidden, bias=False) + for p, num_heads in ((self.q_proj.weight, heads), (self.k_proj.weight, heads), (self.mlp.weight, None)): + p.muon_num_heads = num_heads + p.muon_head_dim = p.shape[0] // num_heads if num_heads else None + + def shard(self, tp: int, rows=None, leaves=("q_proj", "k_proj")): + """Replace the weights with column-parallel shards, as AutoTP's `_tp_partition` does.""" + for leaf in leaves: + weight = getattr(self, leaf).weight + keep = rows if rows is not None else weight.shape[0] // tp + weight.data = weight.data[:keep].clone() + + +def test_the_head_count_follows_the_shard(): + """tp=2 leaves 4 heads on this rank; the tag has to say 4, not the model's 8.""" + attn = _Attn() + attn.shard(tp=2) + + resolve_per_head_muon_after_sharding(attn) + + assert attn.q_proj.weight.shape == (128, 256) + assert attn.q_proj.weight.muon_num_heads == 4 + assert attn.k_proj.weight.muon_num_heads == 4 + + +def test_the_stale_count_would_have_split_heads_in_half(): + """Why this is not caught by the existing shape check: 128 % 8 == 0. + + The divisibility assert in `_per_head_orthogonalize` passes on the stale count, so without + this pass Newton-Schulz runs on 8 blocks of 16 - half of each head - and says nothing. + """ + attn = _Attn() + attn.shard(tp=2) + + rows, stale = attn.q_proj.weight.shape[0], 8 + assert rows % stale == 0, "the stale count divides, which is why it needs correcting rather than asserting" + assert rows // stale == 16 != 32 + + +def test_a_shard_that_splits_a_head_is_dropped(): + """A shard that does not hold whole heads has no per-head structure to use.""" + attn = _Attn() + attn.shard(tp=2, rows=144, leaves=("q_proj", )) # 4.5 heads of 32 + + resolve_per_head_muon_after_sharding(attn) + + assert attn.q_proj.weight.muon_num_heads is None + assert attn.k_proj.weight.muon_num_heads == 8, "one bad shard does not turn the feature off elsewhere" + + +def test_an_unsharded_model_keeps_its_count(): + attn = _Attn() + + resolve_per_head_muon_after_sharding(attn) + + assert attn.q_proj.weight.muon_num_heads == 8 + assert attn.q_proj.weight.shape[0] == 256 + + +def test_untagged_parameters_are_left_alone(): + attn = _Attn() + attn.shard(tp=2) + + resolve_per_head_muon_after_sharding(attn) + + assert getattr(attn.mlp.weight, "muon_num_heads", "MISSING") is None + + +def test_a_model_with_no_tags_at_all_is_a_no_op(): + """Muon without `per_head_muon`, and every non-Muon model: nothing to resolve, no error.""" + model = torch.nn.Linear(8, 8, bias=False) + + resolve_per_head_muon_after_sharding(model) # must not raise + + +def test_it_raises_when_the_flag_ends_up_doing_nothing(): + """Same contract as the tagging pass: an opt-in that silently does nothing is the failure.""" + attn = _Attn() + attn.shard(tp=2, rows=144) + + with pytest.raises(ValueError, match="sharded across head boundaries"): + resolve_per_head_muon_after_sharding(attn) From 5d164954cba8d782a26c88022371e5ff0bbda9bc Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Sun, 6 Sep 2026 20:13:31 +0800 Subject: [PATCH 10/11] [muon] Read the layer's shape, not the ZeRO-3 partition's zero.Init replaces a partitioned parameter's data with a flat placeholder and records the layer's shape as ds_shape, so param.shape is 1-D for every parameter in the model. The width check then confirms nothing and the flag raises on a model whose layout it could read perfectly well. Same root cause as #8438, which fixes the use_muon test on master; this applies it to the tagger's shape reads as well. Signed-off-by: alanhuangyoo --- deepspeed/__init__.py | 25 ++++++++++++++----- .../zero/test_per_head_muon_tagging.py | 19 ++++++++++++++ 2 files changed, 38 insertions(+), 6 deletions(-) diff --git a/deepspeed/__init__.py b/deepspeed/__init__.py index e1dc66277df4..dc1460c5bfd6 100755 --- a/deepspeed/__init__.py +++ b/deepspeed/__init__.py @@ -102,6 +102,18 @@ def _parse_version(version_str): QUERY, KV, MLA_Q, MLA_KV, NOT_HEAD_BLOCKED = "query", "kv", "mla_q", "mla_kv", "not-head-blocked" +def _layer_shape(param: torch.Tensor): + """The parameter's shape as a layer, rather than as a ZeRO-3 partition. + + Under ``deepspeed.zero.Init`` a partitioned parameter's data is a flat placeholder - + ``torch.Size([0])`` on the ranks that do not hold it - and the shape it has as a layer is + recorded as ``ds_shape``. Reading ``param.shape`` there sees a 1-D tensor for every + parameter in the model. + """ + ds_shape = getattr(param, "ds_shape", None) + return tuple(param.shape) if ds_shape is None else tuple(ds_shape) + + def _per_head_muon_meta(model: torch.nn.Module): """The head-count reader and the config the widths come from, built once per model. @@ -199,12 +211,13 @@ def _confirm(param: torch.Tensor, candidates): at once; that is only an ambiguity if they disagree on the head count, which is the whole output, so agreeing candidates are not a conflict. """ - if param.ndim != 2: + shape = _layer_shape(param) + if len(shape) != 2: return None, "not-2d" if not candidates: return None, "no-candidate-geometry" - rows = param.shape[0] + rows = shape[0] exact = [c for c in candidates if rows == c[0] * c[1]] if not exact: return None, "width-mismatch" @@ -276,7 +289,7 @@ def set_optimizer_flags(config_class: DeepSpeedConfig, model: torch.nn.Module) - skipped: dict = {} for name, p in model.named_parameters(): - if p.ndim >= 2 and not any(keyword in name.lower() for keyword in ("embed", "lm_head")): + if len(_layer_shape(p)) >= 2 and not any(keyword in name.lower() for keyword in ("embed", "lm_head")): setattr(p, "use_muon", True) else: setattr(p, "use_muon", False) @@ -286,13 +299,13 @@ def set_optimizer_flags(config_class: DeepSpeedConfig, model: torch.nn.Module) - num_heads, reason = _resolve_attention_head_count(name, p, meta, text_config) leaf = _leaf_module_name(name) if num_heads is not None: - tagged[leaf] = f"{num_heads} heads of {p.shape[0] // num_heads} ({reason})" + tagged[leaf] = f"{num_heads} heads of {_layer_shape(p)[0] // num_heads} ({reason})" else: skipped[leaf] = reason setattr(p, "muon_num_heads", num_heads) # The width, not the count, is what survives a column-parallel split; see # `resolve_per_head_muon_after_sharding`. - setattr(p, "muon_head_dim", p.shape[0] // num_heads if num_heads else None) + setattr(p, "muon_head_dim", _layer_shape(p)[0] // num_heads if num_heads else None) if per_head: _report_per_head_tagging(tagged, skipped) @@ -320,7 +333,7 @@ def resolve_per_head_muon_after_sharding(model: torch.nn.Module) -> None: if head_dim is None: continue leaf = _leaf_module_name(name) - rows = p.shape[0] + rows = _layer_shape(p)[0] if rows % head_dim: setattr(p, "muon_num_heads", None) dropped[leaf] = f"{rows} rows do not divide into heads of {head_dim}" diff --git a/tests/unit/runtime/zero/test_per_head_muon_tagging.py b/tests/unit/runtime/zero/test_per_head_muon_tagging.py index f5ff1ea78200..cacae1c6a6d5 100644 --- a/tests/unit/runtime/zero/test_per_head_muon_tagging.py +++ b/tests/unit/runtime/zero/test_per_head_muon_tagging.py @@ -414,3 +414,22 @@ def __init__(self): # ...and with the flag off it is simply not asked for. assert _flags(_NoAttention(), per_head=False)["mlp.weight"] is None + + +def test_the_head_count_is_read_from_the_layer_shape_under_zero_init(): + """`zero.Init` leaves a flat placeholder and records the layer's shape as `ds_shape`. + + Reading `param.shape` there sees a 1-D tensor for every parameter, so nothing confirms and + the flag raises on a model it could describe perfectly well. + """ + model = _Attn(hidden=64, q_heads=8, kv_heads=2, head_dim=8) + for p in model.parameters(): + p.ds_shape = torch.Size(p.shape) + p.data = torch.zeros(0, dtype=p.dtype) + + tags = _flags(model) + + assert model.q_proj.weight.ndim == 1, "the partitioned parameter really is 1-D here" + assert tags["q_proj.weight"] == 8 + assert tags["k_proj.weight"] == 2 + assert tags["o_proj.weight"] is None From 789a1e7c3ec98e44f13a4b0d687ea90e13cade57 Mon Sep 17 00:00:00 2001 From: alanhuangyoo Date: Wed, 9 Sep 2026 03:21:22 +0800 Subject: [PATCH 11/11] Consolidate the per-head Muon tests into two files @pengdurice asked for fewer files. Five, one per concern, become two split by what they need to run: tests/unit/runtime/zero/test_per_head_muon.py CPU the arithmetic (was test_per_head_muon.py) which parameters are tagged, and with how many heads (was test_per_head_muon_tagging.py) re-resolving the count against a tensor-parallel shard (was test_per_head_muon_tensor_parallel.py) tests/unit/ops/muon/test_per_head_muon_accelerator.py GPU a shard of the per-head result is the per-head result of the shard (was test_per_head_muon_under_sharding.py) end-to-end training (was test_per_head_muon_e2e.py) Nothing is dropped or rewritten: 77 tests before, 77 after, same names. The tagging module's _Attn and the tensor-parallel module's _Attn were different classes with the same name, so the latter is now _ShardedAttn. The accelerator file is not called test_per_head_muon.py because these directories have no __init__.py, and two same-named modules break collection when both are selected in one run. Signed-off-by: alanhuangyoo --- .../muon/test_per_head_muon_accelerator.py | 267 +++++++++ tests/unit/ops/muon/test_per_head_muon_e2e.py | 122 ---- .../muon/test_per_head_muon_under_sharding.py | 140 ----- tests/unit/runtime/zero/test_per_head_muon.py | 546 +++++++++++++++++- .../zero/test_per_head_muon_tagging.py | 435 -------------- .../test_per_head_muon_tensor_parallel.py | 106 ---- 6 files changed, 811 insertions(+), 805 deletions(-) create mode 100644 tests/unit/ops/muon/test_per_head_muon_accelerator.py delete mode 100644 tests/unit/ops/muon/test_per_head_muon_e2e.py delete mode 100644 tests/unit/ops/muon/test_per_head_muon_under_sharding.py delete mode 100644 tests/unit/runtime/zero/test_per_head_muon_tagging.py delete mode 100644 tests/unit/runtime/zero/test_per_head_muon_tensor_parallel.py diff --git a/tests/unit/ops/muon/test_per_head_muon_accelerator.py b/tests/unit/ops/muon/test_per_head_muon_accelerator.py new file mode 100644 index 000000000000..d1d183b69f53 --- /dev/null +++ b/tests/unit/ops/muon/test_per_head_muon_accelerator.py @@ -0,0 +1,267 @@ +# Copyright (c) DeepSpeed Team. +# SPDX-License-Identifier: Apache-2.0 + +# DeepSpeed Team +"""Per-head Muon on the accelerator: sharding equivalence, and end-to-end training. + +Two concerns, in order: + +1. **Why per-head survives a column-parallel split and the whole-matrix path does not.** + Column-parallel tensor parallelism splits an attention projection on dim 0, which is the + axis per-head Newton-Schulz batches over. Each rank therefore holds whole heads, and + orthogonalizing them is the same computation whether the other ranks' heads are present or + not. The whole-matrix path has no such property: orthogonalizing a block of rows is not the + same as taking that block out of the orthogonalization of all of them. + +2. **The whole path, running.** `deepspeed.initialize` tags the parameters, the ZeRO call + sites carry the tag into `muon_update`, and a real training loop takes steps with it, + across ZeRO stages and world size > 1. + +See #8367. The arithmetic and the tagging are pinned on CPU in +`tests/unit/runtime/zero/test_per_head_muon.py`. +""" + +from types import SimpleNamespace + +import pytest +import torch + +import deepspeed +from deepspeed.accelerator import get_accelerator +from deepspeed.runtime.zero.muon.original_muon import ( + _per_head_orthogonalize, + zeropower_via_gram_newtonschulz, + zeropower_via_newtonschulz5, +) +from unit.common import DistributedTest + +# --------------------------------------------------------------------------- +# 1. A shard of the per-head result is the per-head result of the shard +# --------------------------------------------------------------------------- + +HEADS, HEAD_DIM, HIDDEN, STEPS = 8, 32, 256, 5 + + +@pytest.fixture +def grad(): + torch.manual_seed(0) + return torch.randn(HEADS * HEAD_DIM, HIDDEN, device=get_accelerator().device_name()) + + +def _relative(a, b): + return ((a - b).norm() / b.norm()).item() + + +def _column_parallel_shards(tensor, tp): + rows = tensor.shape[0] // tp + return [tensor[r * rows:(r + 1) * rows].contiguous() for r in range(tp)] + + +@pytest.mark.parametrize("ns_method", ["gram", "standard"]) +@pytest.mark.parametrize("tp", [2, 4]) +def test_per_head_on_a_shard_is_the_shard_of_per_head(grad, ns_method, tp): + """Exactly equal, not close: the shards are the same matrices in the same batch.""" + whole = _per_head_orthogonalize(grad.clone(), HEADS, STEPS, ns_method) + sharded = torch.cat([ + _per_head_orthogonalize(shard.clone(), HEADS // tp, STEPS, ns_method) + for shard in _column_parallel_shards(grad, tp) + ]) + + assert torch.equal(sharded, whole) + + +@pytest.mark.parametrize("ns_method", ["gram", "standard"]) +def test_the_whole_matrix_path_does_not_survive_the_split(grad, ns_method): + """The comparison this is measured against, so "exact" above means something. + + Newton-Schulz on a block of rows is a different computation from the same block of + Newton-Schulz on every row, and the difference is not small. + """ + ns_fn = zeropower_via_gram_newtonschulz if ns_method == "gram" else zeropower_via_newtonschulz5 + whole = ns_fn(grad.clone(), steps=STEPS) + sharded = torch.cat([ns_fn(shard.clone(), steps=STEPS) for shard in _column_parallel_shards(grad, 2)]) + + assert _relative(sharded, whole) > 0.1 + + +def test_a_stale_head_count_splits_heads_in_half(grad): + """What the tag does under tp=2 if it is not re-resolved against the shard. + + 128 rows still divide by 8, so the divisibility check passes and each head is cut in two. + """ + whole = _per_head_orthogonalize(grad.clone(), HEADS, STEPS, "gram") + shard = _column_parallel_shards(grad, 2)[0] + + correct = _per_head_orthogonalize(shard.clone(), HEADS // 2, STEPS, "gram") + stale = _per_head_orthogonalize(shard.clone(), HEADS, STEPS, "gram") + + assert torch.equal(correct, whole[:shard.shape[0]]) + assert _relative(stale, whole[:shard.shape[0]]) > 0.1 + + +class TestPerHeadMuonUnderAutoTP(DistributedTest): + """The tags AutoTP leaves behind, through a real `deepspeed.initialize`. + + `set_optimizer_flags` runs before `_configure_tensor_parallel`, and AutoTP replaces the + parameter's `.data` in place, so the tag made against the whole model rides onto a shard. + """ + world_size = 2 + + def test_the_tags_describe_the_shard_and_not_the_model(self): + transformers = pytest.importorskip("transformers") + heads, head_dim = 8, 32 + config = transformers.LlamaConfig(hidden_size=heads * head_dim, + num_attention_heads=heads, + num_key_value_heads=heads, + num_hidden_layers=2, + intermediate_size=2 * heads * head_dim, + vocab_size=128) + model = transformers.AutoModelForCausalLM.from_config(config) + + engine, _, _, _ = deepspeed.initialize(model=model, + model_parameters=model.parameters(), + config={ + "train_micro_batch_size_per_gpu": 1, + "gradient_accumulation_steps": 1, + "bf16": { + "enabled": True + }, + "zero_optimization": { + "stage": 1 + }, + "tensor_parallel": { + "autotp_size": 2 + }, + "optimizer": { + "type": "Muon", + "params": { + "lr": 1e-3, + "per_head_muon": True + } + }, + }) + + tagged = {n: p for n, p in model.named_parameters() if getattr(p, "muon_num_heads", None)} + assert tagged, "AutoTP left nothing tagged" + for name, param in tagged.items(): + assert param.shape[0] // param.muon_num_heads == head_dim, \ + f"{name}: {param.shape[0]} rows over {param.muon_num_heads} heads is not {head_dim} wide" + assert param.muon_num_heads == heads // 2, \ + f"{name}: tp=2 leaves {heads // 2} heads on this rank, tagged {param.muon_num_heads}" + + # AutoTP asserts every rank in the TP group sees the same batch. + ids = torch.arange(8, device=engine.device).unsqueeze(0) % 128 + out = engine(input_ids=ids, labels=ids) + engine.backward(out.loss) + engine.step() + assert all(torch.isfinite(p).all() for p in model.parameters()) + + +# --------------------------------------------------------------------------- +# 2. End-to-end training +# --------------------------------------------------------------------------- + + +class AttentionModel(torch.nn.Module): + """Small GQA-shaped model: split QKV, 8 query heads over 2 kv heads.""" + + def __init__(self, hidden_dim=64, q_heads=8, kv_heads=2, head_dim=8, nlayers=2): + super().__init__() + self.q_heads, self.kv_heads, self.head_dim = q_heads, kv_heads, head_dim + self.blocks = torch.nn.ModuleList() + for _ in range(nlayers): + self.blocks.append( + torch.nn.ModuleDict({ + "q_proj": torch.nn.Linear(hidden_dim, q_heads * head_dim, bias=False), + "k_proj": torch.nn.Linear(hidden_dim, kv_heads * head_dim, bias=False), + "v_proj": torch.nn.Linear(hidden_dim, kv_heads * head_dim, bias=False), + "o_proj": torch.nn.Linear(q_heads * head_dim, hidden_dim, bias=False), + "mlp": torch.nn.Linear(hidden_dim, hidden_dim, bias=False), + })) + self.cross_entropy_loss = torch.nn.CrossEntropyLoss() + self.config = SimpleNamespace(num_attention_heads=q_heads, + num_key_value_heads=kv_heads, + hidden_size=hidden_dim, + head_dim=head_dim) + + def forward(self, x, y): + for b in self.blocks: + q, k, v = b["q_proj"](x), b["k_proj"](x), b["v_proj"](x) + rep = self.q_heads // self.kv_heads + attn = q * k.repeat(1, rep) + v.repeat(1, rep) + x = x + b["mlp"](b["o_proj"](attn)) + return self.cross_entropy_loss(x, y) + + +def _config(zero_stage, per_head, lr=0.01): + return { + "train_batch_size": 4, + "optimizer": { + "type": "muon", + "params": { + "lr": lr, + "adam_lr": lr, + "per_head_muon": per_head + } + }, + "zero_optimization": { + "stage": zero_stage, + # Muon does not support reduce-scatter; the existing Muon suite disables it the + # same way (see TestMuonRejectsReduceScatter). + "reduce_scatter": False, + }, + "fp16": { + "enabled": False + }, + "bf16": { + "enabled": True + }, + } + + +def _train(model, config, steps=6, hidden_dim=64, seed=1234): + engine, *_ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config) + tags = {n: getattr(p, "muon_num_heads", "MISSING") for n, p in model.named_parameters()} + gen = torch.Generator().manual_seed(seed) + losses = [] + for _ in range(steps): + x = torch.randn(4, hidden_dim, generator=gen).to(engine.device).to(torch.bfloat16) + y = torch.randint(0, hidden_dim, (4, ), generator=gen).to(engine.device) + loss = engine(x, y) + engine.backward(loss) + engine.step() + losses.append(loss.item()) + return tags, losses + + +@pytest.mark.parametrize("zero_stage", [1, 2, 3]) +class TestPerHeadMuonEndToEnd(DistributedTest): + world_size = 2 + + def test_tags_reach_the_optimizer(self, zero_stage): + """Per-parameter tags have to survive `deepspeed.initialize` into the ZeRO call sites.""" + torch.manual_seed(1234) + tags, losses = _train(AttentionModel(), _config(zero_stage, per_head=True)) + + assert tags["blocks.0.q_proj.weight"] == 8 + assert tags["blocks.0.k_proj.weight"] == 2, "GQA: kv projections carry the kv head count" + assert tags["blocks.0.v_proj.weight"] == 2 + assert tags["blocks.0.o_proj.weight"] is None, "o_proj's heads are on the input axis" + assert tags["blocks.0.mlp.weight"] is None + assert all(torch.isfinite(torch.tensor(loss)) for loss in losses) + + def test_opt_in_is_off_by_default(self, zero_stage): + torch.manual_seed(1234) + tags, _ = _train(AttentionModel(), _config(zero_stage, per_head=False)) + + assert all(v is None for v in tags.values()), {k: v for k, v in tags.items() if v is not None} + + def test_training_makes_progress_either_way(self, zero_stage): + """Both paths have to train; this is the baseline delock asked for alongside per-head.""" + torch.manual_seed(1234) + _, full = _train(AttentionModel(), _config(zero_stage, per_head=False)) + torch.manual_seed(1234) + _, per_head = _train(AttentionModel(), _config(zero_stage, per_head=True)) + + assert full[-1] < full[0], f"baseline did not train: {full}" + assert per_head[-1] < per_head[0], f"per-head did not train: {per_head}" diff --git a/tests/unit/ops/muon/test_per_head_muon_e2e.py b/tests/unit/ops/muon/test_per_head_muon_e2e.py deleted file mode 100644 index 280a84d4a4cf..000000000000 --- a/tests/unit/ops/muon/test_per_head_muon_e2e.py +++ /dev/null @@ -1,122 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -"""End-to-end training with per-head Muon, across ZeRO stages and world size > 1. - -The unit tests next to this pin the arithmetic and the tagging. These run the whole path: -`deepspeed.initialize` tags the parameters, the ZeRO call sites carry the tag into -`muon_update`, and a real training loop takes steps with it. See #8367. -""" - -from types import SimpleNamespace - -import pytest -import torch - -import deepspeed -from unit.common import DistributedTest - - -class AttentionModel(torch.nn.Module): - """Small GQA-shaped model: split QKV, 8 query heads over 2 kv heads.""" - - def __init__(self, hidden_dim=64, q_heads=8, kv_heads=2, head_dim=8, nlayers=2): - super().__init__() - self.q_heads, self.kv_heads, self.head_dim = q_heads, kv_heads, head_dim - self.blocks = torch.nn.ModuleList() - for _ in range(nlayers): - self.blocks.append( - torch.nn.ModuleDict({ - "q_proj": torch.nn.Linear(hidden_dim, q_heads * head_dim, bias=False), - "k_proj": torch.nn.Linear(hidden_dim, kv_heads * head_dim, bias=False), - "v_proj": torch.nn.Linear(hidden_dim, kv_heads * head_dim, bias=False), - "o_proj": torch.nn.Linear(q_heads * head_dim, hidden_dim, bias=False), - "mlp": torch.nn.Linear(hidden_dim, hidden_dim, bias=False), - })) - self.cross_entropy_loss = torch.nn.CrossEntropyLoss() - self.config = SimpleNamespace(num_attention_heads=q_heads, - num_key_value_heads=kv_heads, - hidden_size=hidden_dim, - head_dim=head_dim) - - def forward(self, x, y): - for b in self.blocks: - q, k, v = b["q_proj"](x), b["k_proj"](x), b["v_proj"](x) - rep = self.q_heads // self.kv_heads - attn = q * k.repeat(1, rep) + v.repeat(1, rep) - x = x + b["mlp"](b["o_proj"](attn)) - return self.cross_entropy_loss(x, y) - - -def _config(zero_stage, per_head, lr=0.01): - return { - "train_batch_size": 4, - "optimizer": { - "type": "muon", - "params": { - "lr": lr, - "adam_lr": lr, - "per_head_muon": per_head - } - }, - "zero_optimization": { - "stage": zero_stage, - # Muon does not support reduce-scatter; the existing Muon suite disables it the - # same way (see TestMuonRejectsReduceScatter). - "reduce_scatter": False, - }, - "fp16": { - "enabled": False - }, - "bf16": { - "enabled": True - }, - } - - -def _train(model, config, steps=6, hidden_dim=64, seed=1234): - engine, *_ = deepspeed.initialize(model=model, model_parameters=model.parameters(), config=config) - tags = {n: getattr(p, "muon_num_heads", "MISSING") for n, p in model.named_parameters()} - gen = torch.Generator().manual_seed(seed) - losses = [] - for _ in range(steps): - x = torch.randn(4, hidden_dim, generator=gen).to(engine.device).to(torch.bfloat16) - y = torch.randint(0, hidden_dim, (4, ), generator=gen).to(engine.device) - loss = engine(x, y) - engine.backward(loss) - engine.step() - losses.append(loss.item()) - return tags, losses - - -@pytest.mark.parametrize("zero_stage", [1, 2, 3]) -class TestPerHeadMuonEndToEnd(DistributedTest): - world_size = 2 - - def test_tags_reach_the_optimizer(self, zero_stage): - """Per-parameter tags have to survive `deepspeed.initialize` into the ZeRO call sites.""" - torch.manual_seed(1234) - tags, losses = _train(AttentionModel(), _config(zero_stage, per_head=True)) - - assert tags["blocks.0.q_proj.weight"] == 8 - assert tags["blocks.0.k_proj.weight"] == 2, "GQA: kv projections carry the kv head count" - assert tags["blocks.0.v_proj.weight"] == 2 - assert tags["blocks.0.o_proj.weight"] is None, "o_proj's heads are on the input axis" - assert tags["blocks.0.mlp.weight"] is None - assert all(torch.isfinite(torch.tensor(loss)) for loss in losses) - - def test_opt_in_is_off_by_default(self, zero_stage): - torch.manual_seed(1234) - tags, _ = _train(AttentionModel(), _config(zero_stage, per_head=False)) - - assert all(v is None for v in tags.values()), {k: v for k, v in tags.items() if v is not None} - - def test_training_makes_progress_either_way(self, zero_stage): - """Both paths have to train; this is the baseline delock asked for alongside per-head.""" - torch.manual_seed(1234) - _, full = _train(AttentionModel(), _config(zero_stage, per_head=False)) - torch.manual_seed(1234) - _, per_head = _train(AttentionModel(), _config(zero_stage, per_head=True)) - - assert full[-1] < full[0], f"baseline did not train: {full}" - assert per_head[-1] < per_head[0], f"per-head did not train: {per_head}" diff --git a/tests/unit/ops/muon/test_per_head_muon_under_sharding.py b/tests/unit/ops/muon/test_per_head_muon_under_sharding.py deleted file mode 100644 index 7b2a57a92e1e..000000000000 --- a/tests/unit/ops/muon/test_per_head_muon_under_sharding.py +++ /dev/null @@ -1,140 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -"""Why per-head Newton-Schulz survives a column-parallel split and the whole-matrix path does not. - -Column-parallel tensor parallelism splits an attention projection on dim 0, which is the axis -per-head Newton-Schulz batches over. Each rank therefore holds whole heads, and orthogonalizing -them is the same computation whether the other ranks' heads are present or not. The whole-matrix -path has no such property: orthogonalizing a block of rows is not the same as taking that block -out of the orthogonalization of all of them. See #8367. -""" - -import pytest -import torch - -import deepspeed -from deepspeed.accelerator import get_accelerator -from deepspeed.runtime.zero.muon.original_muon import ( - _per_head_orthogonalize, - zeropower_via_gram_newtonschulz, - zeropower_via_newtonschulz5, -) -from unit.common import DistributedTest - -HEADS, HEAD_DIM, HIDDEN, STEPS = 8, 32, 256, 5 - - -@pytest.fixture -def grad(): - torch.manual_seed(0) - return torch.randn(HEADS * HEAD_DIM, HIDDEN, device=get_accelerator().device_name()) - - -def _relative(a, b): - return ((a - b).norm() / b.norm()).item() - - -def _column_parallel_shards(tensor, tp): - rows = tensor.shape[0] // tp - return [tensor[r * rows:(r + 1) * rows].contiguous() for r in range(tp)] - - -@pytest.mark.parametrize("ns_method", ["gram", "standard"]) -@pytest.mark.parametrize("tp", [2, 4]) -def test_per_head_on_a_shard_is_the_shard_of_per_head(grad, ns_method, tp): - """Exactly equal, not close: the shards are the same matrices in the same batch.""" - whole = _per_head_orthogonalize(grad.clone(), HEADS, STEPS, ns_method) - sharded = torch.cat([ - _per_head_orthogonalize(shard.clone(), HEADS // tp, STEPS, ns_method) - for shard in _column_parallel_shards(grad, tp) - ]) - - assert torch.equal(sharded, whole) - - -@pytest.mark.parametrize("ns_method", ["gram", "standard"]) -def test_the_whole_matrix_path_does_not_survive_the_split(grad, ns_method): - """The comparison this is measured against, so "exact" above means something. - - Newton-Schulz on a block of rows is a different computation from the same block of - Newton-Schulz on every row, and the difference is not small. - """ - ns_fn = zeropower_via_gram_newtonschulz if ns_method == "gram" else zeropower_via_newtonschulz5 - whole = ns_fn(grad.clone(), steps=STEPS) - sharded = torch.cat([ns_fn(shard.clone(), steps=STEPS) for shard in _column_parallel_shards(grad, 2)]) - - assert _relative(sharded, whole) > 0.1 - - -def test_a_stale_head_count_splits_heads_in_half(grad): - """What the tag does under tp=2 if it is not re-resolved against the shard. - - 128 rows still divide by 8, so the divisibility check passes and each head is cut in two. - """ - whole = _per_head_orthogonalize(grad.clone(), HEADS, STEPS, "gram") - shard = _column_parallel_shards(grad, 2)[0] - - correct = _per_head_orthogonalize(shard.clone(), HEADS // 2, STEPS, "gram") - stale = _per_head_orthogonalize(shard.clone(), HEADS, STEPS, "gram") - - assert torch.equal(correct, whole[:shard.shape[0]]) - assert _relative(stale, whole[:shard.shape[0]]) > 0.1 - - -class TestPerHeadMuonUnderAutoTP(DistributedTest): - """The tags AutoTP leaves behind, through a real `deepspeed.initialize`. - - `set_optimizer_flags` runs before `_configure_tensor_parallel`, and AutoTP replaces the - parameter's `.data` in place, so the tag made against the whole model rides onto a shard. - """ - world_size = 2 - - def test_the_tags_describe_the_shard_and_not_the_model(self): - transformers = pytest.importorskip("transformers") - heads, head_dim = 8, 32 - config = transformers.LlamaConfig(hidden_size=heads * head_dim, - num_attention_heads=heads, - num_key_value_heads=heads, - num_hidden_layers=2, - intermediate_size=2 * heads * head_dim, - vocab_size=128) - model = transformers.AutoModelForCausalLM.from_config(config) - - engine, _, _, _ = deepspeed.initialize(model=model, - model_parameters=model.parameters(), - config={ - "train_micro_batch_size_per_gpu": 1, - "gradient_accumulation_steps": 1, - "bf16": { - "enabled": True - }, - "zero_optimization": { - "stage": 1 - }, - "tensor_parallel": { - "autotp_size": 2 - }, - "optimizer": { - "type": "Muon", - "params": { - "lr": 1e-3, - "per_head_muon": True - } - }, - }) - - tagged = {n: p for n, p in model.named_parameters() if getattr(p, "muon_num_heads", None)} - assert tagged, "AutoTP left nothing tagged" - for name, param in tagged.items(): - assert param.shape[0] // param.muon_num_heads == head_dim, \ - f"{name}: {param.shape[0]} rows over {param.muon_num_heads} heads is not {head_dim} wide" - assert param.muon_num_heads == heads // 2, \ - f"{name}: tp=2 leaves {heads // 2} heads on this rank, tagged {param.muon_num_heads}" - - # AutoTP asserts every rank in the TP group sees the same batch. - ids = torch.arange(8, device=engine.device).unsqueeze(0) % 128 - out = engine(input_ids=ids, labels=ids) - engine.backward(out.loss) - engine.step() - assert all(torch.isfinite(p).all() for p in model.parameters()) diff --git a/tests/unit/runtime/zero/test_per_head_muon.py b/tests/unit/runtime/zero/test_per_head_muon.py index 237004cf103e..8557ea0e5dca 100644 --- a/tests/unit/runtime/zero/test_per_head_muon.py +++ b/tests/unit/runtime/zero/test_per_head_muon.py @@ -1,18 +1,37 @@ +# Copyright (c) DeepSpeed Team. # SPDX-License-Identifier: Apache-2.0 # DeepSpeed Team """Per-head Muon: Newton-Schulz on each attention head instead of the whole projection. Full-matrix orthogonalization treats every head as one coupled block, so heads with larger -momentum dominate the shared update direction. Kimi K3 (arXiv:2607.24653 §2.5) and GLM-5 +momentum dominate the shared update direction. Kimi K3 (arXiv:2607.24653 5.2.5) and GLM-5 "Muon Split" (arXiv:2602.15763) both orthogonalize per head instead. See #8367. -CPU-only: these pin the arithmetic, not the accelerator path. +Three concerns, in order: + +1. **The arithmetic** -- per-head Newton-Schulz equals orthogonalizing each head alone. +2. **The tagging** -- which parameters get a head count, and what it is. `set_optimizer_flags` + already tags `use_muon` per parameter; head structure rides along the same way, so it does + not depend on AutoTP being enabled. +3. **Tensor parallelism** -- `set_optimizer_flags` runs before the engine partitions the model, + so the count it records counts the model's heads. Column-parallel TP splits attention + projections on dim 0, the axis the heads are on: the per-head width survives the split and + the count does not. Nothing catches that on its own, because the stale count usually still + divides the shard. + +CPU-only: these pin the arithmetic and the bookkeeping, not the accelerator path. The +multi-GPU side lives in `tests/unit/ops/muon/test_per_head_muon_accelerator.py`. """ +from types import SimpleNamespace + import pytest import torch +import deepspeed +from deepspeed import _attention_head_count, resolve_per_head_muon_after_sharding +from deepspeed.runtime.config import MUON_OPTIMIZER from deepspeed.runtime.zero.muon.original_muon import ( muon_update, zeropower_via_gram_newtonschulz, @@ -20,6 +39,10 @@ zeropower_via_newtonschulz5, ) +# --------------------------------------------------------------------------- +# 1. The arithmetic +# --------------------------------------------------------------------------- + def _ns_tolerance(ns_method): """A few ulps of whatever dtype the kernel iterates in. @@ -114,3 +137,522 @@ def test_rejects_shapes_that_do_not_split_into_heads(): conv_like = torch.randn(4, 4, 3, 3) with pytest.raises(ValueError, match="expects a 2D attention projection"): muon_update(conv_like.clone(), torch.zeros_like(conv_like), num_heads=4) + + +# --------------------------------------------------------------------------- +# 2. Which parameters are tagged, and with how many heads +# --------------------------------------------------------------------------- + + +class _Attn(torch.nn.Module): + + def __init__(self, hidden=64, q_heads=8, kv_heads=2, head_dim=8, fused=False): + super().__init__() + self.q_proj = torch.nn.Linear(hidden, q_heads * head_dim, bias=False) + self.k_proj = torch.nn.Linear(hidden, kv_heads * head_dim, bias=False) + self.v_proj = torch.nn.Linear(hidden, kv_heads * head_dim, bias=False) + self.o_proj = torch.nn.Linear(q_heads * head_dim, hidden, bias=False) + self.mlp = torch.nn.Linear(hidden, hidden, bias=False) + self.embed_tokens = torch.nn.Embedding(16, hidden) + if fused: + self.qkv_proj = torch.nn.Linear(hidden, (q_heads + 2 * kv_heads) * head_dim, bias=False) + # Real configs carry the per-head width, either as head_dim or derivably from + # hidden_size. Without one the shape cannot confirm the name, and the tagger declines. + self.config = SimpleNamespace(num_attention_heads=q_heads, + num_key_value_heads=kv_heads, + hidden_size=hidden, + head_dim=head_dim) + + +def _flags(model, per_head=True): + cfg = SimpleNamespace(optimizer_name=MUON_OPTIMIZER, + optimizer_params={"per_head_muon": per_head} if per_head else {}) + deepspeed.set_optimizer_flags(cfg, model) + return {name: getattr(p, "muon_num_heads", "MISSING") for name, p in model.named_parameters()} + + +def test_query_projection_uses_the_query_head_count(): + tags = _flags(_Attn(q_heads=8, kv_heads=2)) + + assert tags["q_proj.weight"] == 8 + + +def test_output_projection_is_left_alone(): + """o_proj is `[hidden, num_heads * head_dim]` - its heads are on the input axis. + + The split is on dim 0, so tagging it would cut across the wrong axis, and with the usual + hidden == num_heads * head_dim it still divides evenly, i.e. silently wrong rather than an + error. Regression test: it was tagged in the first version of this. + """ + tags = _flags(_Attn(q_heads=8, kv_heads=2)) + + assert tags["o_proj.weight"] is None + + +@pytest.mark.parametrize("mlp_name", [ + "intermediate.dense.weight", + "output.dense.weight", + "mlp.dense_h_to_4h.weight", + "mlp.dense_4h_to_h.weight", +]) +def test_mlp_matrices_named_dense_are_not_treated_as_attention(mlp_name): + """`dense` names an MLP matrix as often as an attention one. + + Matching it anywhere in the path tagged `intermediate.dense` and `dense_h_to_4h` with a head + count, splitting a matrix that has no head structure. Regression test: the first version of + this matched on the full path and did exactly that. + """ + from deepspeed import _attention_head_count + + model = _Attn(q_heads=8, kv_heads=2) + weight = torch.zeros(4 * 64, 64) + + assert _attention_head_count(f"encoder.layer.0.{mlp_name}", weight, model) is None + + +def test_kv_projections_use_the_kv_head_count_under_gqa(): + """K/V have fewer heads than Q under GQA, and splitting them by the query count would be wrong.""" + tags = _flags(_Attn(q_heads=8, kv_heads=2)) + + assert tags["k_proj.weight"] == 2 + assert tags["v_proj.weight"] == 2 + + +def test_non_attention_parameters_are_left_on_the_full_matrix_path(): + tags = _flags(_Attn()) + + assert tags["mlp.weight"] is None + assert tags["embed_tokens.weight"] is None + + +def test_fused_qkv_is_skipped(): + """One matrix holding Q, K and V does not split into uniform heads under GQA.""" + tags = _flags(_Attn(fused=True)) + + assert tags["qkv_proj.weight"] is None + + +def test_opt_in_is_required(): + tags = _flags(_Attn(), per_head=False) + + assert all(v is None for v in tags.values()), tags + + +def test_shape_that_does_not_divide_is_skipped(): + """A projection whose output dim is not a multiple of the head count is not that layout.""" + model = _Attn(q_heads=8, kv_heads=2) + model.q_proj = torch.nn.Linear(64, 63, bias=False) # 63 % 8 != 0 + + assert _flags(model)["q_proj.weight"] is None + + +def test_use_muon_tagging_is_unchanged(): + model = _Attn() + _flags(model) + + assert model.q_proj.weight.use_muon is True + assert model.embed_tokens.weight.use_muon is False + + +@pytest.mark.parametrize("q_heads,kv_heads", [(8, 8), (8, 1), (12, 4)]) +def test_head_counts_track_the_config(q_heads, kv_heads): + tags = _flags(_Attn(q_heads=q_heads, kv_heads=kv_heads, hidden=64, head_dim=8)) + + assert tags["q_proj.weight"] == q_heads + assert tags["k_proj.weight"] == kv_heads + + +@pytest.mark.parametrize("arch", ["llama", "qwen2", "mistral"]) +def test_split_qkv_architectures_tag_only_qkv(arch): + """Real HF configs rather than a stand-in, so the leaf names are the ones models actually use.""" + transformers = pytest.importorskip("transformers") + cfg_cls = { + "llama": transformers.LlamaConfig, + "qwen2": transformers.Qwen2Config, + "mistral": transformers.MistralConfig, + }[arch] + cfg = cfg_cls(hidden_size=64, + num_attention_heads=8, + num_key_value_heads=2, + num_hidden_layers=1, + intermediate_size=128, + vocab_size=32) + model = transformers.AutoModelForCausalLM.from_config(cfg) + + tags = {n.split(".")[-2]: _attention_head_count(n, p, model) for n, p in model.named_parameters() if p.ndim == 2} + + assert tags["q_proj"] == 8 + assert tags["k_proj"] == 2, "GQA: k/v are blocked by num_key_value_heads, not the query count" + assert tags["v_proj"] == 2 + assert tags["o_proj"] is None, "o_proj's heads are on the input axis" + for mlp_leaf in ("gate_proj", "up_proj", "down_proj"): + assert tags[mlp_leaf] is None, f"{mlp_leaf} has no head structure" + + +@pytest.mark.parametrize("arch", ["gpt_neox", "falcon"]) +def test_fused_qkv_architectures_tag_nothing(arch): + """These name their MLP matrices `dense_h_to_4h` / `dense_4h_to_h` and their output proj `dense`. + + Matching `dense` anywhere in the path tagged all three; this pins that none of them are. + """ + transformers = pytest.importorskip("transformers") + cfg_cls = {"gpt_neox": transformers.GPTNeoXConfig, "falcon": transformers.FalconConfig}[arch] + kwargs = dict(hidden_size=64, num_attention_heads=8, num_hidden_layers=1, vocab_size=32) + if arch == "gpt_neox": + kwargs["intermediate_size"] = 128 + model = transformers.AutoModelForCausalLM.from_config(cfg_cls(**kwargs)) + + tags = {n: _attention_head_count(n, p, model) for n, p in model.named_parameters() if p.ndim == 2} + + assert all(v is None for v in tags.values()), \ + {k: v for k, v in tags.items() if v is not None} + + +# Shapes and config values read off the real checkpoints delock pointed at in #8367: +# inference-optimization/GLM-5.2-0.8B-A0.8B and inference-optimization/Kimi-K3-0.40B. +def _glm52_mla_config(): + return SimpleNamespace(num_attention_heads=16, + num_key_value_heads=16, + hidden_size=2048, + head_dim=64, + q_lora_rank=512, + kv_lora_rank=128, + qk_nope_head_dim=192, + qk_rope_head_dim=64, + v_head_dim=128) + + +@pytest.mark.parametrize( + "leaf,shape,expected", + [ + ("q_b_proj", (4096, 512), 16), # 16 * (qk_nope 192 + qk_rope 64) + ("kv_b_proj", (5120, 128), 16), # 16 * (qk_nope 192 + v_head_dim 128) + ("q_a_proj", (512, 2048), None), # down-projection, no head structure + ("kv_a_proj_with_mqa", (192, 2048), None), # latent + rope, does not split into heads + ("o_proj", (2048, 2048), None), # heads on the input axis + ]) +def test_mla_tags_only_the_up_projections(leaf, shape, expected): + """MLA blocks its two up-projections by head, which is the split GLM-5's Muon Split applies. + + Their per-head width is not `head_dim`: q_b is qk_nope + qk_rope and kv_b is + qk_nope + v_head_dim, so a tagger that assumes `num_heads * head_dim` rejects both. + """ + model = SimpleNamespace(config=_glm52_mla_config()) + name = f"model.layers.0.self_attn.{leaf}.weight" + + assert _attention_head_count(name, torch.zeros(shape), model) == expected + + +def test_mla_head_width_is_checked_not_assumed(): + """A shape that does not equal num_heads * per-head width is not the layout we think it is.""" + model = SimpleNamespace(config=_glm52_mla_config()) + + ok = _attention_head_count("l.0.self_attn.q_b_proj.weight", torch.zeros(4096, 512), model) + wrong = _attention_head_count("l.0.self_attn.q_b_proj.weight", torch.zeros(4080, 512), model) + + assert ok == 16 + assert wrong is None + + +# Shapes measured by instantiating DeepseekV2Attention on the released +# deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct config under transformers 5.16.1. +def _deepseek_v2_lite_mla_config(): + return SimpleNamespace(num_attention_heads=16, + num_key_value_heads=16, + hidden_size=2048, + head_dim=64, + q_lora_rank=None, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=128) + + +@pytest.mark.parametrize( + "leaf,shape,expected", + [ + ("q_proj", (3072, 2048), 16), # 16 * (qk_nope 128 + qk_rope 64) + ("kv_b_proj", (4096, 512), 16), # 16 * (qk_nope 128 + v_head_dim 128) + ("kv_a_proj_with_mqa", (576, 2048), None), # kv_lora_rank + qk_rope, no head structure + ("o_proj", (2048, 2048), None), # heads on the input axis, and 2048 still divides by 16 + ]) +def test_mla_without_q_lora_rank_tags_the_plain_q_proj(leaf, shape, expected): + """Without a q_lora_rank there is no q_a/q_b pair; the up-projection is `q_proj` itself. + + Its per-head width stays `qk_nope + qk_rope`, so reading `head_dim` gives 16 * 64 = 1024 + against a real 3072 and drops the model off the per-head path. + """ + model = SimpleNamespace(config=_deepseek_v2_lite_mla_config()) + name = f"model.layers.0.self_attn.{leaf}.weight" + + assert _attention_head_count(name, torch.zeros(shape), model) == expected + + +def test_head_dim_alone_would_reject_the_mla_q_proj(): + """Guards the width source rather than the outcome: head_dim is present and wrong here.""" + config = _deepseek_v2_lite_mla_config() + + assert config.head_dim is not None + assert config.num_attention_heads * config.head_dim == 1024 + assert _attention_head_count("l.0.self_attn.q_proj.weight", torch.zeros(3072, 2048), + SimpleNamespace(config=config)) == 16 + + +def test_q_proj_on_a_non_mla_config_still_uses_head_dim(): + """The MLA width only applies where the config carries the MLA head dimensions.""" + config = SimpleNamespace(num_attention_heads=8, num_key_value_heads=8, hidden_size=512, head_dim=64) + model = SimpleNamespace(config=config) + + assert _attention_head_count("l.0.self_attn.q_proj.weight", torch.zeros(512, 512), model) == 8 + assert _attention_head_count("l.0.self_attn.q_proj.weight", torch.zeros(768, 512), model) is None + + +def test_linear_attention_named_like_standard_attention_is_rejected(): + """Kimi-K3-0.40B is `kimi_linear`, not MLA: q_proj is [256, 1024] with 8 heads of 74. + + The names match the standard-attention list, so only the shape check keeps it off the + per-head path. + """ + text = SimpleNamespace(num_attention_heads=8, + num_key_value_heads=8, + hidden_size=1024, + head_dim=74, + qk_nope_head_dim=64, + qk_rope_head_dim=32, + v_head_dim=64) + model = SimpleNamespace(config=SimpleNamespace(text_config=text)) + + for leaf in ("q_proj", "k_proj", "v_proj"): + name = f"model.layers.0.self_attn.{leaf}.weight" + assert _attention_head_count(name, torch.zeros(256, 1024), model) is None + + +def test_head_count_comes_from_the_shared_extractor(): + """Head counts are read through AutoTPMeta, so alternative config spellings work.""" + model = SimpleNamespace(config=SimpleNamespace(n_head=8, hidden_size=64, head_dim=8)) + + assert _attention_head_count("l.0.attn.q_proj.weight", torch.zeros(64, 64), model) == 8 + + +# --- candidate resolution ------------------------------------------------------ +# +# `q_proj` can be either the standard query projection or, on an MLA model without a +# q_lora_rank, the query up-projection. Both candidates are evaluated and the shape decides, +# so the outcome does not depend on which config fields happen to be present. + + +def _kimi_k3_hybrid_config(): + """Kimi-K3-0.40B: linear-attention layers on a config that also carries MLA fields. + + `linear_attn_config` gives `num_heads: 8, head_dim: 32`, so its `q_proj` is (256, 1024). + The top-level MLA dimensions belong to the model's two MLA layers, and `head_dim` is 74. + Neither top-level geometry describes the KDA projection. + """ + return SimpleNamespace(num_attention_heads=8, + num_key_value_heads=8, + hidden_size=1024, + head_dim=74, + qk_nope_head_dim=64, + qk_rope_head_dim=32, + v_head_dim=64) + + +@pytest.mark.parametrize("leaf", ["q_proj", "k_proj", "v_proj"]) +def test_linear_attention_on_a_config_with_mla_leftovers_is_declined(leaf): + """No candidate confirms, so it stays on the full-matrix path. + + 8 x (qk_nope 64 + qk_rope 32) = 768 and 8 x head_dim 74 = 592, against 256 rows. This is + the case tracked in #8420; until the linear-attention geometry is read, declining is the + correct outcome and it must come from the shape rather than from branch ordering. + """ + model = SimpleNamespace(config=_kimi_k3_hybrid_config()) + + assert _attention_head_count(f"model.layers.0.self_attn.{leaf}.weight", torch.zeros(256, 1024), model) is None + + +def test_two_candidates_agreeing_on_the_head_count_are_not_ambiguous(): + """Ambiguity is about the answer, not the route. + + The output is a head count, so two candidates that confirm with the same count give the + same answer and there is nothing to be ambiguous about. + """ + config = SimpleNamespace(num_attention_heads=8, + num_key_value_heads=8, + hidden_size=512, + head_dim=96, + qk_nope_head_dim=64, + qk_rope_head_dim=32, + v_head_dim=64) + model = SimpleNamespace(config=config) + + # 8 x 96 = 768 by head_dim, and 8 x (64 + 32) = 768 by the MLA width. + assert _attention_head_count("l.0.self_attn.q_proj.weight", torch.zeros(768, 512), model) == 8 + + +def test_candidates_that_disagree_on_the_head_count_are_skipped(): + """A real ambiguity: both confirm the shape, and they give different answers.""" + from deepspeed import _confirm + + candidates = [(8, 96, "mla-q"), (12, 64, "head-dim")] + num_heads, reason = _confirm(torch.zeros(768, 512), candidates) + + assert num_heads is None + assert reason.startswith("ambiguous:") + assert "mla-q=8" in reason and "head-dim=12" in reason + + +def test_a_config_without_a_per_head_width_is_declined(): + """Divisibility alone is not confirmation. + + `rows % heads == 0` holds for matrices that are not head-blocked at all, which is how + o_proj used to slip through. Without a width there is nothing to confirm against. + """ + config = SimpleNamespace(num_attention_heads=8, num_key_value_heads=8) + model = SimpleNamespace(config=config) + + assert _attention_head_count("l.0.self_attn.q_proj.weight", torch.zeros(512, 512), model) is None + + +def test_head_dim_is_derived_when_the_config_omits_it(): + """Configs that leave head_dim out still define it as hidden_size // num_attention_heads.""" + config = SimpleNamespace(num_attention_heads=8, num_key_value_heads=8, hidden_size=512) + model = SimpleNamespace(config=config) + + assert _attention_head_count("l.0.self_attn.q_proj.weight", torch.zeros(512, 512), model) == 8 + + +def test_the_flag_errors_rather_than_silently_doing_nothing(): + """An explicit opt-in that tags nothing is the tensor-parallel failure mode. + + Under TP the config describes the whole model while each rank holds a shard, so every + projection fails its width check and per-head is off model-wide while the user believes it + is on. There is no partial result to keep, so this is an error. + """ + + class _NoAttention(torch.nn.Module): + + def __init__(self): + super().__init__() + self.mlp = torch.nn.Linear(64, 64, bias=False) + self.config = SimpleNamespace(num_attention_heads=8, num_key_value_heads=8, hidden_size=64, head_dim=8) + + with pytest.raises(ValueError, match="no attention projection could be tagged"): + _flags(_NoAttention()) + + # ...and with the flag off it is simply not asked for. + assert _flags(_NoAttention(), per_head=False)["mlp.weight"] is None + + +def test_the_head_count_is_read_from_the_layer_shape_under_zero_init(): + """`zero.Init` leaves a flat placeholder and records the layer's shape as `ds_shape`. + + Reading `param.shape` there sees a 1-D tensor for every parameter, so nothing confirms and + the flag raises on a model it could describe perfectly well. + """ + model = _Attn(hidden=64, q_heads=8, kv_heads=2, head_dim=8) + for p in model.parameters(): + p.ds_shape = torch.Size(p.shape) + p.data = torch.zeros(0, dtype=p.dtype) + + tags = _flags(model) + + assert model.q_proj.weight.ndim == 1, "the partitioned parameter really is 1-D here" + assert tags["q_proj.weight"] == 8 + assert tags["k_proj.weight"] == 2 + assert tags["o_proj.weight"] is None + + +# --------------------------------------------------------------------------- +# 3. Tensor parallelism: re-resolving the count against the shard +# --------------------------------------------------------------------------- + + +class _ShardedAttn(torch.nn.Module): + """8 heads of 32, as the whole model sees it.""" + + def __init__(self, hidden=256, heads=8, head_dim=32): + super().__init__() + self.q_proj = torch.nn.Linear(hidden, heads * head_dim, bias=False) + self.k_proj = torch.nn.Linear(hidden, heads * head_dim, bias=False) + self.mlp = torch.nn.Linear(hidden, hidden, bias=False) + for p, num_heads in ((self.q_proj.weight, heads), (self.k_proj.weight, heads), (self.mlp.weight, None)): + p.muon_num_heads = num_heads + p.muon_head_dim = p.shape[0] // num_heads if num_heads else None + + def shard(self, tp: int, rows=None, leaves=("q_proj", "k_proj")): + """Replace the weights with column-parallel shards, as AutoTP's `_tp_partition` does.""" + for leaf in leaves: + weight = getattr(self, leaf).weight + keep = rows if rows is not None else weight.shape[0] // tp + weight.data = weight.data[:keep].clone() + + +def test_the_head_count_follows_the_shard(): + """tp=2 leaves 4 heads on this rank; the tag has to say 4, not the model's 8.""" + attn = _ShardedAttn() + attn.shard(tp=2) + + resolve_per_head_muon_after_sharding(attn) + + assert attn.q_proj.weight.shape == (128, 256) + assert attn.q_proj.weight.muon_num_heads == 4 + assert attn.k_proj.weight.muon_num_heads == 4 + + +def test_the_stale_count_would_have_split_heads_in_half(): + """Why this is not caught by the existing shape check: 128 % 8 == 0. + + The divisibility assert in `_per_head_orthogonalize` passes on the stale count, so without + this pass Newton-Schulz runs on 8 blocks of 16 - half of each head - and says nothing. + """ + attn = _ShardedAttn() + attn.shard(tp=2) + + rows, stale = attn.q_proj.weight.shape[0], 8 + assert rows % stale == 0, "the stale count divides, which is why it needs correcting rather than asserting" + assert rows // stale == 16 != 32 + + +def test_a_shard_that_splits_a_head_is_dropped(): + """A shard that does not hold whole heads has no per-head structure to use.""" + attn = _ShardedAttn() + attn.shard(tp=2, rows=144, leaves=("q_proj", )) # 4.5 heads of 32 + + resolve_per_head_muon_after_sharding(attn) + + assert attn.q_proj.weight.muon_num_heads is None + assert attn.k_proj.weight.muon_num_heads == 8, "one bad shard does not turn the feature off elsewhere" + + +def test_an_unsharded_model_keeps_its_count(): + attn = _ShardedAttn() + + resolve_per_head_muon_after_sharding(attn) + + assert attn.q_proj.weight.muon_num_heads == 8 + assert attn.q_proj.weight.shape[0] == 256 + + +def test_untagged_parameters_are_left_alone(): + attn = _ShardedAttn() + attn.shard(tp=2) + + resolve_per_head_muon_after_sharding(attn) + + assert getattr(attn.mlp.weight, "muon_num_heads", "MISSING") is None + + +def test_a_model_with_no_tags_at_all_is_a_no_op(): + """Muon without `per_head_muon`, and every non-Muon model: nothing to resolve, no error.""" + model = torch.nn.Linear(8, 8, bias=False) + + resolve_per_head_muon_after_sharding(model) # must not raise + + +def test_it_raises_when_the_flag_ends_up_doing_nothing(): + """Same contract as the tagging pass: an opt-in that silently does nothing is the failure.""" + attn = _ShardedAttn() + attn.shard(tp=2, rows=144) + + with pytest.raises(ValueError, match="sharded across head boundaries"): + resolve_per_head_muon_after_sharding(attn) diff --git a/tests/unit/runtime/zero/test_per_head_muon_tagging.py b/tests/unit/runtime/zero/test_per_head_muon_tagging.py deleted file mode 100644 index cacae1c6a6d5..000000000000 --- a/tests/unit/runtime/zero/test_per_head_muon_tagging.py +++ /dev/null @@ -1,435 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -"""Which parameters per-head Muon tags, and with how many heads. See #8367. - -`set_optimizer_flags` already tags `use_muon` per parameter; head structure rides along the -same way so it does not depend on AutoTP being enabled. CPU-only. -""" - -from types import SimpleNamespace - -import pytest -import torch - -import deepspeed -from deepspeed import _attention_head_count -from deepspeed.runtime.config import MUON_OPTIMIZER - - -class _Attn(torch.nn.Module): - - def __init__(self, hidden=64, q_heads=8, kv_heads=2, head_dim=8, fused=False): - super().__init__() - self.q_proj = torch.nn.Linear(hidden, q_heads * head_dim, bias=False) - self.k_proj = torch.nn.Linear(hidden, kv_heads * head_dim, bias=False) - self.v_proj = torch.nn.Linear(hidden, kv_heads * head_dim, bias=False) - self.o_proj = torch.nn.Linear(q_heads * head_dim, hidden, bias=False) - self.mlp = torch.nn.Linear(hidden, hidden, bias=False) - self.embed_tokens = torch.nn.Embedding(16, hidden) - if fused: - self.qkv_proj = torch.nn.Linear(hidden, (q_heads + 2 * kv_heads) * head_dim, bias=False) - # Real configs carry the per-head width, either as head_dim or derivably from - # hidden_size. Without one the shape cannot confirm the name, and the tagger declines. - self.config = SimpleNamespace(num_attention_heads=q_heads, - num_key_value_heads=kv_heads, - hidden_size=hidden, - head_dim=head_dim) - - -def _flags(model, per_head=True): - cfg = SimpleNamespace(optimizer_name=MUON_OPTIMIZER, - optimizer_params={"per_head_muon": per_head} if per_head else {}) - deepspeed.set_optimizer_flags(cfg, model) - return {name: getattr(p, "muon_num_heads", "MISSING") for name, p in model.named_parameters()} - - -def test_query_projection_uses_the_query_head_count(): - tags = _flags(_Attn(q_heads=8, kv_heads=2)) - - assert tags["q_proj.weight"] == 8 - - -def test_output_projection_is_left_alone(): - """o_proj is `[hidden, num_heads * head_dim]` - its heads are on the input axis. - - The split is on dim 0, so tagging it would cut across the wrong axis, and with the usual - hidden == num_heads * head_dim it still divides evenly, i.e. silently wrong rather than an - error. Regression test: it was tagged in the first version of this. - """ - tags = _flags(_Attn(q_heads=8, kv_heads=2)) - - assert tags["o_proj.weight"] is None - - -@pytest.mark.parametrize("mlp_name", [ - "intermediate.dense.weight", - "output.dense.weight", - "mlp.dense_h_to_4h.weight", - "mlp.dense_4h_to_h.weight", -]) -def test_mlp_matrices_named_dense_are_not_treated_as_attention(mlp_name): - """`dense` names an MLP matrix as often as an attention one. - - Matching it anywhere in the path tagged `intermediate.dense` and `dense_h_to_4h` with a head - count, splitting a matrix that has no head structure. Regression test: the first version of - this matched on the full path and did exactly that. - """ - from deepspeed import _attention_head_count - - model = _Attn(q_heads=8, kv_heads=2) - weight = torch.zeros(4 * 64, 64) - - assert _attention_head_count(f"encoder.layer.0.{mlp_name}", weight, model) is None - - -def test_kv_projections_use_the_kv_head_count_under_gqa(): - """K/V have fewer heads than Q under GQA, and splitting them by the query count would be wrong.""" - tags = _flags(_Attn(q_heads=8, kv_heads=2)) - - assert tags["k_proj.weight"] == 2 - assert tags["v_proj.weight"] == 2 - - -def test_non_attention_parameters_are_left_on_the_full_matrix_path(): - tags = _flags(_Attn()) - - assert tags["mlp.weight"] is None - assert tags["embed_tokens.weight"] is None - - -def test_fused_qkv_is_skipped(): - """One matrix holding Q, K and V does not split into uniform heads under GQA.""" - tags = _flags(_Attn(fused=True)) - - assert tags["qkv_proj.weight"] is None - - -def test_opt_in_is_required(): - tags = _flags(_Attn(), per_head=False) - - assert all(v is None for v in tags.values()), tags - - -def test_shape_that_does_not_divide_is_skipped(): - """A projection whose output dim is not a multiple of the head count is not that layout.""" - model = _Attn(q_heads=8, kv_heads=2) - model.q_proj = torch.nn.Linear(64, 63, bias=False) # 63 % 8 != 0 - - assert _flags(model)["q_proj.weight"] is None - - -def test_use_muon_tagging_is_unchanged(): - model = _Attn() - _flags(model) - - assert model.q_proj.weight.use_muon is True - assert model.embed_tokens.weight.use_muon is False - - -@pytest.mark.parametrize("q_heads,kv_heads", [(8, 8), (8, 1), (12, 4)]) -def test_head_counts_track_the_config(q_heads, kv_heads): - tags = _flags(_Attn(q_heads=q_heads, kv_heads=kv_heads, hidden=64, head_dim=8)) - - assert tags["q_proj.weight"] == q_heads - assert tags["k_proj.weight"] == kv_heads - - -@pytest.mark.parametrize("arch", ["llama", "qwen2", "mistral"]) -def test_split_qkv_architectures_tag_only_qkv(arch): - """Real HF configs rather than a stand-in, so the leaf names are the ones models actually use.""" - transformers = pytest.importorskip("transformers") - cfg_cls = { - "llama": transformers.LlamaConfig, - "qwen2": transformers.Qwen2Config, - "mistral": transformers.MistralConfig, - }[arch] - cfg = cfg_cls(hidden_size=64, - num_attention_heads=8, - num_key_value_heads=2, - num_hidden_layers=1, - intermediate_size=128, - vocab_size=32) - model = transformers.AutoModelForCausalLM.from_config(cfg) - - tags = {n.split(".")[-2]: _attention_head_count(n, p, model) for n, p in model.named_parameters() if p.ndim == 2} - - assert tags["q_proj"] == 8 - assert tags["k_proj"] == 2, "GQA: k/v are blocked by num_key_value_heads, not the query count" - assert tags["v_proj"] == 2 - assert tags["o_proj"] is None, "o_proj's heads are on the input axis" - for mlp_leaf in ("gate_proj", "up_proj", "down_proj"): - assert tags[mlp_leaf] is None, f"{mlp_leaf} has no head structure" - - -@pytest.mark.parametrize("arch", ["gpt_neox", "falcon"]) -def test_fused_qkv_architectures_tag_nothing(arch): - """These name their MLP matrices `dense_h_to_4h` / `dense_4h_to_h` and their output proj `dense`. - - Matching `dense` anywhere in the path tagged all three; this pins that none of them are. - """ - transformers = pytest.importorskip("transformers") - cfg_cls = {"gpt_neox": transformers.GPTNeoXConfig, "falcon": transformers.FalconConfig}[arch] - kwargs = dict(hidden_size=64, num_attention_heads=8, num_hidden_layers=1, vocab_size=32) - if arch == "gpt_neox": - kwargs["intermediate_size"] = 128 - model = transformers.AutoModelForCausalLM.from_config(cfg_cls(**kwargs)) - - tags = {n: _attention_head_count(n, p, model) for n, p in model.named_parameters() if p.ndim == 2} - - assert all(v is None for v in tags.values()), \ - {k: v for k, v in tags.items() if v is not None} - - -# Shapes and config values read off the real checkpoints delock pointed at in #8367: -# inference-optimization/GLM-5.2-0.8B-A0.8B and inference-optimization/Kimi-K3-0.40B. -def _glm52_mla_config(): - return SimpleNamespace(num_attention_heads=16, - num_key_value_heads=16, - hidden_size=2048, - head_dim=64, - q_lora_rank=512, - kv_lora_rank=128, - qk_nope_head_dim=192, - qk_rope_head_dim=64, - v_head_dim=128) - - -@pytest.mark.parametrize( - "leaf,shape,expected", - [ - ("q_b_proj", (4096, 512), 16), # 16 * (qk_nope 192 + qk_rope 64) - ("kv_b_proj", (5120, 128), 16), # 16 * (qk_nope 192 + v_head_dim 128) - ("q_a_proj", (512, 2048), None), # down-projection, no head structure - ("kv_a_proj_with_mqa", (192, 2048), None), # latent + rope, does not split into heads - ("o_proj", (2048, 2048), None), # heads on the input axis - ]) -def test_mla_tags_only_the_up_projections(leaf, shape, expected): - """MLA blocks its two up-projections by head, which is the split GLM-5's Muon Split applies. - - Their per-head width is not `head_dim`: q_b is qk_nope + qk_rope and kv_b is - qk_nope + v_head_dim, so a tagger that assumes `num_heads * head_dim` rejects both. - """ - model = SimpleNamespace(config=_glm52_mla_config()) - name = f"model.layers.0.self_attn.{leaf}.weight" - - assert _attention_head_count(name, torch.zeros(shape), model) == expected - - -def test_mla_head_width_is_checked_not_assumed(): - """A shape that does not equal num_heads * per-head width is not the layout we think it is.""" - model = SimpleNamespace(config=_glm52_mla_config()) - - ok = _attention_head_count("l.0.self_attn.q_b_proj.weight", torch.zeros(4096, 512), model) - wrong = _attention_head_count("l.0.self_attn.q_b_proj.weight", torch.zeros(4080, 512), model) - - assert ok == 16 - assert wrong is None - - -# Shapes measured by instantiating DeepseekV2Attention on the released -# deepseek-ai/DeepSeek-Coder-V2-Lite-Instruct config under transformers 5.16.1. -def _deepseek_v2_lite_mla_config(): - return SimpleNamespace(num_attention_heads=16, - num_key_value_heads=16, - hidden_size=2048, - head_dim=64, - q_lora_rank=None, - kv_lora_rank=512, - qk_nope_head_dim=128, - qk_rope_head_dim=64, - v_head_dim=128) - - -@pytest.mark.parametrize( - "leaf,shape,expected", - [ - ("q_proj", (3072, 2048), 16), # 16 * (qk_nope 128 + qk_rope 64) - ("kv_b_proj", (4096, 512), 16), # 16 * (qk_nope 128 + v_head_dim 128) - ("kv_a_proj_with_mqa", (576, 2048), None), # kv_lora_rank + qk_rope, no head structure - ("o_proj", (2048, 2048), None), # heads on the input axis, and 2048 still divides by 16 - ]) -def test_mla_without_q_lora_rank_tags_the_plain_q_proj(leaf, shape, expected): - """Without a q_lora_rank there is no q_a/q_b pair; the up-projection is `q_proj` itself. - - Its per-head width stays `qk_nope + qk_rope`, so reading `head_dim` gives 16 * 64 = 1024 - against a real 3072 and drops the model off the per-head path. - """ - model = SimpleNamespace(config=_deepseek_v2_lite_mla_config()) - name = f"model.layers.0.self_attn.{leaf}.weight" - - assert _attention_head_count(name, torch.zeros(shape), model) == expected - - -def test_head_dim_alone_would_reject_the_mla_q_proj(): - """Guards the width source rather than the outcome: head_dim is present and wrong here.""" - config = _deepseek_v2_lite_mla_config() - - assert config.head_dim is not None - assert config.num_attention_heads * config.head_dim == 1024 - assert _attention_head_count("l.0.self_attn.q_proj.weight", torch.zeros(3072, 2048), - SimpleNamespace(config=config)) == 16 - - -def test_q_proj_on_a_non_mla_config_still_uses_head_dim(): - """The MLA width only applies where the config carries the MLA head dimensions.""" - config = SimpleNamespace(num_attention_heads=8, num_key_value_heads=8, hidden_size=512, head_dim=64) - model = SimpleNamespace(config=config) - - assert _attention_head_count("l.0.self_attn.q_proj.weight", torch.zeros(512, 512), model) == 8 - assert _attention_head_count("l.0.self_attn.q_proj.weight", torch.zeros(768, 512), model) is None - - -def test_linear_attention_named_like_standard_attention_is_rejected(): - """Kimi-K3-0.40B is `kimi_linear`, not MLA: q_proj is [256, 1024] with 8 heads of 74. - - The names match the standard-attention list, so only the shape check keeps it off the - per-head path. - """ - text = SimpleNamespace(num_attention_heads=8, - num_key_value_heads=8, - hidden_size=1024, - head_dim=74, - qk_nope_head_dim=64, - qk_rope_head_dim=32, - v_head_dim=64) - model = SimpleNamespace(config=SimpleNamespace(text_config=text)) - - for leaf in ("q_proj", "k_proj", "v_proj"): - name = f"model.layers.0.self_attn.{leaf}.weight" - assert _attention_head_count(name, torch.zeros(256, 1024), model) is None - - -def test_head_count_comes_from_the_shared_extractor(): - """Head counts are read through AutoTPMeta, so alternative config spellings work.""" - model = SimpleNamespace(config=SimpleNamespace(n_head=8, hidden_size=64, head_dim=8)) - - assert _attention_head_count("l.0.attn.q_proj.weight", torch.zeros(64, 64), model) == 8 - - -# --- candidate resolution ------------------------------------------------------ -# -# `q_proj` can be either the standard query projection or, on an MLA model without a -# q_lora_rank, the query up-projection. Both candidates are evaluated and the shape decides, -# so the outcome does not depend on which config fields happen to be present. - - -def _kimi_k3_hybrid_config(): - """Kimi-K3-0.40B: linear-attention layers on a config that also carries MLA fields. - - `linear_attn_config` gives `num_heads: 8, head_dim: 32`, so its `q_proj` is (256, 1024). - The top-level MLA dimensions belong to the model's two MLA layers, and `head_dim` is 74. - Neither top-level geometry describes the KDA projection. - """ - return SimpleNamespace(num_attention_heads=8, - num_key_value_heads=8, - hidden_size=1024, - head_dim=74, - qk_nope_head_dim=64, - qk_rope_head_dim=32, - v_head_dim=64) - - -@pytest.mark.parametrize("leaf", ["q_proj", "k_proj", "v_proj"]) -def test_linear_attention_on_a_config_with_mla_leftovers_is_declined(leaf): - """No candidate confirms, so it stays on the full-matrix path. - - 8 x (qk_nope 64 + qk_rope 32) = 768 and 8 x head_dim 74 = 592, against 256 rows. This is - the case tracked in #8420; until the linear-attention geometry is read, declining is the - correct outcome and it must come from the shape rather than from branch ordering. - """ - model = SimpleNamespace(config=_kimi_k3_hybrid_config()) - - assert _attention_head_count(f"model.layers.0.self_attn.{leaf}.weight", torch.zeros(256, 1024), model) is None - - -def test_two_candidates_agreeing_on_the_head_count_are_not_ambiguous(): - """Ambiguity is about the answer, not the route. - - The output is a head count, so two candidates that confirm with the same count give the - same answer and there is nothing to be ambiguous about. - """ - config = SimpleNamespace(num_attention_heads=8, - num_key_value_heads=8, - hidden_size=512, - head_dim=96, - qk_nope_head_dim=64, - qk_rope_head_dim=32, - v_head_dim=64) - model = SimpleNamespace(config=config) - - # 8 x 96 = 768 by head_dim, and 8 x (64 + 32) = 768 by the MLA width. - assert _attention_head_count("l.0.self_attn.q_proj.weight", torch.zeros(768, 512), model) == 8 - - -def test_candidates_that_disagree_on_the_head_count_are_skipped(): - """A real ambiguity: both confirm the shape, and they give different answers.""" - from deepspeed import _confirm - - candidates = [(8, 96, "mla-q"), (12, 64, "head-dim")] - num_heads, reason = _confirm(torch.zeros(768, 512), candidates) - - assert num_heads is None - assert reason.startswith("ambiguous:") - assert "mla-q=8" in reason and "head-dim=12" in reason - - -def test_a_config_without_a_per_head_width_is_declined(): - """Divisibility alone is not confirmation. - - `rows % heads == 0` holds for matrices that are not head-blocked at all, which is how - o_proj used to slip through. Without a width there is nothing to confirm against. - """ - config = SimpleNamespace(num_attention_heads=8, num_key_value_heads=8) - model = SimpleNamespace(config=config) - - assert _attention_head_count("l.0.self_attn.q_proj.weight", torch.zeros(512, 512), model) is None - - -def test_head_dim_is_derived_when_the_config_omits_it(): - """Configs that leave head_dim out still define it as hidden_size // num_attention_heads.""" - config = SimpleNamespace(num_attention_heads=8, num_key_value_heads=8, hidden_size=512) - model = SimpleNamespace(config=config) - - assert _attention_head_count("l.0.self_attn.q_proj.weight", torch.zeros(512, 512), model) == 8 - - -def test_the_flag_errors_rather_than_silently_doing_nothing(): - """An explicit opt-in that tags nothing is the tensor-parallel failure mode. - - Under TP the config describes the whole model while each rank holds a shard, so every - projection fails its width check and per-head is off model-wide while the user believes it - is on. There is no partial result to keep, so this is an error. - """ - - class _NoAttention(torch.nn.Module): - - def __init__(self): - super().__init__() - self.mlp = torch.nn.Linear(64, 64, bias=False) - self.config = SimpleNamespace(num_attention_heads=8, num_key_value_heads=8, hidden_size=64, head_dim=8) - - with pytest.raises(ValueError, match="no attention projection could be tagged"): - _flags(_NoAttention()) - - # ...and with the flag off it is simply not asked for. - assert _flags(_NoAttention(), per_head=False)["mlp.weight"] is None - - -def test_the_head_count_is_read_from_the_layer_shape_under_zero_init(): - """`zero.Init` leaves a flat placeholder and records the layer's shape as `ds_shape`. - - Reading `param.shape` there sees a 1-D tensor for every parameter, so nothing confirms and - the flag raises on a model it could describe perfectly well. - """ - model = _Attn(hidden=64, q_heads=8, kv_heads=2, head_dim=8) - for p in model.parameters(): - p.ds_shape = torch.Size(p.shape) - p.data = torch.zeros(0, dtype=p.dtype) - - tags = _flags(model) - - assert model.q_proj.weight.ndim == 1, "the partitioned parameter really is 1-D here" - assert tags["q_proj.weight"] == 8 - assert tags["k_proj.weight"] == 2 - assert tags["o_proj.weight"] is None diff --git a/tests/unit/runtime/zero/test_per_head_muon_tensor_parallel.py b/tests/unit/runtime/zero/test_per_head_muon_tensor_parallel.py deleted file mode 100644 index cb12d608bade..000000000000 --- a/tests/unit/runtime/zero/test_per_head_muon_tensor_parallel.py +++ /dev/null @@ -1,106 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 - -# DeepSpeed Team -"""Per-head Muon under tensor parallelism. See #8367. - -`set_optimizer_flags` runs before the engine partitions the model, so the head count it records -counts the model's heads. Column-parallel TP then splits attention projections on dim 0, which -is the axis the heads are on: the per-head width survives the split and the count does not. -Nothing catches that on its own, because the stale count usually still divides the shard. -""" - -import pytest -import torch - -from deepspeed import resolve_per_head_muon_after_sharding - - -class _Attn(torch.nn.Module): - """8 heads of 32, as the whole model sees it.""" - - def __init__(self, hidden=256, heads=8, head_dim=32): - super().__init__() - self.q_proj = torch.nn.Linear(hidden, heads * head_dim, bias=False) - self.k_proj = torch.nn.Linear(hidden, heads * head_dim, bias=False) - self.mlp = torch.nn.Linear(hidden, hidden, bias=False) - for p, num_heads in ((self.q_proj.weight, heads), (self.k_proj.weight, heads), (self.mlp.weight, None)): - p.muon_num_heads = num_heads - p.muon_head_dim = p.shape[0] // num_heads if num_heads else None - - def shard(self, tp: int, rows=None, leaves=("q_proj", "k_proj")): - """Replace the weights with column-parallel shards, as AutoTP's `_tp_partition` does.""" - for leaf in leaves: - weight = getattr(self, leaf).weight - keep = rows if rows is not None else weight.shape[0] // tp - weight.data = weight.data[:keep].clone() - - -def test_the_head_count_follows_the_shard(): - """tp=2 leaves 4 heads on this rank; the tag has to say 4, not the model's 8.""" - attn = _Attn() - attn.shard(tp=2) - - resolve_per_head_muon_after_sharding(attn) - - assert attn.q_proj.weight.shape == (128, 256) - assert attn.q_proj.weight.muon_num_heads == 4 - assert attn.k_proj.weight.muon_num_heads == 4 - - -def test_the_stale_count_would_have_split_heads_in_half(): - """Why this is not caught by the existing shape check: 128 % 8 == 0. - - The divisibility assert in `_per_head_orthogonalize` passes on the stale count, so without - this pass Newton-Schulz runs on 8 blocks of 16 - half of each head - and says nothing. - """ - attn = _Attn() - attn.shard(tp=2) - - rows, stale = attn.q_proj.weight.shape[0], 8 - assert rows % stale == 0, "the stale count divides, which is why it needs correcting rather than asserting" - assert rows // stale == 16 != 32 - - -def test_a_shard_that_splits_a_head_is_dropped(): - """A shard that does not hold whole heads has no per-head structure to use.""" - attn = _Attn() - attn.shard(tp=2, rows=144, leaves=("q_proj", )) # 4.5 heads of 32 - - resolve_per_head_muon_after_sharding(attn) - - assert attn.q_proj.weight.muon_num_heads is None - assert attn.k_proj.weight.muon_num_heads == 8, "one bad shard does not turn the feature off elsewhere" - - -def test_an_unsharded_model_keeps_its_count(): - attn = _Attn() - - resolve_per_head_muon_after_sharding(attn) - - assert attn.q_proj.weight.muon_num_heads == 8 - assert attn.q_proj.weight.shape[0] == 256 - - -def test_untagged_parameters_are_left_alone(): - attn = _Attn() - attn.shard(tp=2) - - resolve_per_head_muon_after_sharding(attn) - - assert getattr(attn.mlp.weight, "muon_num_heads", "MISSING") is None - - -def test_a_model_with_no_tags_at_all_is_a_no_op(): - """Muon without `per_head_muon`, and every non-Muon model: nothing to resolve, no error.""" - model = torch.nn.Linear(8, 8, bias=False) - - resolve_per_head_muon_after_sharding(model) # must not raise - - -def test_it_raises_when_the_flag_ends_up_doing_nothing(): - """Same contract as the tagging pass: an opt-in that silently does nothing is the failure.""" - attn = _Attn() - attn.shard(tp=2, rows=144) - - with pytest.raises(ValueError, match="sharded across head boundaries"): - resolve_per_head_muon_after_sharding(attn)