Skip to content

Commit 6473dad

Browse files
author
Ronald Tse
committed
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.
1 parent ee65ec8 commit 6473dad

8 files changed

Lines changed: 729 additions & 17 deletions

File tree

‎scripts/probe_speculative.py‎

Lines changed: 231 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,231 @@
1+
"""Speculative-decoding acceptance probe (TODO.qwen-next/10 §1).
2+
3+
Measures whether ara-diac-layerdrop-1.0-int4 (drafter) can draft for
4+
ara-diac-small-2.1-int8 (verifier) on the golden-v1 Arabic rows: block
5+
acceptance rate, tokens per verifier pass, and exactness of the
6+
speculative loop against verifier-only greedy (must be identical —
7+
greedy verification is output-preserving by construction, so a
8+
mismatch means a probe bug, not a model property).
9+
10+
CPU-only, no training. Results land beside the log in ~/ml-logs.
11+
"""
12+
13+
from __future__ import annotations
14+
15+
import json
16+
import sys
17+
import time
18+
from pathlib import Path
19+
20+
import numpy as np
21+
22+
sys.path.insert(0, str(Path(__file__).resolve().parents[1] / "runtime" / "src"))
23+
from interscript_ml.model import Model # noqa: E402
24+
from interscript_ml.tokens import EOS_ID, PAD_ID, encode # noqa: E402
25+
26+
DRAFTER = "ara-diac-layerdrop-1.0-int4"
27+
VERIFIER = "ara-diac-small-2.1-int8"
28+
GOLDEN = Path.home() / "ml-logs" / "golden" / f"{DRAFTER}.jsonl"
29+
OUT = Path.home() / "ml-logs" / "spec_probe"
30+
K = 8
31+
# The plain-path exactness reference is O(T^2); cap it to short rows so
32+
# the probe finishes. Longer rows still report acceptance statistics.
33+
EXACT_LIMIT_BYTES = 600
34+
35+
36+
def plain_logits(model: Model, hidden, feed: list[int]) -> np.ndarray:
37+
"""Full-sequence run with zero pasts; works for both kv and plain
38+
decoder graphs. Returns logits [L, V] (batch dropped)."""
39+
outputs = model._decoder.run(
40+
None,
41+
{
42+
"input_ids": np.array([feed], dtype=np.int64),
43+
"encoder_hidden_states": hidden,
44+
**model._pasts,
45+
},
46+
)
47+
return outputs[0][0]
48+
49+
50+
def plain_greedy(model: Model, hidden, max_len: int) -> list[int]:
51+
"""Full-sequence greedy over the same plain path the verifier
52+
decisions use (zero pasts; works for kv and plain graphs). The
53+
exactness theorem compares against THIS, not the KV path — int8
54+
near-ties can flip between execution paths."""
55+
feed = [PAD_ID]
56+
out: list[int] = []
57+
for _ in range(max_len):
58+
logits = plain_logits(model, hidden, feed)
59+
token = int(np.argmax(logits[-1]))
60+
if token == EOS_ID:
61+
break
62+
out.append(token)
63+
feed.append(token)
64+
return out
65+
66+
67+
def draft(model: Model, hidden, seq: list[int], k: int, max_len: int) -> list[int]:
68+
"""Greedily draft k tokens conditioned on seq (the verifier-
69+
authoritative prefix)."""
70+
if len(seq) >= max_len:
71+
return []
72+
feed = [PAD_ID] + seq
73+
if model._kv_session:
74+
pasts = dict(model._pasts)
75+
draft_out: list[int] = []
76+
current = np.array([feed], dtype=np.int64)
77+
while len(draft_out) < k and len(seq) + len(draft_out) < max_len:
78+
outputs = model._decoder.run(
79+
None,
80+
{"input_ids": current, "encoder_hidden_states": hidden, **pasts},
81+
)
82+
results = dict(zip(model._output_names, outputs, strict=True))
83+
token = int(np.argmax(results["logits"][0, -1]))
84+
if token == EOS_ID:
85+
draft_out.append(token)
86+
break
87+
draft_out.append(token)
88+
pasts = {
89+
name: results[name.replace("past_", "present_", 1)] for name in pasts
90+
}
91+
current = np.array([[token]], dtype=np.int64)
92+
return draft_out
93+
# plain graph: stepwise full re-run
94+
draft_out = []
95+
while len(draft_out) < k and len(seq) + len(draft_out) < max_len:
96+
logits = plain_logits(model, hidden, feed + draft_out)
97+
token = int(np.argmax(logits[-1]))
98+
if token == EOS_ID:
99+
draft_out.append(token)
100+
break
101+
draft_out.append(token)
102+
return draft_out
103+
104+
105+
def spec_decode(drafter: Model, verifier: Model, hidden_d, hidden_v, max_len: int):
106+
seq: list[int] = []
107+
drafted = 0
108+
accepted = 0 # draft tokens the verifier kept (bonus excluded)
109+
bonus = 0
110+
blocks = 0
111+
while len(seq) < max_len:
112+
block = draft(drafter, hidden_d, seq, K, max_len)
113+
if not block:
114+
break
115+
blocks += 1
116+
drafted += len(block)
117+
feed = [PAD_ID] + seq + block
118+
logits = plain_logits(verifier, hidden_v, feed)
119+
n_acc = 0
120+
correction = None
121+
for i, d in enumerate(block):
122+
v = int(np.argmax(logits[len(seq) + i]))
123+
if v == d:
124+
n_acc += 1
125+
else:
126+
correction = v
127+
break
128+
if correction is not None:
129+
seq.extend(block[:n_acc])
130+
accepted += n_acc
131+
if correction == EOS_ID:
132+
return seq, dict(
133+
drafted=drafted, accepted=accepted, bonus=bonus, blocks=blocks
134+
)
135+
seq.append(correction)
136+
continue
137+
accepted += n_acc
138+
seq.extend(block[:-1] if block[-1] == EOS_ID else block)
139+
if block[-1] == EOS_ID:
140+
return seq, dict(
141+
drafted=drafted, accepted=accepted, bonus=bonus, blocks=blocks
142+
)
143+
# seq already contains the block; the last logits position has
144+
# seen all of it and predicts the bonus token
145+
next_tok = int(np.argmax(logits[len(seq)]))
146+
bonus += 1
147+
if next_tok == EOS_ID:
148+
return seq, dict(
149+
drafted=drafted, accepted=accepted, bonus=bonus, blocks=blocks
150+
)
151+
seq.append(next_tok)
152+
return seq, dict(drafted=drafted, accepted=accepted, bonus=bonus, blocks=blocks)
153+
154+
155+
def main() -> None:
156+
OUT.mkdir(parents=True, exist_ok=True)
157+
drafter = Model.load(DRAFTER)
158+
verifier = Model.load(VERIFIER)
159+
rows = [
160+
json.loads(line)
161+
for line in GOLDEN.read_text(encoding="utf-8").splitlines()
162+
if line.strip()
163+
]
164+
results = []
165+
for idx, row in enumerate(rows):
166+
text = row["input"]
167+
max_len = max(256, 4 * len(text.encode("utf-8")))
168+
ids_d = np.array([encode(text)], dtype=np.int64)
169+
ids_v = np.array([encode(text)], dtype=np.int64)
170+
hidden_d = drafter._encoder.run(None, {"input_ids": ids_d})[0]
171+
hidden_v = verifier._encoder.run(None, {"input_ids": ids_v})[0]
172+
t0 = time.time()
173+
seq, stats = spec_decode(drafter, verifier, hidden_d, hidden_v, max_len)
174+
dt = time.time() - t0
175+
# exactness vs the SAME execution path the verifier decisions
176+
# came from: plain full-sequence greedy. (KV-greedy can differ
177+
# from plain at int8 near-ties — that is the quantized-parity
178+
# phenomenon, not a probe bug; reported separately.)
179+
in_bytes = len(text.encode("utf-8"))
180+
if in_bytes <= EXACT_LIMIT_BYTES:
181+
ref_plain = plain_greedy(verifier, hidden_v, max_len)
182+
exact: bool | None = seq == ref_plain
183+
kv_match = verifier.generate(text, max_len=max_len) == ref_plain
184+
else:
185+
exact = None
186+
kv_match = None
187+
progress = len(seq)
188+
results.append(
189+
{
190+
"row": idx,
191+
"in_bytes": in_bytes,
192+
"out_tokens": progress,
193+
"exact": exact,
194+
"kv_match": kv_match,
195+
"acceptance": stats["accepted"] / max(stats["drafted"], 1),
196+
"tokens_per_verify": progress / max(stats["blocks"], 1),
197+
"blocks": stats["blocks"],
198+
"bonus": stats["bonus"],
199+
"sec": round(dt, 2),
200+
}
201+
)
202+
print(
203+
f"row {idx:2d} in={results[-1]['in_bytes']:5d}B out={progress:4d} "
204+
f"acc={results[-1]['acceptance']:.3f} tok/verify="
205+
f"{results[-1]['tokens_per_verify']:.2f} exact={exact} "
206+
f"kv_match={kv_match} {dt:.1f}s",
207+
flush=True,
208+
)
209+
tot_out = sum(r["out_tokens"] for r in results)
210+
tot_blocks = sum(r["blocks"] for r in results)
211+
exact_rows = [r for r in results if r["exact"] is not None]
212+
summary = {
213+
"drafter": DRAFTER,
214+
"verifier": VERIFIER,
215+
"k": K,
216+
"rows": len(results),
217+
"exact_checked_rows": len(exact_rows),
218+
"all_exact": all(r["exact"] for r in exact_rows),
219+
"plain_kv_identical": all(r["kv_match"] for r in exact_rows),
220+
"mean_acceptance": sum(r["acceptance"] for r in results) / len(results),
221+
"tokens_per_verify_overall": tot_out / max(tot_blocks, 1),
222+
"results": results,
223+
}
224+
(OUT / "results.json").write_text(
225+
json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8"
226+
)
227+
print(json.dumps({k: v for k, v in summary.items() if k != "results"}, indent=2))
228+
229+
230+
if __name__ == "__main__":
231+
main()

‎src/gpu/distill_specs.yaml‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -320,6 +320,38 @@ ara-diac-small-layerdrop-6ep:
320320
mode: sequence
321321
optimizer: muon
322322
note: A2 rung; gate tracks G2a (4.57) within 0.1pp at half depth
323+
ara-diac-small-lite2:
324+
# TODO.impl/04: single-variable vs run-009 (5.78) - the INIT SOURCE.
325+
# Byte-identical recipe except the layer-drop bridge copies kept
326+
# layers from the TRAINED 2.1 student instead of generic byt5-small
327+
# (DeepSeek's compress-from-trained-weights pattern).
328+
teacher: rababa_arabic_byt5/run-007-news/best
329+
teacher_volume: rababa
330+
out_volume: rababa
331+
student_init: /checkpoints/rababa_arabic_distill_small/run-007-r7-muon-6ep/best
332+
layer_drop: 'true'
333+
student_config:
334+
d_model: 1472
335+
d_kv: 64
336+
d_ff: 3584
337+
num_heads: 6
338+
enc_layers: 6
339+
dec_layers: 4
340+
feed_forward_proj: gated-gelu
341+
train: r5-units/domain.txt
342+
train_extra:
343+
- r5-units/replay.txt
344+
unit_limits:
345+
- 24000
346+
- 6000
347+
max_len: 1450
348+
label_beams: '1'
349+
out: rababa_arabic_distill_small/run-012-lite2
350+
labels_file: teacher_labels_r7.jsonl
351+
labels_complete: 'true'
352+
mode: sequence
353+
optimizer: muon
354+
note: TODO.impl/04 init-source rung; gate vs run-009 (5.78)
323355
ara-diac-tiny-max:
324356
# THE gapless title test: the 30M class with the campaign's FULL
325357
# lever set — full corpus (24k+6k, not the 12k subset the collapse

‎src/gpu/modal_distill.py‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1078,13 +1078,24 @@ def __getitem__(self, i):
10781078

10791079
named += list(mtp_named(mtp_head))
10801080
muon_params, adamw_params = split_parameters(named)
1081+
headwise = []
1082+
if spec.get("headwise_muon"):
1083+
from gpu.muon import qk_named
1084+
1085+
headwise = [p for _, p in qk_named(named)]
1086+
headwise_ids = {id(p) for p in headwise}
1087+
muon_params = [p for p in muon_params if id(p) not in headwise_ids]
10811088
optimizer = Muon(
10821089
muon_params, lr=float(spec.get("muon_lr", 0.01)),
10831090
momentum=0.95, weight_decay=0.01,
10841091
)
1092+
if headwise:
1093+
heads = int(spec.get("student_config", {}).get("num_heads", 6))
1094+
optimizer.add_headwise_group(headwise, heads=heads)
10851095
optimizer.add_adamw_group(adamw_params, lr=1e-4, weight_decay=0.0)
10861096
print(
10871097
f"[{spec_id}] muon: {len(muon_params)} matrix / "
1098+
f"{len(headwise)} headwise q/k / "
10881099
f"{len(adamw_params)} embedding-like params",
10891100
flush=True,
10901101
)

‎src/gpu/modal_teacher_sadeed.py‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
"""Modal probe: teacher-only SadeedDiac-25 predictions (TODO.qwen-next/10 §2).
2+
3+
Runs the canonical teacher (r6 = run-006-morph) over the benchmark
4+
under the published windowed protocol and writes per-paragraph
5+
predictions to the checkpoints volume, so per-domain DER slicing can
6+
be done offline (r7's slice comes from run-007's teacher column).
7+
8+
modal run --detach src/gpu/modal_teacher_sadeed.py
9+
"""
10+
11+
from __future__ import annotations
12+
13+
from pathlib import Path
14+
15+
import modal
16+
17+
REPO_ROOT = Path(__file__).resolve().parent.parent.parent
18+
19+
IMAGE = (
20+
modal.Image.debian_slim(python_version="3.11")
21+
.pip_install(
22+
"torch>=2.4,<3",
23+
"transformers==5.14.1",
24+
"numpy>=1.26",
25+
"pyarrow",
26+
)
27+
.add_local_dir(str(REPO_ROOT), "/root/interscript-ml", copy=True)
28+
.add_local_dir(
29+
"/Users/mulgogi/src/interscript/rababa/data/sadeed-diac-25",
30+
"/opt/rababa/data/sadeed-diac-25",
31+
copy=True,
32+
)
33+
.workdir("/root/interscript-ml")
34+
)
35+
36+
CHECKPOINTS = modal.Volume.from_name("rababa-checkpoints")
37+
38+
TEACHER = "/checkpoints/rababa_arabic_byt5/run-006-morph/best"
39+
OUT = "/checkpoints/probes/r6_sadeed_preds.jsonl"
40+
41+
app = modal.App("interscript-ml-teacher-sadeed", image=IMAGE)
42+
43+
44+
@app.function(
45+
gpu="A10G",
46+
cpu=8,
47+
memory=32 * 1024,
48+
timeout=2 * 3600,
49+
volumes={"/checkpoints": CHECKPOINTS},
50+
)
51+
def teacher_preds() -> dict:
52+
import json
53+
import sys
54+
55+
import pyarrow.parquet as pq
56+
import torch
57+
from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
58+
59+
# Modal copies the entry file to /root/<name>.py while the repo
60+
# image sits at /root/interscript-ml — cover both layouts.
61+
for cand in (Path.cwd() / "src", Path("/root/interscript-ml/src")):
62+
if (cand / "harness" / "sadeed.py").exists() and str(cand) not in sys.path:
63+
sys.path.insert(0, str(cand))
64+
from harness.sadeed import strip_diacritics, windowed_paragraphs
65+
66+
tok = AutoTokenizer.from_pretrained("google/byt5-small")
67+
teacher = AutoModelForSeq2SeqLM.from_pretrained(TEACHER).to("cuda").eval()
68+
teacher.generation_config.max_length = 100_000
69+
70+
table = pq.read_table("/opt/rababa/data/sadeed-diac-25/train.parquet")
71+
inputs = [strip_diacritics(t) for t in table.column("input").to_pylist()]
72+
73+
preds = windowed_paragraphs(teacher, tok, inputs, window=1400)
74+
75+
out = Path(OUT)
76+
out.parent.mkdir(parents=True, exist_ok=True)
77+
with open(out, "w", encoding="utf-8") as f:
78+
for i, (src, pred) in enumerate(zip(inputs, preds, strict=True)):
79+
f.write(
80+
json.dumps({"idx": i, "src": src, "teacher": pred}, ensure_ascii=False)
81+
+ "\n"
82+
)
83+
CHECKPOINTS.commit()
84+
return {"rows": len(preds), "out": OUT}
85+
86+
87+
@app.local_entrypoint()
88+
def main() -> None:
89+
print(teacher_preds.remote())

0 commit comments

Comments
 (0)