diff --git a/deepspeed/__init__.py b/deepspeed/__init__.py index 26bfc7c77df6..dc1460c5bfd6 100755 --- a/deepspeed/__init__.py +++ b/deepspeed/__init__.py @@ -80,15 +80,278 @@ def _parse_version(version_str): # Set to torch's distributed package or deepspeed.comm based inside DeepSpeedEngine init dist = None +# 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") +# 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 _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. + + 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`. + + 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 _classify_leaf(leaf: str): + """Which kind of attention matrix this leaf name claims to be, or None if it claims none. + + A name is a claim, not a layout. What the leaf resolves to is decided later, by the shape. + """ + 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): + 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 _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 + + +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 + num_kv_heads = meta.num_kv_heads or num_attention_heads + 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) + + 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] + + +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. + """ + shape = _layer_shape(param) + if len(shape) != 2: + return None, "not-2d" + if not candidates: + return None, "no-candidate-geometry" + + rows = 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: 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 + 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")): + 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) + 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 {_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", _layer_shape(p)[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 = _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}" + 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, diff --git a/deepspeed/runtime/engine.py b/deepspeed/runtime/engine.py index 86918bd71c5a..f72bc57f8c6f 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/deepspeed/runtime/zero/muon/original_muon.py b/deepspeed/runtime/zero/muon/original_muon.py index 1cbc46392410..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 @@ -142,11 +153,47 @@ def zeropower_via_gram_newtonschulz(G, steps: int): NS_METHODS = {"standard", "gram"} +def _per_head_orthogonalize(update, num_heads, ns_steps, ns_method): + """Newton-Schulz per attention head, then fold the head dim back.""" + if update.ndim != 2: + raise ValueError(f"Per-head Muon expects a 2D attention projection, got shape {tuple(update.shape)}.") + + out_features, in_features = update.shape + if num_heads < 1 or out_features % num_heads != 0: + raise ValueError(f"Per-head Muon needs the output dim to split evenly across heads, but " + f"{out_features} is not divisible by num_heads={num_heads}.") + + head_dim = out_features // num_heads + ns_fn = zeropower_via_gram_newtonschulz if ns_method == "gram" else zeropower_via_newtonschulz5 + # Scale per head block, matching what the full-matrix path does for the whole matrix. + scale = max(1, head_dim / in_features)**0.5 + per_head = ns_fn(update.view(num_heads, head_dim, in_features), steps=ns_steps) * scale + + return per_head.reshape(out_features, in_features) + + @compiler.compile() -def muon_update(grad, momentum, beta=0.95, ns_steps=5, nesterov=True, ns_method="gram", is_expert_group=False): +def muon_update(grad, + momentum, + beta=0.95, + ns_steps=5, + nesterov=True, + ns_method="gram", + is_expert_group=False, + num_heads=None): + """Muon update, optionally orthogonalizing each attention head separately. + + With ``num_heads`` set, the update for an attention projection of shape + ``[num_heads * head_dim, in_features]`` is viewed as ``[num_heads, head_dim, in_features]`` + and Newton-Schulz runs on that batch, so each head is orthogonalized against itself instead + of sharing one update direction with every other head. Both NS kernels are already batched, + so this is the same path the expert-group branch takes. + """ orig_dtype = grad.dtype momentum.lerp_(grad, 1 - beta) update = grad.lerp_(momentum, beta) if nesterov else momentum + if num_heads is not None: + return _per_head_orthogonalize(update, num_heads, ns_steps, ns_method).to(orig_dtype) if is_expert_group: ns_fn = zeropower_via_gram_newtonschulz if ns_method == "gram" else zeropower_via_newtonschulz5 scale = max(1, update.size(-2) / update.size(-1))**0.5 @@ -219,7 +266,8 @@ def step(self, closure=None): state["momentum_buffer"], beta=group["momentum"], ns_method=group.get("ns_method", "gram"), - is_expert_group=getattr(p, 'is_expert_group', False)) + is_expert_group=getattr(p, 'is_expert_group', False), + num_heads=getattr(p, 'muon_num_heads', None)) p.mul_(1 - group["lr"] * group["weight_decay"]) p.add_(update.reshape(p.shape), alpha=-group["lr"]) dist.all_gather(params_pad[base_i:base_i + dist.get_world_size()], @@ -257,7 +305,8 @@ def step(self, closure=None): state["momentum_buffer"], beta=group["momentum"], ns_method=group.get("ns_method", "gram"), - is_expert_group=getattr(p, 'is_expert_group', False)) + is_expert_group=getattr(p, 'is_expert_group', False), + num_heads=getattr(p, 'muon_num_heads', None)) p.mul_(1 - group["lr"] * group["weight_decay"]) p.add_(update.reshape(p.shape), alpha=-group["lr"]) @@ -349,7 +398,8 @@ def step(self, closure=None): state["momentum_buffer"], beta=group["momentum"], ns_method=group.get("ns_method", "gram"), - is_expert_group=getattr(p, 'is_expert_group', False)) + is_expert_group=getattr(p, 'is_expert_group', False), + num_heads=getattr(p, 'muon_num_heads', None)) p.mul_(1 - group["lr"] * group["weight_decay"]) p.add_(update.reshape(p.shape), alpha=-group["lr"]) dist.all_gather(params_pad[base_i:base_i + dist.get_world_size()], @@ -421,7 +471,8 @@ def step(self, closure=None): state["momentum_buffer"], beta=group["momentum"], ns_method=group.get("ns_method", "gram"), - is_expert_group=getattr(p, 'is_expert_group', False)) + is_expert_group=getattr(p, 'is_expert_group', False), + num_heads=getattr(p, 'muon_num_heads', None)) p.mul_(1 - group["lr"] * group["weight_decay"]) p.add_(update.reshape(p.shape), alpha=-group["lr"]) else: 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/docs/_pages/config-json.md b/docs/_pages/config-json.md index cd1c2def9562..141d1becd859 100644 --- a/docs/_pages/config-json.md +++ b/docs/_pages/config-json.md @@ -60,6 +60,47 @@ 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. + +**Tensor parallelism.** Column-parallel TP splits an attention projection on dim 0, which is +the axis the heads are on, so a rank holds whole heads and the per-head width is unchanged. That +makes the per-head split exact under TP: Newton-Schulz on a rank's heads is the same computation +whether the other ranks' heads are present or not. The head *count* is not invariant, so with +AutoTP the counts are re-resolved against the shards after partitioning; a shard whose rows are +not a multiple of the per-head width does not hold whole heads and stays on the full-matrix path. +A model that arrives already sharded by an external tensor-parallel implementation cannot be +tagged at all, because the config then describes a width no parameter has. + +**The flag reports what it did.** Because it is an explicit opt-in, DeepSpeed raises at +`deepspeed.initialize` if it is enabled and no attention projection could be tagged, rather than +training on without it. Parameters that match an attention name but confirm no geometry are +reported as a warning and stay on the full-matrix path, so a hybrid model still gets per-head on +its recognized layers. By default, non-Muon parameters use `FusedAdam`. When optimizer state is offloaded to the CPU, DeepSpeed selects `DeepSpeedCPUAdam`. This is the same backend selection used by the Adam and AdamW optimizer types. 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/runtime/zero/test_per_head_muon.py b/tests/unit/runtime/zero/test_per_head_muon.py new file mode 100644 index 000000000000..8557ea0e5dca --- /dev/null +++ b/tests/unit/runtime/zero/test_per_head_muon.py @@ -0,0 +1,658 @@ +# 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 5.2.5) and GLM-5 +"Muon Split" (arXiv:2602.15763) both orthogonalize per head instead. See #8367. + +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, + ns_compute_dtype, + zeropower_via_newtonschulz5, +) + +# --------------------------------------------------------------------------- +# 1. The arithmetic +# --------------------------------------------------------------------------- + + +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. + """ + 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%.""" + return max(1e-2, 2 * torch.finfo(ns_compute_dtype(ns_method)).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) + + +# --------------------------------------------------------------------------- +# 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)