From 6473dadf8b5319ef18a771d5fffcfd9e418de7e8 Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Fri, 11 Sep 2026 21:24:46 +0800 Subject: [PATCH 1/2] feat(gpu): headwise Muon, Sinkhorn-balanced update, DeepSeek probes TODO.impl/03, /06 and the two measurement probes from the DeepSeek-V4.1-Flash learnings pass: - muon: headwise group (per-head Newton-Schulz for Q/K, config-gated by the headwise_muon spec flag, off by default), qk_named selector, shared momentum helper. Off-state routing bit-identical. - sinkhorn_update: Algorithm 1 as an optimizer for embedding/head matrices (Nesterov momentum, alternating row/col L2, near-zero row mask, sqrt(hidden) RMS conversion, gamma=0.18 lr correction, no weight decay, eps=1e-20 guards). - probe_speculative: draft/verify acceptance measurement over golden-v1 rows between the shipped layerdrop-int4 and small-2.1-int8 artifacts (early rows: acceptance 0.97-1.00, ~8.9 tokens/verify). - modal_teacher_sadeed: teacher-only SadeedDiac-25 predictions under the windowed protocol (used for the r6 per-domain slice, negative verdict recorded in rababa TODO.impl/02). - distill_specs: ara-diac-small-lite2 - the lite rung's untested variable, layer-drop init from the trained 2.1 student. --- scripts/probe_speculative.py | 231 ++++++++++++++++++++++++++++++++ src/gpu/distill_specs.yaml | 32 +++++ src/gpu/modal_distill.py | 11 ++ src/gpu/modal_teacher_sadeed.py | 89 ++++++++++++ src/gpu/muon.py | 89 +++++++++--- src/gpu/sinkhorn_update.py | 77 +++++++++++ tests/test_muon_headwise.py | 122 +++++++++++++++++ tests/test_sinkhorn_update.py | 95 +++++++++++++ 8 files changed, 729 insertions(+), 17 deletions(-) create mode 100644 scripts/probe_speculative.py create mode 100644 src/gpu/modal_teacher_sadeed.py create mode 100644 src/gpu/sinkhorn_update.py create mode 100644 tests/test_muon_headwise.py create mode 100644 tests/test_sinkhorn_update.py diff --git a/scripts/probe_speculative.py b/scripts/probe_speculative.py new file mode 100644 index 0000000..d0f717f --- /dev/null +++ b/scripts/probe_speculative.py @@ -0,0 +1,231 @@ +"""Speculative-decoding acceptance probe (TODO.qwen-next/10 §1). + +Measures whether ara-diac-layerdrop-1.0-int4 (drafter) can draft for +ara-diac-small-2.1-int8 (verifier) on the golden-v1 Arabic rows: block +acceptance rate, tokens per verifier pass, and exactness of the +speculative loop against verifier-only greedy (must be identical — +greedy verification is output-preserving by construction, so a +mismatch means a probe bug, not a model property). + +CPU-only, no training. Results land beside the log in ~/ml-logs. +""" + +from __future__ import annotations + +import json +import sys +import time +from pathlib import Path + +import numpy as np + +sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "runtime" / "src")) +from interscript_ml.model import Model # noqa: E402 +from interscript_ml.tokens import EOS_ID, PAD_ID, encode # noqa: E402 + +DRAFTER = "ara-diac-layerdrop-1.0-int4" +VERIFIER = "ara-diac-small-2.1-int8" +GOLDEN = Path.home() / "ml-logs" / "golden" / f"{DRAFTER}.jsonl" +OUT = Path.home() / "ml-logs" / "spec_probe" +K = 8 +# The plain-path exactness reference is O(T^2); cap it to short rows so +# the probe finishes. Longer rows still report acceptance statistics. +EXACT_LIMIT_BYTES = 600 + + +def plain_logits(model: Model, hidden, feed: list[int]) -> np.ndarray: + """Full-sequence run with zero pasts; works for both kv and plain + decoder graphs. Returns logits [L, V] (batch dropped).""" + outputs = model._decoder.run( + None, + { + "input_ids": np.array([feed], dtype=np.int64), + "encoder_hidden_states": hidden, + **model._pasts, + }, + ) + return outputs[0][0] + + +def plain_greedy(model: Model, hidden, max_len: int) -> list[int]: + """Full-sequence greedy over the same plain path the verifier + decisions use (zero pasts; works for kv and plain graphs). The + exactness theorem compares against THIS, not the KV path — int8 + near-ties can flip between execution paths.""" + feed = [PAD_ID] + out: list[int] = [] + for _ in range(max_len): + logits = plain_logits(model, hidden, feed) + token = int(np.argmax(logits[-1])) + if token == EOS_ID: + break + out.append(token) + feed.append(token) + return out + + +def draft(model: Model, hidden, seq: list[int], k: int, max_len: int) -> list[int]: + """Greedily draft k tokens conditioned on seq (the verifier- + authoritative prefix).""" + if len(seq) >= max_len: + return [] + feed = [PAD_ID] + seq + if model._kv_session: + pasts = dict(model._pasts) + draft_out: list[int] = [] + current = np.array([feed], dtype=np.int64) + while len(draft_out) < k and len(seq) + len(draft_out) < max_len: + outputs = model._decoder.run( + None, + {"input_ids": current, "encoder_hidden_states": hidden, **pasts}, + ) + results = dict(zip(model._output_names, outputs, strict=True)) + token = int(np.argmax(results["logits"][0, -1])) + if token == EOS_ID: + draft_out.append(token) + break + draft_out.append(token) + pasts = { + name: results[name.replace("past_", "present_", 1)] for name in pasts + } + current = np.array([[token]], dtype=np.int64) + return draft_out + # plain graph: stepwise full re-run + draft_out = [] + while len(draft_out) < k and len(seq) + len(draft_out) < max_len: + logits = plain_logits(model, hidden, feed + draft_out) + token = int(np.argmax(logits[-1])) + if token == EOS_ID: + draft_out.append(token) + break + draft_out.append(token) + return draft_out + + +def spec_decode(drafter: Model, verifier: Model, hidden_d, hidden_v, max_len: int): + seq: list[int] = [] + drafted = 0 + accepted = 0 # draft tokens the verifier kept (bonus excluded) + bonus = 0 + blocks = 0 + while len(seq) < max_len: + block = draft(drafter, hidden_d, seq, K, max_len) + if not block: + break + blocks += 1 + drafted += len(block) + feed = [PAD_ID] + seq + block + logits = plain_logits(verifier, hidden_v, feed) + n_acc = 0 + correction = None + for i, d in enumerate(block): + v = int(np.argmax(logits[len(seq) + i])) + if v == d: + n_acc += 1 + else: + correction = v + break + if correction is not None: + seq.extend(block[:n_acc]) + accepted += n_acc + if correction == EOS_ID: + return seq, dict( + drafted=drafted, accepted=accepted, bonus=bonus, blocks=blocks + ) + seq.append(correction) + continue + accepted += n_acc + seq.extend(block[:-1] if block[-1] == EOS_ID else block) + if block[-1] == EOS_ID: + return seq, dict( + drafted=drafted, accepted=accepted, bonus=bonus, blocks=blocks + ) + # seq already contains the block; the last logits position has + # seen all of it and predicts the bonus token + next_tok = int(np.argmax(logits[len(seq)])) + bonus += 1 + if next_tok == EOS_ID: + return seq, dict( + drafted=drafted, accepted=accepted, bonus=bonus, blocks=blocks + ) + seq.append(next_tok) + return seq, dict(drafted=drafted, accepted=accepted, bonus=bonus, blocks=blocks) + + +def main() -> None: + OUT.mkdir(parents=True, exist_ok=True) + drafter = Model.load(DRAFTER) + verifier = Model.load(VERIFIER) + rows = [ + json.loads(line) + for line in GOLDEN.read_text(encoding="utf-8").splitlines() + if line.strip() + ] + results = [] + for idx, row in enumerate(rows): + text = row["input"] + max_len = max(256, 4 * len(text.encode("utf-8"))) + ids_d = np.array([encode(text)], dtype=np.int64) + ids_v = np.array([encode(text)], dtype=np.int64) + hidden_d = drafter._encoder.run(None, {"input_ids": ids_d})[0] + hidden_v = verifier._encoder.run(None, {"input_ids": ids_v})[0] + t0 = time.time() + seq, stats = spec_decode(drafter, verifier, hidden_d, hidden_v, max_len) + dt = time.time() - t0 + # exactness vs the SAME execution path the verifier decisions + # came from: plain full-sequence greedy. (KV-greedy can differ + # from plain at int8 near-ties — that is the quantized-parity + # phenomenon, not a probe bug; reported separately.) + in_bytes = len(text.encode("utf-8")) + if in_bytes <= EXACT_LIMIT_BYTES: + ref_plain = plain_greedy(verifier, hidden_v, max_len) + exact: bool | None = seq == ref_plain + kv_match = verifier.generate(text, max_len=max_len) == ref_plain + else: + exact = None + kv_match = None + progress = len(seq) + results.append( + { + "row": idx, + "in_bytes": in_bytes, + "out_tokens": progress, + "exact": exact, + "kv_match": kv_match, + "acceptance": stats["accepted"] / max(stats["drafted"], 1), + "tokens_per_verify": progress / max(stats["blocks"], 1), + "blocks": stats["blocks"], + "bonus": stats["bonus"], + "sec": round(dt, 2), + } + ) + print( + f"row {idx:2d} in={results[-1]['in_bytes']:5d}B out={progress:4d} " + f"acc={results[-1]['acceptance']:.3f} tok/verify=" + f"{results[-1]['tokens_per_verify']:.2f} exact={exact} " + f"kv_match={kv_match} {dt:.1f}s", + flush=True, + ) + tot_out = sum(r["out_tokens"] for r in results) + tot_blocks = sum(r["blocks"] for r in results) + exact_rows = [r for r in results if r["exact"] is not None] + summary = { + "drafter": DRAFTER, + "verifier": VERIFIER, + "k": K, + "rows": len(results), + "exact_checked_rows": len(exact_rows), + "all_exact": all(r["exact"] for r in exact_rows), + "plain_kv_identical": all(r["kv_match"] for r in exact_rows), + "mean_acceptance": sum(r["acceptance"] for r in results) / len(results), + "tokens_per_verify_overall": tot_out / max(tot_blocks, 1), + "results": results, + } + (OUT / "results.json").write_text( + json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8" + ) + print(json.dumps({k: v for k, v in summary.items() if k != "results"}, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/src/gpu/distill_specs.yaml b/src/gpu/distill_specs.yaml index 0fa9d30..841fd55 100644 --- a/src/gpu/distill_specs.yaml +++ b/src/gpu/distill_specs.yaml @@ -320,6 +320,38 @@ ara-diac-small-layerdrop-6ep: mode: sequence optimizer: muon note: A2 rung; gate tracks G2a (4.57) within 0.1pp at half depth +ara-diac-small-lite2: + # TODO.impl/04: single-variable vs run-009 (5.78) - the INIT SOURCE. + # Byte-identical recipe except the layer-drop bridge copies kept + # layers from the TRAINED 2.1 student instead of generic byt5-small + # (DeepSeek's compress-from-trained-weights pattern). + teacher: rababa_arabic_byt5/run-007-news/best + teacher_volume: rababa + out_volume: rababa + student_init: /checkpoints/rababa_arabic_distill_small/run-007-r7-muon-6ep/best + layer_drop: 'true' + student_config: + d_model: 1472 + d_kv: 64 + d_ff: 3584 + num_heads: 6 + enc_layers: 6 + dec_layers: 4 + feed_forward_proj: gated-gelu + train: r5-units/domain.txt + train_extra: + - r5-units/replay.txt + unit_limits: + - 24000 + - 6000 + max_len: 1450 + label_beams: '1' + out: rababa_arabic_distill_small/run-012-lite2 + labels_file: teacher_labels_r7.jsonl + labels_complete: 'true' + mode: sequence + optimizer: muon + note: TODO.impl/04 init-source rung; gate vs run-009 (5.78) ara-diac-tiny-max: # THE gapless title test: the 30M class with the campaign's FULL # lever set — full corpus (24k+6k, not the 12k subset the collapse diff --git a/src/gpu/modal_distill.py b/src/gpu/modal_distill.py index 6024e8b..f583f61 100644 --- a/src/gpu/modal_distill.py +++ b/src/gpu/modal_distill.py @@ -1078,13 +1078,24 @@ def __getitem__(self, i): named += list(mtp_named(mtp_head)) muon_params, adamw_params = split_parameters(named) + headwise = [] + if spec.get("headwise_muon"): + from gpu.muon import qk_named + + headwise = [p for _, p in qk_named(named)] + headwise_ids = {id(p) for p in headwise} + muon_params = [p for p in muon_params if id(p) not in headwise_ids] optimizer = Muon( muon_params, lr=float(spec.get("muon_lr", 0.01)), momentum=0.95, weight_decay=0.01, ) + if headwise: + heads = int(spec.get("student_config", {}).get("num_heads", 6)) + optimizer.add_headwise_group(headwise, heads=heads) optimizer.add_adamw_group(adamw_params, lr=1e-4, weight_decay=0.0) print( f"[{spec_id}] muon: {len(muon_params)} matrix / " + f"{len(headwise)} headwise q/k / " f"{len(adamw_params)} embedding-like params", flush=True, ) diff --git a/src/gpu/modal_teacher_sadeed.py b/src/gpu/modal_teacher_sadeed.py new file mode 100644 index 0000000..f7a0a66 --- /dev/null +++ b/src/gpu/modal_teacher_sadeed.py @@ -0,0 +1,89 @@ +"""Modal probe: teacher-only SadeedDiac-25 predictions (TODO.qwen-next/10 §2). + +Runs the canonical teacher (r6 = run-006-morph) over the benchmark +under the published windowed protocol and writes per-paragraph +predictions to the checkpoints volume, so per-domain DER slicing can +be done offline (r7's slice comes from run-007's teacher column). + + modal run --detach src/gpu/modal_teacher_sadeed.py +""" + +from __future__ import annotations + +from pathlib import Path + +import modal + +REPO_ROOT = Path(__file__).resolve().parent.parent.parent + +IMAGE = ( + modal.Image.debian_slim(python_version="3.11") + .pip_install( + "torch>=2.4,<3", + "transformers==5.14.1", + "numpy>=1.26", + "pyarrow", + ) + .add_local_dir(str(REPO_ROOT), "/root/interscript-ml", copy=True) + .add_local_dir( + "/Users/mulgogi/src/interscript/rababa/data/sadeed-diac-25", + "/opt/rababa/data/sadeed-diac-25", + copy=True, + ) + .workdir("/root/interscript-ml") +) + +CHECKPOINTS = modal.Volume.from_name("rababa-checkpoints") + +TEACHER = "/checkpoints/rababa_arabic_byt5/run-006-morph/best" +OUT = "/checkpoints/probes/r6_sadeed_preds.jsonl" + +app = modal.App("interscript-ml-teacher-sadeed", image=IMAGE) + + +@app.function( + gpu="A10G", + cpu=8, + memory=32 * 1024, + timeout=2 * 3600, + volumes={"/checkpoints": CHECKPOINTS}, +) +def teacher_preds() -> dict: + import json + import sys + + import pyarrow.parquet as pq + import torch + from transformers import AutoModelForSeq2SeqLM, AutoTokenizer + + # Modal copies the entry file to /root/.py while the repo + # image sits at /root/interscript-ml — cover both layouts. + for cand in (Path.cwd() / "src", Path("/root/interscript-ml/src")): + if (cand / "harness" / "sadeed.py").exists() and str(cand) not in sys.path: + sys.path.insert(0, str(cand)) + from harness.sadeed import strip_diacritics, windowed_paragraphs + + tok = AutoTokenizer.from_pretrained("google/byt5-small") + teacher = AutoModelForSeq2SeqLM.from_pretrained(TEACHER).to("cuda").eval() + teacher.generation_config.max_length = 100_000 + + table = pq.read_table("/opt/rababa/data/sadeed-diac-25/train.parquet") + inputs = [strip_diacritics(t) for t in table.column("input").to_pylist()] + + preds = windowed_paragraphs(teacher, tok, inputs, window=1400) + + out = Path(OUT) + out.parent.mkdir(parents=True, exist_ok=True) + with open(out, "w", encoding="utf-8") as f: + for i, (src, pred) in enumerate(zip(inputs, preds, strict=True)): + f.write( + json.dumps({"idx": i, "src": src, "teacher": pred}, ensure_ascii=False) + + "\n" + ) + CHECKPOINTS.commit() + return {"rows": len(preds), "out": OUT} + + +@app.local_entrypoint() +def main() -> None: + print(teacher_preds.remote()) diff --git a/src/gpu/muon.py b/src/gpu/muon.py index b27d836..cb1ae1a 100644 --- a/src/gpu/muon.py +++ b/src/gpu/muon.py @@ -35,21 +35,18 @@ def zeropower_via_newtonschulz5(g: torch.Tensor, steps: int = 5) -> torch.Tensor return x.to(g.dtype) +_MUON_DEFAULTS = {"lr": 0.01, "momentum": 0.95, "nesterov": True, + "ns_steps": 5, "weight_decay": 0.0} + + class Muon(torch.optim.Optimizer): def __init__(self, params, lr: float = 0.01, momentum: float = 0.95, nesterov: bool = True, ns_steps: int = 5, weight_decay: float = 0.0) -> None: - super().__init__( - list(params), - { - "lr": lr, - "momentum": momentum, - "nesterov": nesterov, - "ns_steps": ns_steps, - "weight_decay": weight_decay, - "adamw": False, - }, - ) + settings = {**_MUON_DEFAULTS, "lr": lr, "momentum": momentum, + "nesterov": nesterov, "ns_steps": ns_steps, + "weight_decay": weight_decay, "adamw": False} + super().__init__(list(params), settings) def add_adamw_group(self, params, lr: float = 1e-4, betas=(0.9, 0.999), weight_decay: float = 0.0) -> None: @@ -63,29 +60,77 @@ def add_adamw_group(self, params, lr: float = 1e-4, betas=(0.9, 0.999), "adamw": True, }) + def add_headwise_group(self, params, heads: int) -> None: + """Per-head preconditioning (head-wise Muon): each attention + head's rows of a Q/K weight are orthogonalized as their own + matrix. Accepts params already inside a Muon group (their + group is converted in place) or all-new params (a new group is + added with the constructor's default settings).""" + params = list(params) + owned = {id(p) for group in self.param_groups for p in group["params"]} + hits = {id(p) for p in params if id(p) in owned} + if hits and len(hits) != len(params): + raise ValueError("headwise routing takes all-new or all-existing params") + if hits: + for group in self.param_groups: + members = {id(q) for q in group["params"]} + if hits & members: + if hits != members: + raise ValueError( + "headwise routing converts whole groups; " + "split the group first") + group["headwise"] = True + group["heads"] = heads + return + settings = {**_MUON_DEFAULTS, "adamw": False, + "params": params, "headwise": True, "heads": heads} + self.add_param_group(settings) + @torch.no_grad() def step(self, closure=None): # noqa: ARG002 for group in self.param_groups: if group.get("adamw"): self._adamw_step(group) + elif group.get("headwise"): + self._muon_step_headwise(group) else: self._muon_step(group) + def _momentum_direction(self, p, group): + state = self.state[p] + if "momentum_buffer" not in state: + state["momentum_buffer"] = torch.zeros_like(p.grad) + buf = state["momentum_buffer"] + buf.lerp_(p.grad, 1 - group["momentum"]) + return p.grad.lerp(buf, group["momentum"]) if group["nesterov"] else buf + def _muon_step(self, group) -> None: for p in group["params"]: if p.grad is None: continue - state = self.state[p] - if "momentum_buffer" not in state: - state["momentum_buffer"] = torch.zeros_like(p.grad) - buf = state["momentum_buffer"] - buf.lerp_(p.grad, 1 - group["momentum"]) - g = p.grad.lerp(buf, group["momentum"]) if group["nesterov"] else buf + g = self._momentum_direction(p, group) u = zeropower_via_newtonschulz5(g, steps=group["ns_steps"]) if group["weight_decay"]: p.mul_(1 - group["lr"] * group["weight_decay"]) p.add_(u.to(p.dtype), alpha=-group["lr"] * max(1, p.size(-2) / p.size(-1)) ** 0.5) + def _muon_step_headwise(self, group) -> None: + heads = group["heads"] + for p in group["params"]: + if p.grad is None: + continue + if p.size(0) % heads: + raise ValueError(f"param rows {p.size(0)} not divisible by {heads} heads") + g = self._momentum_direction(p, group) + if group["weight_decay"]: + p.mul_(1 - group["lr"] * group["weight_decay"]) + d = p.size(0) // heads + for h in range(heads): + rows = slice(h * d, (h + 1) * d) + u = zeropower_via_newtonschulz5(g[rows], steps=group["ns_steps"]) + scale = max(1, g[rows].size(-2) / g[rows].size(-1)) ** 0.5 + p[rows].add_(u.to(p.dtype), alpha=-group["lr"] * scale) + def _adamw_step(self, group) -> None: beta1, beta2 = group["betas"] for p in group["params"]: @@ -128,3 +173,13 @@ def split_parameters(named_params): ) (adamw if embedding_like else muon).append(p) return muon, adamw + + +def qk_named(named_params): + """The Q/K projection weights head-wise Muon applies to (DeepSeek + V4.1 Flash sec 2.5; GLM-5 and Kimi-K3 validate the same split). + T5 names both self- and cross-attention projections.""" + import re + + pattern = re.compile(r"(SelfAttention|EncDecAttention)\.(q|k)\.weight$") + return [(name, p) for name, p in named_params if pattern.search(name)] diff --git a/src/gpu/sinkhorn_update.py b/src/gpu/sinkhorn_update.py new file mode 100644 index 0000000..f7a37b4 --- /dev/null +++ b/src/gpu/sinkhorn_update.py @@ -0,0 +1,77 @@ +"""Sinkhorn-balanced momentum update for embedding tables and +prediction heads (DeepSeek-V4.1-Flash Algorithm 1, TODO.impl/06). + +Replaces Adam for large row-structured matrices: Nesterov momentum, +then alternating row/column L2 normalization (odd number of steps, +ending on rows), near-zero row masking, sqrt(n) to convert unit row +L2 into unit row RMS, and a gamma-corrected learning rate to match +Adam's update magnitude. Momentum-only state; no weight decay. + +Numerical note (the mHC lesson): alternating direct division produced +NaNs for ~10% of inits when magnitudes shrink; the eps here is 1e-20 +and every division is guarded. +""" + +from __future__ import annotations + +import math + +import torch + + +def sinkhorn_balance( + update: torch.Tensor, + k: int = 11, + tau: float = 1e-3, + eps: float = 1e-20, +) -> torch.Tensor: + """Alternating row/column L2 normalization ending on rows (k odd), + with near-zero rows masked. Returns the balanced update; the + caller applies the sqrt(n) RMS conversion.""" + if k % 2 == 0: + k += 1 # the algorithm requires an odd count (ends row-wise) + rows = update.norm(dim=1) + mean_row = rows.mean() + work = update.clone() + work[(rows <= tau * mean_row)] = 0.0 + for step in range(1, k + 1): + if step % 2 == 1: # odd: rows + norms = work.norm(dim=1, keepdim=True).add_(eps) + work = work / norms + else: # even: columns + norms = work.norm(dim=0, keepdim=True).add_(eps) + work = work / norms + work[(rows <= tau * mean_row)] = 0.0 + return work + + +class SinkhornUpdate(torch.optim.Optimizer): + """Momentum + Sinkhorn balancing in place of Adam's second moment, + for embedding tables and prediction heads (row = token/n-gram, + column = hidden feature).""" + + def __init__(self, params, lr: float = 2.6e-4, beta: float = 0.95, + gamma: float = 0.18, k: int = 11, tau: float = 1e-3, + weight_decay: float = 0.0) -> None: # noqa: ARG002 (ignored by contract) + # rows of the update matrix carry the row structure (token / + # n-gram identity); n = hidden feature count (columns), whose + # sqrt converts unit row L2 norm into unit row RMS + settings = {"lr": lr, "momentum": beta, "gamma": gamma, + "k": k, "tau": tau, "weight_decay": 0.0} + super().__init__(list(params), settings) + + @torch.no_grad() + def step(self, closure=None) -> None: # noqa: ARG002 + for group in self.param_groups: + beta, lr, gamma = group["momentum"], group["lr"], group["gamma"] + for p in group["params"]: + if p.grad is None: + continue + state = self.state[p] + if "momentum_buffer" not in state: + state["momentum_buffer"] = torch.zeros_like(p.grad) + buf = state["momentum_buffer"] + buf.lerp_(p.grad, 1 - beta) + g = p.grad.lerp(buf, beta) # Nesterov lookahead + u = sinkhorn_balance(g, k=group["k"], tau=group["tau"]) + p.add_(u.to(p.dtype), alpha=-lr * gamma * math.sqrt(p.size(1))) diff --git a/tests/test_muon_headwise.py b/tests/test_muon_headwise.py new file mode 100644 index 0000000..a60e040 --- /dev/null +++ b/tests/test_muon_headwise.py @@ -0,0 +1,122 @@ +"""Head-wise Muon specs (TODO.impl/03): per-head preconditioning for +Q/K weights — the update each attention head would get from vanilla +Muon applied to its slice alone, reassembled into the full matrix. +Validated externally by DeepSeek-V4.1-Flash, GLM-5, Kimi-K3.""" + +from __future__ import annotations + +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src")) + +import pytest + +torch = pytest.importorskip("torch") + +from gpu.muon import Muon, qk_named, split_parameters # noqa: E402 + + +def _grad_like(shape, seed): + g = torch.Generator().manual_seed(seed) + return torch.randn(*shape, generator=g) + + +def test_headwise_equals_per_slice_vanilla() -> None: + """One step of head-wise Muon on [H*d, D] must equal vanilla Muon + applied to each [d, D] slice with the same grad — that is the + entire semantic content of 'different preconditioners per head'.""" + heads, d, dim = 3, 8, 16 + weight = torch.randn(heads * d, dim, generator=torch.Generator().manual_seed(7)) + grad = _grad_like((heads * d, dim), seed=11) + + slices = list(weight.chunk(heads)) + slice_grads = list(grad.chunk(heads)) + # clone BEFORE vanilla mutates `weight` through its chunk views + hw_weight = weight.clone() + hw = Muon([hw_weight], lr=0.01, momentum=0.95) + hw.add_headwise_group(list(hw.param_groups[0]["params"]), heads=heads) + # the base group must be replaced, not duplicated: the param moves + # into the headwise group + assert hw.param_groups[0].get("headwise") is True + hw_weight.grad = grad.clone() + hw.step() + + vanilla = Muon(slices, lr=0.01, momentum=0.95) + for p, g in zip(slices, slice_grads): + p.grad = g.clone() + vanilla.step() + + for h in range(heads): + torch.testing.assert_close(hw_weight.chunk(heads)[h], weight.chunk(heads)[h]) + + +def test_heads_one_matches_vanilla_whole() -> None: + weight = torch.randn(12, 16, generator=torch.Generator().manual_seed(3)) + grad = _grad_like((12, 16), seed=5) + + v = Muon([weight.clone()], lr=0.01, momentum=0.95) + v.param_groups[0]["params"][0].grad = grad.clone() + v.step() + + hw_w = weight.clone() + hw = Muon([hw_w], lr=0.01, momentum=0.95) + hw.add_headwise_group(list(hw.param_groups[0]["params"]), heads=1) + hw_w.grad = grad.clone() + hw.step() + + torch.testing.assert_close(hw_w, v.param_groups[0]["params"][0]) + + +def test_headwise_differs_from_vanilla_for_heterogeneous_heads() -> None: + """If every head were preconditioned identically the split would be + a no-op; give the slices different singular structure and the two + updates must diverge.""" + heads, d, dim = 2, 8, 16 + base = _grad_like((heads * d, dim), seed=13) + base[:d] *= 0.01 # near-isotropic head vs anisotropic head + weight = torch.zeros(heads * d, dim) + + v = Muon([weight.clone()], lr=0.01, momentum=0.95) + v.param_groups[0]["params"][0].grad = base.clone() + v.step() + + hw_w = weight.clone() + hw = Muon([hw_w], lr=0.01, momentum=0.95) + hw.add_headwise_group(list(hw.param_groups[0]["params"]), heads=heads) + hw_w.grad = base.clone() + hw.step() + + assert not torch.allclose(hw_w, v.param_groups[0]["params"][0]) + + +def test_qk_named_selects_only_qk_projections() -> None: + named = [ + ("encoder.block.0.layer.0.SelfAttention.q.weight", torch.zeros(2)), + ("encoder.block.0.layer.0.SelfAttention.k.weight", torch.zeros(2)), + ("encoder.block.0.layer.0.SelfAttention.v.weight", torch.zeros(2)), + ("encoder.block.0.layer.0.SelfAttention.o.weight", torch.zeros(2)), + ("decoder.block.0.layer.0.EncDecAttention.q.weight", torch.zeros(2)), + ("decoder.block.0.layer.0.EncDecAttention.k.weight", torch.zeros(2)), + ("encoder.block.0.layer.1.DenseReluDense.wi_0.weight", torch.zeros(2)), + ] + selected = {n for n, _ in qk_named(named)} + assert selected == { + "encoder.block.0.layer.0.SelfAttention.q.weight", + "encoder.block.0.layer.0.SelfAttention.k.weight", + "decoder.block.0.layer.0.EncDecAttention.q.weight", + "decoder.block.0.layer.0.EncDecAttention.k.weight", + } + + +def test_split_parameters_unchanged_when_headwise_unused() -> None: + """Off-state must be bit-identical routing: the feature adds a + group kind, it does not touch the default split.""" + named = [ + ("shared.weight", torch.zeros(3, 4, requires_grad=True)), + ("encoder.block.0.layer.1.DenseReluDense.wi_0.weight", torch.zeros(4, 3, requires_grad=True)), + ("encoder.block.0.layer.0.layer_norm.weight", torch.zeros(3, requires_grad=True)), + ] + muon, adamw = split_parameters(named) + assert [p.shape for p in muon] == [torch.Size([4, 3])] + assert [p.shape for p in adamw] == [torch.Size([3, 4]), torch.Size([3])] diff --git a/tests/test_sinkhorn_update.py b/tests/test_sinkhorn_update.py new file mode 100644 index 0000000..c78e7be --- /dev/null +++ b/tests/test_sinkhorn_update.py @@ -0,0 +1,95 @@ +"""Sinkhorn-balanced momentum update specs (TODO.impl/06, DeepSeek +V4.1-Flash Algorithm 1): Nesterov momentum, alternating row/column L2 +normalization over K steps, near-zero row masking, sqrt(n) RMS +conversion, gamma-corrected learning rate, no weight decay. Replaces +Adam for embedding tables and prediction heads.""" + +from __future__ import annotations + +import math +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src")) + +import pytest + +torch = pytest.importorskip("torch") + +from gpu.sinkhorn_update import SinkhornUpdate, sinkhorn_balance # noqa: E402 + + +def _randn(rows, cols, seed): + return torch.randn(rows, cols, generator=torch.Generator().manual_seed(seed)) + + +def test_balance_yields_unit_row_norms() -> None: + w = _randn(6, 4, seed=1) + u = sinkhorn_balance(w, k=11) + row_norms = u.norm(dim=1) + assert torch.allclose(row_norms, torch.ones_like(row_norms), atol=1e-4) + + +def test_balance_masks_near_zero_rows() -> None: + """Rows at or below tau * mean-row-norm contribute nothing — their + update stays zero after balancing.""" + w = _randn(6, 4, seed=2) + w[0] = 1e-9 # dead row + u = sinkhorn_balance(w, k=11, tau=1e-3) + assert u[0].abs().max().item() == 0.0 + live = u[1:] + assert torch.allclose(live.norm(dim=1), torch.ones(live.size(0)), atol=1e-4) + + +def test_balance_survives_small_magnitudes() -> None: + """The mHC log-domain lesson: direct division NaNs when magnitudes + shrink across alternating normalizations; eps guards must hold at + 1e-20-class scales.""" + w = _randn(8, 5, seed=3) * 1e-8 + u = sinkhorn_balance(w, k=11, eps=1e-20) + assert torch.isfinite(u).all() + + +def test_balance_matches_hand_computed_case() -> None: + """k=1 is a single ROW step: each row is L2-normalized; the sqrt(n) + RMS conversion is the optimizer's job, not the balance's.""" + w = torch.tensor([[3.0, 4.0], [6.0, 8.0]]) + u = sinkhorn_balance(w, k=1) + expect = torch.tensor([[0.6, 0.8], [0.6, 0.8]]) + torch.testing.assert_close(u, expect) + + +def test_optimizer_step_matches_manual_algorithm() -> None: + """One SinkhornUpdate step on an embedding-shaped weight equals the + algorithm written out by hand: Nesterov momentum, balance, sqrt(n), + gamma-corrected lr, no weight decay.""" + rows, cols = 5, 3 + weight = _randn(rows, cols, seed=4) + grad = _randn(rows, cols, seed=5) + beta, lr, gamma, n = 0.9, 0.02, 0.18, rows + + before = weight.clone() + opt = SinkhornUpdate([weight], lr=lr, beta=beta, gamma=gamma) + weight.grad = grad.clone() + opt.step() + + momentum = grad.clone() * (1 - beta) # first step: buf = (1-beta)*g + lookahead = beta * momentum + (1 - beta) * grad + u = sinkhorn_balance(lookahead, k=11) + n = before.size(1) # hidden feature count + delta = math.sqrt(n) * u + expect = before - (gamma * lr) * delta + torch.testing.assert_close(weight, expect) + + +def test_no_weight_decay_applied() -> None: + weight = _randn(4, 3, seed=6) + before = weight.clone() + opt = SinkhornUpdate([weight], weight_decay=0.1) # ignored by contract + weight.grad = _randn(4, 3, seed=7) + opt.step() + # balanced rows have unit L2 norm; movement is exactly + # lr * gamma * sqrt(n_hidden), untouched by weight decay + movement = (weight - before).norm(dim=1) + expect = opt.defaults["lr"] * opt.defaults["gamma"] * math.sqrt(3) + assert torch.allclose(movement, torch.full_like(movement, expect), atol=1e-5) From 76bce0c36e6d67d3adbe2f970592f8a54d2c7620 Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Fri, 11 Sep 2026 21:34:16 +0800 Subject: [PATCH 2/2] style: ruff clean (unused import, line length, zip strict) --- src/gpu/modal_golden_matrix.py | 24 ++++++++++++++++++------ src/gpu/modal_teacher_sadeed.py | 1 - tests/test_muon_headwise.py | 5 +++-- 3 files changed, 21 insertions(+), 9 deletions(-) diff --git a/src/gpu/modal_golden_matrix.py b/src/gpu/modal_golden_matrix.py index 765792c..69353ae 100644 --- a/src/gpu/modal_golden_matrix.py +++ b/src/gpu/modal_golden_matrix.py @@ -39,19 +39,31 @@ # index id -> volume zip path + test pairs (volume-relative) SOURCES = { - "khm-latn-1.0": ("imf/khm-latn/khm-latn-1.0-fp32.zip", "/secryst-datasets/khmer-translit/test.jsonl"), + "khm-latn-1.0": + ("imf/khm-latn/khm-latn-1.0-fp32.zip", + "/secryst-datasets/khmer-translit/test.jsonl"), "urd-g2p-1.0": ("imf/urd-g2p/urd-g2p-1.0-fp32.zip", "/ud-g2p/urdu-g2p/test.jsonl"), "urd-diac-1.0": ("imf/urd-diac/urd-diac-1.0-fp32.zip", "/ud-diacrit/urdu-diacrit/test.jsonl"), - "tha-g2p-base-1.0": ("imf/tha-g2p-base/tha-g2p-base-1.0-fp32.zip", "/secryst-datasets/thai-ipa/test.jsonl"), - "tha-g2p-small-1.0": ("imf/tha-g2p-small/tha-g2p-small-1.0-int8.zip", "/secryst-datasets/thai-ipa/test.jsonl"), + "tha-g2p-base-1.0": + ("imf/tha-g2p-base/tha-g2p-base-1.0-fp32.zip", + "/secryst-datasets/thai-ipa/test.jsonl"), + "tha-g2p-small-1.0": + ("imf/tha-g2p-small/tha-g2p-small-1.0-int8.zip", + "/secryst-datasets/thai-ipa/test.jsonl"), "fas-g2p-1.0": ("imf/fas-g2p/fas-g2p-1.0-fp32.zip", "/persian/persian-g2p/test.jsonl"), "heb-diac-1.0": ("imf/heb-diac/heb-diac-1.0-fp32.zip", "nakdimon/test-imf.jsonl"), "heb-diac-1.1": ("imf/heb-diac/heb.zip", "nakdimon/test-imf.jsonl"), - "heb-diac-small-1.0": ("imf/heb-diac-small/heb-diac-small-1.0-fp32.zip", "nakdimon/test-imf.jsonl"), + "heb-diac-small-1.0": + ("imf/heb-diac-small/heb-diac-small-1.0-fp32.zip", + "nakdimon/test-imf.jsonl"), "ara-diac-1.0": ("imf/ara-diac/ara-diac-1.0-fp32.zip", "arabic-sadeed-imf/test.jsonl"), "ara-diac-2.0-int8": ("imf/ara-diac2/ara-diac-2.0-int8.zip", "arabic-sadeed-imf/test.jsonl"), - "ara-diac-small-2.1-int8": ("imf/ara-diac-small-21/ara-diac-small-2.1-int8.zip", "arabic-sadeed-imf/test.jsonl"), - "ara-diac-layerdrop-1.0-int4": ("imf/ara-diac-layerdrop/ara-diac-layerdrop-1.0-int4.zip", "arabic-sadeed-imf/test.jsonl"), + "ara-diac-small-2.1-int8": + ("imf/ara-diac-small-21/ara-diac-small-2.1-int8.zip", + "arabic-sadeed-imf/test.jsonl"), + "ara-diac-layerdrop-1.0-int4": + ("imf/ara-diac-layerdrop/ara-diac-layerdrop-1.0-int4.zip", + "arabic-sadeed-imf/test.jsonl"), } diff --git a/src/gpu/modal_teacher_sadeed.py b/src/gpu/modal_teacher_sadeed.py index f7a0a66..32377aa 100644 --- a/src/gpu/modal_teacher_sadeed.py +++ b/src/gpu/modal_teacher_sadeed.py @@ -53,7 +53,6 @@ def teacher_preds() -> dict: import sys import pyarrow.parquet as pq - import torch from transformers import AutoModelForSeq2SeqLM, AutoTokenizer # Modal copies the entry file to /root/.py while the repo diff --git a/tests/test_muon_headwise.py b/tests/test_muon_headwise.py index a60e040..a00193c 100644 --- a/tests/test_muon_headwise.py +++ b/tests/test_muon_headwise.py @@ -43,7 +43,7 @@ def test_headwise_equals_per_slice_vanilla() -> None: hw.step() vanilla = Muon(slices, lr=0.01, momentum=0.95) - for p, g in zip(slices, slice_grads): + for p, g in zip(slices, slice_grads, strict=True): p.grad = g.clone() vanilla.step() @@ -114,7 +114,8 @@ def test_split_parameters_unchanged_when_headwise_unused() -> None: group kind, it does not touch the default split.""" named = [ ("shared.weight", torch.zeros(3, 4, requires_grad=True)), - ("encoder.block.0.layer.1.DenseReluDense.wi_0.weight", torch.zeros(4, 3, requires_grad=True)), + ("encoder.block.0.layer.1.DenseReluDense.wi_0.weight", + torch.zeros(4, 3, requires_grad=True)), ("encoder.block.0.layer.0.layer_norm.weight", torch.zeros(3, requires_grad=True)), ] muon, adamw = split_parameters(named)