diff --git a/docs/EXPERIMENTS.md b/docs/EXPERIMENTS.md new file mode 100644 index 0000000..eae0c44 --- /dev/null +++ b/docs/EXPERIMENTS.md @@ -0,0 +1,149 @@ +# Experiment registry + +Pre-registered experiments: hypothesis, protocol, and gate recorded +**before** results exist. Numbers graduate to RESULTS.md only when +harness-verified; paper.adoc claims only after a RESULTS.md entry. +The registry diff is the pre-registration evidence. + +Origin: the Qwen3.8-Flash-Next technique review (2026-08-27; basis paper +arXiv 2601.21204, LongCat-Flash-Lite — verified). Full working notes: +`TODO.qwen-next/` in the rababa working tree (uncommitted by design). + +--- + +## E1 — Margin-aware parity gates + +- **Status:** implemented + validated across every shipped artifact; + policy adopted. +- **Hypothesis:** byte students have flat top-1 margins, so quantization + flips near-tie argmaxes at KLD ~1e-5 — invisible to the CER-delta + release gate. +- **Protocol:** teacher-forced forward on both sides (torch decoder vs + ONNX zip graphs) over the parity probe pairs (first 300 test pairs per + model); per-position top1−top2 reference margins, argmax flip rate, + KL(reference||zip), share of flips below the corpus p10 margin. + `imf.parity.run_margin_analysis`; emitted by every `modal_export + parity` gate and standalone via `modal_export margins` (read-only for + published zips; JSONs on secryst-models:/imf//). +- **Measured across the catalog (2026-08-27):** + +| Artifact | Precision | Flips | Rate | KLD | ref. margin p50 | near-tie share | +|---|---|---|---|---|---|---| +| fas-g2p-1.0 | fp32 | 0/15,103 | 0.00% | 0 | 0.200 | — | +| tha-g2p-base-1.0 | fp32 | 0/6,932 | 0.00% | 0 | 0.421 | — | +| heb-diac-1.1 | fp16 | 6/34,178 | 0.02% | ~0 | 0.205 | 1.00 | +| urd-diac-1.0 | fp16 | 0/7,870 | 0.00% | 0 | 0.213 | — | +| urd-g2p-1.0 | fp16 | 3/6,281 | 0.05% | ~0 | 0.169 | 1.00 | +| khm-latn-1.0 | fp16 | 84/2,869 | 2.93% | 6.3e-06 | 0.121 | 0.77 | +| **heb-diac-1.1** | **int8** | **3,193/34,178** | **9.34%** | 2.4e-05 | 0.205 | **0.20** | +| urd-g2p-1.0 | int8 | 107/6,281 | 1.70% | 5.1e-06 | 0.169 | 0.97 | +| khm-latn-1.0 | int8 | 71/2,869 | 2.47% | 1.1e-05 | 0.121 | 0.90 | +| tha-g2p-small-1.0 | int8 | 63/6,932 | 0.91% | 1.2e-03 | 0.421 | 0.89 | +| urd-diac-1.0 | int8 | 27/7,870 | 0.34% | 1.6e-05 | 0.213 | 1.00 | +| tha-g2p-small-1.0 | int4 | 18/6,932 | 0.26% | 4.6e-04 | 0.421 | 1.00 | + +All rows passed the CER parity gate at release. Readings: +- fp32 is exact everywhere (harness sanity); fp16 is benign (≤0.05%) + except khm-latn — the flattest margins in the catalog (p50 0.121), + 2.93% flips, 77% near-tie. +- **heb-diac-1.1 int8 is the outlier: 9.34% flip rate with only 20% of + flips at near-tie positions** — 80% of its argmax flips occur where + the reference was confident. That is the dangerous class the CER gate + cannot see. +- **Root cause found and fixed at the export default (2026-08-27, + controlled probes on the same 300 pairs):** the culprit is the + quantized *head*. Per-channel weights alone: 8.50% flips, 78% still + confident-position, +25% artifact size — rejected. Keeping + `/lm_head/MatMul` in fp32 (body int8): **0.26% flips (36x fewer), + KLD 47x lower, 100% of remaining flips near-tie, +0.4% size**; + head-fp32 + per-channel adds nothing. Quantizing the node that + computes argmax moves the decision boundary directly. + `export_zips` now excludes the head MatMul from int8 by default + (`imf.export.head_matmul_names` + `nodes_to_exclude`). Shipped int8 + zips predate this; re-exporting them is a release decision. +- The distilled students behave as the decode analysis predicts: flat + but consistent — tha-g2p-small's shipped int4 flips 0.26% of + positions, all near-tie. +- **Policy (adopted):** embedding-like tensors — byte embeddings, tied + lm_head, relative-attention bias, and any memory-layer lookup tables — + are a separate quantization class: fp16 (≥ int8 floor) when the body + is quantized. Precision floors are keyed on access pattern, not + tensor size (the Qwen/Unsloth lesson). + +## E2 — PKM memory-layer student (ara-diac-small run-003-pkm) + +- **Status:** COMPLETE (2026-08-28). **Outcome: positive but below the + pre-registered win bar.** +- **Measured (full 1,200 paragraphs, windowed zero-skip, greedy):** + PKM student **7.5553** DER-CE vs run-002 vanilla ByT5-small **8.259** + (teacher reproduces 2.5815) — 0.704pp of the 5.677pp teacher-student + gap closed (12.4% relative) at +85.9M lookup params and near-zero + added compute (gathers, not matmuls). +- **Verdict per the pre-agreed rule (≥1.0pp = win): NOT met.** The + capacity axis is real but not the dominant term of the gap; the rest + is modeling/optimization (E3 tests the optimization half) and domain + coverage (the subset lesson from run-002). Reported exactly as + measured — no threshold-moving. +- Engagement was verified independent of outcome: gates moved off zero + by step-500 (0.0008-0.0028) and settled at 0.034-0.053 by step-10,500 + (15-20x) — the memory branch contributed measurably but modestly, + matching the 0.70pp outcome. CE 2.07 → 0.02 over 10,995 steps. +- Launched 2026-08-27, A10G, labels reused from run-002 (single-variable + design). Survived one mid-run eviction: resumed from step-2000. +- **Hypothesis:** the 5.68pp teacher→student gap (r6 2.5815 → ByT5-small + 8.259 full-set) is partly a *capacity* gap that lookup memory closes + at near-zero compute — parameters and compute are separable (arXiv + 2601.21204; PKM on character-level LM: Lample et al., NeurIPS 2019). +- **Design:** google/byt5-small backbone + 3 product-key memory layers + on decoder blocks [-2, -4, -6] (128² = 16,384 slots, top-32 reads, + ~+25M params), zero-init output gates (pretrained function preserved + at step 0 — tested). Teacher r6 frozen; identical corpus, labels, + limits, seed, and 3-epoch schedule as run-002. +- **Protocol:** windowed zero-skip Misraj DER-CE, full 1,200 paragraphs + (the published harness; `modal_distill evaluate_der`). +- **Pre-agreed gate:** ≤ 3.07 windowed DER-CE (the run-002 gate). +- **Pre-agreed verdict rule:** PKM wins if it closes ≥ 1.0pp of the + 5.68pp full-set gap at equal decode-time compute. No movement ⇒ the + gap is modeling/optimization, not capacity — publishable negative. +- **Comparison targets:** run-002 student 8.259 / teacher 2.5815 + (full-set, published in RESULTS.md). + +## E3 — Muon optimizer A/B (run-004-pkm-muon) + +- **Status:** COMPLETE (2026-08-28). **Outcome: ADOPTED — the gate is + exceeded 9x.** +- **Measured (full 1,200, windowed zero-skip, greedy):** Muon arm + **4.8287** DER-CE vs the single-variable AdamW arm (run-003-pkm) + 7.555 — **−2.727pp from the optimizer alone** (adopt gate ≥0.3pp). + Against the shipped vanilla student (8.259): 3.430pp of the 5.677pp + canonical gap closed (60.4%). Teacher reproduced at 2.5997 in this + container (range across eval containers 2.5793–2.5997, bf16 + autocast; protocol consistent). +- Training-side corroboration: CE ~0.007 vs the AdamW arm's ~0.02 at + equal steps; step time ~1.2s vs ~3.4s (the <15% overhead gate met + with margin — Newton–Schulz is cheap next to 1450-byte windows); no + stability events. +- **Combined with E2, the gap decomposition at the ByT5-small rung:** + capacity ~0.70pp (PKM) + optimization ~2.73pp (Muon, on the PKM + architecture) + residual ~2.25pp (domain coverage — the subset + lesson). Caveat carried: the A/B isolates the optimizer ON the + memory architecture; the vanilla+Muon cell (run-005-muon) is queued + to complete the 2x2 factorial. +- **Hypothesis:** orthogonalized-momentum updates (Newton–Schulz; the + optimizer Qwen3.8-Flash-Next / LongCat report) help even in + knowledge-limited distillation fine-tunes — unmeasured territory for + byte-level seq2seq students. Prior expectation: small (RL negative; + data-side levers won before). +- **Design:** identical to E2 except optimizer — Muon on 2D hidden + matrices (lr 0.01, momentum 0.95, wd 0.01, cosine), AdamW group for + embedding-like params including the memory tables (random-access + class per E1 policy). Same seed, data, schedule. +- **Pre-agreed adopt gate:** ≥ 0.3pp DER improvement at equal steps, no + stability regressions, step-time overhead < 15%. +- **Report:** either direction goes to the paper's training-methods + appendix. + +## Parked + +- **Speculative decoding** (LongCat converts sparsity→speed): revisit + only if API latency data shows p95 decode binding. No code, by design. diff --git a/docs/PUBLICATION-NOTES.md b/docs/PUBLICATION-NOTES.md new file mode 100644 index 0000000..3bdbf05 --- /dev/null +++ b/docs/PUBLICATION-NOTES.md @@ -0,0 +1,103 @@ +# Publication notes — what is worth publishing, and where it lives + +Status: 2026-08-28. Rule: a claim is publishable when it is +harness-measured and recorded in RESULTS.md (models/protocol numbers) +or EXPERIMENTS.md (registered experiments). Nothing below cites an +un-landed number. + +## Complete, measured, citable now + +### 1. The artifact contract (IMF v1) and cross-runtime parity +One zip, three ONNX graphs, fixed byte table, per-member SHA-256, +byte-identical output from Ruby/Python/TypeScript; opset-14 floor for +old consumer runtimes. Paper: sections 3–4 (section-imf). The framing +contribution — neural models under the same discipline as +deterministic transliteration maps. + +### 2. The decode-protocol correction +Beam-4 with length normalization inflates flat byte-student PER 4.2× +(12.06 published vs 2.85 real, same artifact, greedy). Language- +dependent at the teacher tier (Hebrew gains ~12 DER points from beam, +Arabic nothing). Paper: section-decode. This is the paper's most +quotable calibration result. + +### 3. Pretrained-or-collapse (the client-tier frontier law) +Random-init byte students collapse at every capacity tested (33M +75.80 PER, 70M 78.51); the pretrained rung is the whole cliff (300M +→ 2.85). The Arabic replica was label-corrupted in transit — +disclosed, retracted as evidence, kept as a reproducibility lesson. +Paper: section-frontier + section-repro. + +### 4. Controlled aux-representation ablation (r8, 2026-08-27) +Identical teacher/corpus/seed/init; only the aux stream's output +representation varies. Morphology 2.5793 < phonemic IPA 2.6588 < none +2.6775 (full 1,200-paragraph windowed zero-skip); the IPA projection +itself was learned (2.3% CER probe). **The strong form of "phonemes +help diacritization" fails; lexical-morphological knowledge is the +active ingredient.** Paper: leaderboard section. RESULTS.md (rababa): +r8 section. + +### 5. Margin-aware parity + the head-fp32 quantization fix (E1) +Teacher-forced margin analysis (flip rate / KLD / near-tie share) +measured across the entire catalog — invisible to the CER gate. +Headline diagnosis: the shipped Hebrew int8 artifact flipped 9.34% of +positions, 80% at confident margins. Controlled probe matrix isolated +the quantized *head* (not weight granularity: per-channel recovered +almost nothing at +25% size); keeping the 1.2M-parameter head in fp32 +cut flips 36× (0.26%, all near-tie) at +0.4% size. Now the export +default + a release gate + a regression test. Paper: IMF section +"Margin-aware parity" bullet. EXPERIMENTS.md E1. + +### 6. The memory-layer capacity experiment (E2, 2026-08-28) +Single-variable: shipped ByT5-small student vs +3 product-key memory +layers (+85.9M params, near-zero compute; gates verified engaged). +8.259 → 7.555 full-set DER — 0.704pp of the 5.677pp teacher-student +gap (12.4% relative), **below the pre-registered ≥1.0pp bar**. Honest +conclusion: the distillation gap is dominated by optimization and +domain coverage, not parameter capacity. First controlled test of the +LongCat/Qwen "embedding scaling" axis on a byte-level seq2seq student. +Paper: frontier section (fourth question). RESULTS.md run-003-pkm. + +### 7. Measurement-discipline findings +- Subset-selection artifact: first-300-paragraphs 3.66 vs full-set + 8.26 (teacher reproduces 2.5815/2.5793 — harness soundness proven). + Paper: leaderboard section. Standing rule: full-set-only publication. +- Leaderboard positioning: r6 2.5793 is the best dedicated model under + the protocol (behind Claude-3.7-Sonnet's 1.3941, ahead of GLM-5.2 + 2.6911, Gemini-Flash 3.1926, GPT-4 3.8645, Sadeed-1.5B 7.2915). + +### 8. Negative results (honesty assets, all in the log) +RL teacher polishing flat/negative ×3; microkimi bridges improve +structure but not accuracy; teacher beam-search unnecessary for +Arabic; per-channel int8 rejected on measurement; the 30 MiB tier +closed as infeasible without pretraining. + +### 9. Muon optimizer A/B on the memory student (E3) — LANDED 2026-08-28 +**4.8287 vs 7.5553 full-set (−2.727pp from the optimizer alone); adopt +gate (≥0.3pp) exceeded 9×. ADOPTED.** Training CE ~0.007 vs ~0.02 at +equal steps, ~1.2s/step vs ~3.4s, no stability events. With E2 this +completes a controlled decomposition of the distillation gap at the +ByT5-small rung: ≈0.70pp capacity + 2.73pp optimization + 2.25pp +residual (domain coverage). The strongest training-methods result in +the paper — first controlled Muon measurement on byte-level seq2seq +distillation. Paper: frontier section. EXPERIMENTS.md E3, RESULTS.md +run-004. Factorial cell 4 (vanilla+Muon, run-005) queued. + +### 10. Arabic news-domain adaptation (r7) — ID LANDED 2026-08-28 +**2.2864 / 1.3343 full-set windowed zero-skip — new best dedicated +model, −0.29pp over r6** (the news mix improved in-domain, not just +OOD; behind only Claude-3.7-Sonnet's published 1.3941). OOD half +(WikiNews-2024 multi-ref, gate: beat r6's 19.82/12.46) running via the +auto-launched actor. → rababa RESULTS.md (recorded) → paper leaderboard +table + OOD note once OOD confirms; canonical-teacher promotion +pending that check. + +## Venue fit (working notes) + +The spine is systems-with-measurements: an artifact contract + +distillation discipline + three calibration/measurement corrections +(decode, subset, quantization-margins) that generalize beyond our +stack. The frontier law (pretrained-or-collapse) and the controlled +capacity/aux/optimizer experiments give it empirical heft. Package as +one paper (current paper.adoc); the margin/head finding alone is also +a strong short workshop paper if a split is ever wanted. diff --git a/docs/RESULTS.md b/docs/RESULTS.md index 588b6f8..756b92c 100644 --- a/docs/RESULTS.md +++ b/docs/RESULTS.md @@ -222,6 +222,44 @@ mojibake labels (byt5 decode_joined bug), the second silently resumed from the poisoned lineage's checkpoints (now guarded by labels.sha digest matching), this one clean end-to-end. CE plateaued at ~0.016. +### run-003-pkm — memory-layer student (2026-08-28, research run) + +The qwen-next capacity experiment (EXPERIMENTS.md E2): identical to +run-002 except three product-key memory layers (+85.9M lookup params, +zero-init gates) on the ByT5-small decoder — single-variable. + +| Model | DER-CE (full 1,200) | +|---|---| +| Teacher (r6, full-set) | 2.5815% | +| **Student + PKM memory (run-003-pkm)** | **7.5553%** | +| Student vanilla (run-002) | 8.2590% | + +0.704pp of the 5.677pp teacher-student gap closed (12.4% relative) at +near-zero added compute — below the pre-registered ≥1.0pp win bar, so +the memory axis is real but not the dominant term of the gap. Gates +verified engaged (0.034-0.053 at completion). Not shipped; the vanilla +client rung stands. + +### run-004-pkm-muon — optimizer A/B on the memory student (2026-08-28) + +Identical to run-003-pkm except the optimizer (Muon on 2D hidden +matrices, AdamW group for embedding-like params; EXPERIMENTS.md E3): + +| Model | DER-CE (full 1,200) | +|---|---| +| Teacher (r6, this container) | 2.5997% | +| **ByT5-small + PKM + Muon (run-004)** | **4.8287%** | +| ByT5-small + PKM + AdamW (run-003) | 7.5553% | +| ByT5-small vanilla + AdamW (run-002) | 8.2590% | + +**−2.727pp from the optimizer alone** — the adopt gate (≥0.3pp) +exceeded 9x; 3.430pp of the 5.677pp canonical gap closed (60.4%) +combining memory + optimizer. Training CE ~0.007 vs ~0.02 at equal +steps; ~1.2s/step vs ~3.4s; no stability events. The teacher-student +gap at this rung decomposes: ~0.70pp capacity + ~2.73pp optimization ++ ~2.25pp residual (domain coverage). The vanilla+Muon factorial cell +(run-005-muon) completes the decomposition. + Leaderboard context (SadeedDiac-25, Misraj evaluator, zero-skip, harakat-projected DER-CE): the teacher tier (r6, 580M) at 2.5793 (reproduced at 2.5815, 2026-08-26) is the best dedicated model measured diff --git a/docs/paper.adoc b/docs/paper.adoc index d10b1f9..836caeb 100644 --- a/docs/paper.adoc +++ b/docs/paper.adoc @@ -112,6 +112,7 @@ Design decisions and their rationale: * **KV-cache decoder, self-attention only.** The streaming graph caches self-attention keys/values per layer; cross-attention K/V are recomputed each step as a deterministic projection of the encoder state. This keeps the attention-mask length consistent for any past length and spares runtimes all cross-cache bookkeeping. The graph's batch axis is genuinely dynamic — verified by feeding *different* tokens per row and matching each against its batch-1 reference. * **Per-member SHA-256, verified at load.** metadata.yaml carries a checksum map; every runtime verifies every member before use. Tampering raises; partial downloads cannot load. The metadata block itself is excluded (it carries the parity result written post-gate) — the graphs are what is integrity-protected. * **Precision-aware parity.** The strict release gate compares torch-reference greedy decode against the ONNX artifact's decode over a held-out set, with the bound keyed on declared precision: 0.2pp character-error delta at fp32, 1.0 at fp16, 2.0 at int8, 3.0 at int4. Measured artifacts: 0.0pp (fp32, several models), 0.08pp (int8), 0.07pp (int4). +* **Margin-aware parity.** A CER-delta gate measures what already broke, not what is about to: byte-level models have flat top-1 margins (median top1−top2 logit gap 0.12–0.42 across our catalog), so quantization noise can silently flip near-tie argmaxes at KL divergences (~1e-5) far below anything a decode-comparison gate detects. Every gate therefore also runs a teacher-forced margin analysis — argmax flip rate, KL divergence, and the share of flips at near-tie positions — across both sides of the artifact. Deploying it retroactively across the catalog surfaced exactly this failure: the shipped Hebrew int8 artifact flipped 9.34% of positions with 80% of flips at *confident* positions. A controlled probe matrix isolated the cause — not weight-scale granularity (per-channel quantization recovered almost nothing at +25% size) but the quantized *head*: quantizing the MatMul that computes argmax moves the decision boundary directly. Keeping that one 1.2M-parameter projection in fp32 (body int8) cut flips 36× to 0.26% — all remaining flips near-tie — at +0.4% artifact size; the export path now excludes the head from int8 quantization by default. The general lesson matches the decode study: gates must measure the protocol's actual decision surface (argmax under the shipped numerical path), not a downstream aggregate that can stay flat while the surface degrades. * **Parts for large artifacts.** GitHub caps release assets at 2 GiB. Artifacts above a threshold split into independently checksummed parts; all three runtimes stream parts in order, verify each, reassemble, and verify the whole. Consumers see one logical download. * **opset 14.** The floor imposed by the oldest consumer runtime's bundled ONNX Runtime. All graphs are exported and validated at exactly this opset. @@ -212,6 +213,8 @@ External claims are made only where a public benchmark exists and the full proto The teacher is the best dedicated (task-trained, runnable-locally) model measured under this protocol — second only to a frontier proprietary LLM's published number, ahead of a clean GLM-5.2 reproduction, Gemini Flash, and GPT-4 — and it does so at 580M parameters against Sadeed's 1.5B. +The r6 teacher's auxiliary-task design is validated by a controlled ablation: an identical run differing only in the auxiliary stream's output representation — broad-phonemic IPA of the same training units (deterministic converter) in place of morphological analysis — scores 2.6588 (vs 2.6775 for no auxiliary task and 2.5793 for the morphological one), while a held-out probe confirms the phonemic projection itself was learned (2.3% CER on the IPA stream). The improvement is therefore attributable to lexical-morphological knowledge injection, not to phonemic supervision per se: the strong form of the "phonemes help diacritization" hypothesis fails where its weak form (any structured auxiliary projection beats none) barely holds. + The client student's row carries a measurement lesson we record rather than bury. Its first-published figure, 3.66, was measured on the first 300 paragraphs of the benchmark; the full 1,200-paragraph run scores 8.26 (teacher reproduces at 2.5815 against its documented 2.5793, confirming protocol consistency). The subset was not representative: its paragraphs sit closer to the student's training domain, and the remaining 900 expose a domain-generalization gap the subset hid. Two practices follow, now standing rules: student-tier numbers are published only from full benchmark sets, and a subset figure is labeled as such at first publication. The student ships with the full-set number disclosed — behind Sadeed-1.5B on the leaderboard — and closing its domain gap (more diverse label coverage, or on-policy distillation; <>) is open work. Two protocol notes the benchmark's own authors make and we adopt: self-reported numbers from other protocols (Sadeed's repository reports 1.2% DER under its own split) are not comparable, and multi-reference evaluation (Mohamed & Mubarak 2025) is the honest treatment of valid-alternative diacritizations; on WikiNews-2024 our teacher family scores 19.82 WER against QCRI's in-domain-trained 2.70 — a domain gap (classical hadith vs news text), recorded rather than averaged away. @@ -248,6 +251,19 @@ Three findings: The engineering consequence: the client tier ships at ByT5-small, quantized. A true 30–70 MiB tier requires a *pretrained* backbone at that scale — byte-level pretraining of a narrow model — which remains open future work. The quantization ladder is the practical size lever today: 246 MiB at int8, 193 MiB at int4, with the 4-bit quality cost measured at +0.17pp CER and certified by the parity gate. +A fourth question completes the frontier analysis: at the shipped rung, is the residual teacher–student gap a *parameter-capacity* gap? Recent sparse-memory architectures (product-key memory, Lample et al. 2019; embedding scaling, Liu et al. 2026) argue parameters and compute are separable — lookup capacity at near-zero FLOPs. We tested this directly on the Arabic client tier with a single-variable design: identical teacher, corpus, teacher-labels, seed, and schedule as the shipped student, plus three product-key memory layers (+85.9M parameters, 16,384 slots each, top-32 reads) injected behind zero-initialized gates on the decoder FFNs (the pretrained function is preserved exactly at step 0, and the gates are verified to engage — settling at 0.03–0.05). + +[%autowidth,cols="1,1,1"] +|=== +|Student (Arabic, identical conditions) |Params |DER-CE (full 1,200) + +|ByT5-small (shipped) |300M |8.259 +|**+ 3 PKM memory layers** |386M |**7.555** +|r6 teacher |580M |2.582 +|=== + +Memory closes 0.70pp of the 5.68pp teacher–student gap — a real, controlled improvement at near-zero added compute, but below the pre-registered 1.0pp bar under which we would have called the capacity axis *the* fix. The same pre-registration discipline then isolated the *optimization* term: an identical arm differing only in optimizer — Muon (orthogonalized momentum via Newton–Schulz) on the hidden matrices, AdamW for embedding-like parameters — scored **4.829**, a further −2.73pp. The gap at this rung decomposes into ≈0.7pp capacity + ≈2.7pp optimization + ≈2.2pp residual (domain coverage; the full-benchmark subset lesson of <> measures exactly this exposure). Two practical notes travel with the number: Muon's training CE was ~3× lower at equal steps *and* ~2.8× faster per step on this workload (1450-byte windows dominate; Newton–Schulz is cheap next to them), and the adopt gate (≥0.3pp, set before training) was exceeded nine-fold. Optimization — not parameter count — is the dominant recoverable term of the distillation gap at this rung; capacity is a modest second lever, now measured. (This speaks to the *distillation* gap on a pretrained backbone; it does not reopen the from-scratch collapse of finding 1, which is an initialization effect.) + == The decode protocol is part of the measurement [[section-decode]] @@ -402,3 +418,5 @@ budget is on the order of 100 A10G-hours. * Bertina, A., Beirami, S., Biniazian, H., Esmaeilnia, E., Shahi, S., Pirnia, M. _Bridging the gap: An intermediate language for enhanced and cost-effective grapheme-to-phoneme conversion with homographs with multiple pronunciations disambiguation._ arXiv:2505.06599, 2025. * Rosenthal, A., Shaked, N. _D-Nikud: Enhancing Hebrew diacritization with LSTM and pretrained models._ arXiv:2402.00075, 2024. * Elboher, Y., Pinter, Y. _Hebrew diacritics restoration using visual representation._ arXiv:2510.26521, 2025. +* Lample, G., Sablayrolles, A., Ranzato, M., Denoyer, L., Jégou, H. _Large memory layers with product keys._ NeurIPS 2019 (arXiv:1907.05242). +* Liu, H. et al. _Scaling embeddings outperforms scaling experts in language models._ arXiv:2601.21204, 2026. diff --git a/src/gpu/modal_distill.py b/src/gpu/modal_distill.py index 7899372..547ffd4 100644 --- a/src/gpu/modal_distill.py +++ b/src/gpu/modal_distill.py @@ -55,6 +55,25 @@ ) CHECKPOINTS = modal.Volume.from_name("rababa-checkpoints") + + +def _ensure_src_path() -> None: + # Modal copies the entry file to /root/.py while the repo image + # sits at /root/interscript-ml — cover both layouts before importing + # sibling modules (gpu.pkm, gpu.muon). + import sys + from pathlib import Path + + for cand in ( + Path(__file__).resolve().parent.parent, + Path.cwd() / "src", + Path("/root/interscript-ml/src"), + ): + if (cand / "gpu" / "pkm.py").exists(): + if str(cand) not in sys.path: + sys.path.insert(0, str(cand)) + return + raise RuntimeError("src/gpu/pkm.py not found on any known layout") DATASETS = modal.Volume.from_name("rababa-datasets") SECRYST_CHECKPOINTS = modal.Volume.from_name("secryst-checkpoints") SECRYST_DATASETS = modal.Volume.from_name("secryst-datasets") @@ -122,6 +141,70 @@ "mode": "sequence", "note": "r6 canonical (2.5793 DER); gate <= 3.07 windowed DER-CE", }, + "ara-diac-small-pkm": { + # TODO.qwen-next/02 — the LongCat/Qwen capacity axis: keep the + # ByT5-small compute, add product-key lookup memory (+~25M + # params). Everything else identical to run-002 (teacher, corpus, + # labels, seed) so the comparison is single-variable. + "teacher": "rababa_arabic_byt5/run-006-morph/best", + "teacher_volume": "rababa", + "student_init": "google/byt5-small", + # byt5-small has only 4 decoder blocks (depth lives in the + # encoder) — all but the first carry memory + "pkm": {"layer_indices": [-1, -2, -3], "n_keys": 128, "topk": 32}, + "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-003-pkm", + "labels_file": "teacher_labels_v2.jsonl", + "labels_complete": "true", + "mode": "sequence", + "note": "PKM memory student; gate <= 3.07 windowed DER-CE; " + "verdict vs run-002's 8.259 full-set (closes >= 1.0pp?)", + }, + "ara-diac-small-pkm-muon": { + # TODO.qwen-next/03 — identical to ara-diac-small-pkm except the + # optimizer (Muon + AdamW groups). The A/B pair is run-003-pkm + # (AdamW) vs run-004-pkm-muon. + "teacher": "rababa_arabic_byt5/run-006-morph/best", + "teacher_volume": "rababa", + "student_init": "google/byt5-small", + "pkm": {"layer_indices": [-1, -2, -3], "n_keys": 128, "topk": 32}, + "optimizer": "muon", + "muon_lr": "0.01", + "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-004-pkm-muon", + "labels_file": "teacher_labels_v2.jsonl", + "labels_complete": "true", + "mode": "sequence", + "note": "Muon A/B arm; same gate and verdict rule as run-003-pkm", + }, + "ara-diac-small-muon": { + # factorial completion (EXPERIMENTS.md E3 caveat): the A/B was + # PKM+Muon vs PKM+AdamW; this cell measures Muon WITHOUT memory + # so the 2x2 {vanilla, pkm} x {adamw, muon} closes cleanly + "teacher": "rababa_arabic_byt5/run-006-morph/best", + "teacher_volume": "rababa", + "student_init": "google/byt5-small", + "optimizer": "muon", + "muon_lr": "0.01", + "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-005-muon", + "labels_file": "teacher_labels_v2.jsonl", + "labels_complete": "true", + "mode": "sequence", + "note": "vanilla ByT5-small + Muon (factorial cell 4)", + }, "ara-diac-tiny": { "teacher": "rababa_arabic_byt5/run-006-morph/best", "teacher_volume": "rababa", @@ -691,6 +774,11 @@ def distill_sequence(spec_id: str, epochs: int = 3) -> dict: print(f"[{spec_id}] tiny student: {n_params:.1f}M params", flush=True) else: student = AutoModelForSeq2SeqLM.from_pretrained(spec["student_init"]) + if spec.get("pkm"): + _ensure_src_path() + from gpu.pkm import inject_pkm + + inject_pkm(student, **spec["pkm"]) student.train() class Pairs(Dataset): @@ -988,7 +1076,23 @@ def __getitem__(self, i): drop_last=True, ) total_steps = len(train_loader) * epochs - optimizer = torch.optim.AdamW(student.parameters(), lr=1e-4) + if spec.get("optimizer") == "muon": + _ensure_src_path() + from gpu.muon import Muon, split_parameters + + muon_params, adamw_params = split_parameters(student.named_parameters()) + optimizer = Muon( + muon_params, lr=float(spec.get("muon_lr", 0.01)), + momentum=0.95, weight_decay=0.01, + ) + optimizer.add_adamw_group(adamw_params, lr=1e-4, weight_decay=0.0) + print( + f"[{spec_id}] muon: {len(muon_params)} matrix / " + f"{len(adamw_params)} embedding-like params", + flush=True, + ) + else: + optimizer = torch.optim.AdamW(student.parameters(), lr=1e-4) scheduler = get_cosine_schedule_with_warmup( optimizer, total_steps // 20, total_steps ) @@ -1114,7 +1218,13 @@ def evaluate_der(spec_id: str, window: int = 1400, limit: int = 0) -> dict: tok = AutoTokenizer.from_pretrained("google/byt5-small") teacher = AutoModelForSeq2SeqLM.from_pretrained(teacher_path).to("cuda").eval() - student = AutoModelForSeq2SeqLM.from_pretrained(str(student_path)).to("cuda").eval() + if spec.get("pkm"): + _ensure_src_path() + from gpu.pkm import load_student_with_pkm + + student = load_student_with_pkm(student_path, spec["pkm"]).to("cuda").eval() + else: + student = AutoModelForSeq2SeqLM.from_pretrained(str(student_path)).to("cuda").eval() # custom students carry T5's default max_length=20; windowed inputs # run to 1400 bytes, clamping max_new_tokens to zero for m in (teacher, student): @@ -1197,6 +1307,18 @@ def der_ce(model) -> dict: result = {"teacher": der_ce(teacher), "student": der_ce(student)} result["gate_delta"] = round(result["student"]["der_ce"] - result["teacher"]["der_ce"], 4) result["gate_pass"] = result["gate_delta"] <= 0.5 + # durable verdict marker: the run dir is the provenance record (also + # what r7-style _init_choice probes read) + import json + + out_root = Path( + vol_map[spec.get("out_volume", spec.get("teacher_volume", "rababa"))] + ) / spec["out"] + out_root.mkdir(parents=True, exist_ok=True) + (out_root / "final_eval.json").write_text( + json.dumps(result, indent=2), encoding="utf-8" + ) + CHECKPOINTS.commit() return result @@ -1460,3 +1582,134 @@ def mk(spec: str = "tha-g2p-tiny-mk", epochs: int = 3) -> None: @app.local_entrypoint() def eval_der(spec: str = "ara-diac-small", limit: int = 0) -> None: print(evaluate_der.remote(spec, limit=limit)) + + +@app.function( + cpu=4, + memory=16 * 1024, + timeout=30 * 60, + volumes={"/checkpoints": CHECKPOINTS, "/secryst-checkpoints": SECRYST_CHECKPOINTS}, +) +def probe_pkm_gates(spec_id: str = "ara-diac-small-pkm") -> dict: + """E2 engagement probe (r8's IPA-probe analogue): gate values and + memory-table statistics from the latest step checkpoint. Gates moving + off zero = the memory branch is being used, not bypassed.""" + from pathlib import Path + + import torch + from transformers import AutoModelForSeq2SeqLM + + spec = SPECS[spec_id] + out_root = Path("/checkpoints") / spec["out"] + ckpts = sorted( + out_root.glob("step-*"), key=lambda p: int(p.name.split("-")[1]) + ) + if not ckpts: + raise RuntimeError(f"no step checkpoints under {out_root}") + + _ensure_src_path() + from gpu.pkm import ProductKeyMemory, inject_pkm + + student = AutoModelForSeq2SeqLM.from_pretrained(spec["student_init"]) + inject_pkm(student, **spec["pkm"]) + sd = torch.load(ckpts[-1] / "student.pt", map_location="cpu", weights_only=True) + student.load_state_dict(sd) + + report: dict = {"checkpoint": ckpts[-1].name, "gates": {}, "tables": {}} + for i, block in enumerate(student.decoder.block): + wrapped = block.layer[1] + if not hasattr(wrapped, "memory"): + continue + report["gates"][f"dec-{i}"] = round(float(wrapped.gate), 4) + v = wrapped.memory.values + report["tables"][f"dec-{i}"] = { + "rows": int(v.shape[0]), + "row_norm_mean": round(float(v.norm(dim=-1).mean()), 4), + "row_norm_p99": round(float(v.norm(dim=-1).quantile(0.99)), 4), + "key_drift_q1": round(float(wrapped.memory.k1.abs().mean()), 4), + } + return report + + +@app.local_entrypoint() +def pkm_gates(spec: str = "ara-diac-small-pkm") -> None: + print(probe_pkm_gates.remote(spec)) + + +@app.function( + cpu=1, + memory=1024, + timeout=24 * 3600, + volumes={"/checkpoints": CHECKPOINTS, "/secryst-checkpoints": SECRYST_CHECKPOINTS}, +) +def qwen_next_chain() -> dict: + """Server-side orchestrator for the qwen-next experiments (E2/E3): + workstation-independent, self-healing, idempotent — the replacement + for local shell chains that die with the workstation or break when + the repo's checked-out branch changes. + + State machine per arm, driven by durable volume markers only: + best/config.json absent -> watch step-* checkpoints; respawn + training if no progress for 20 min + (distill_sequence resumes from the + latest checkpoint; a redundant spawn + is benign — a finished run saves + best again and exits) + best present, final_eval.json absent -> evaluate_der (which now + writes final_eval.json itself) + final_eval.json present -> arm done + + Audit trail: chain_log.jsonl in each run dir. If this function times + out (24h) or is evicted, relaunching continues from the markers: + + modal run --detach src/gpu/modal_distill.py::qwen_chain + """ + import json + import time + from pathlib import Path + + ARMS = [ + ("ara-diac-small-pkm", "rababa_arabic_distill_small/run-003-pkm"), + ("ara-diac-small-pkm-muon", "rababa_arabic_distill_small/run-004-pkm-muon"), + ("ara-diac-small-muon", "rababa_arabic_distill_small/run-005-muon"), + ] + ROOT = Path("/checkpoints") + + def log(run: str, event: str) -> None: + with (ROOT / run / "chain_log.jsonl").open("a", encoding="utf-8") as fh: + fh.write(json.dumps({"t": round(time.time()), "event": event}) + "\n") + CHECKPOINTS.commit() + + def latest_step(run: str) -> int: + steps = [int(p.name.split("-")[1]) for p in (ROOT / run).glob("step-*")] + return max(steps) if steps else -1 + + status = {} + for spec_id, run in ARMS: + while not (ROOT / run / "best" / "config.json").exists(): + CHECKPOINTS.reload() + before = latest_step(run) + log(run, f"watch step={before}") + time.sleep(1200) + CHECKPOINTS.reload() + after = latest_step(run) + if after == before and not (ROOT / run / "best" / "config.json").exists(): + log(run, f"stalled at step={after}; respawning {spec_id}") + distill_sequence.spawn(spec_id, epochs=3) + log(run, "training complete (best present)") + if not (ROOT / run / "final_eval.json").exists(): + log(run, "evaluating") + evaluate_der.remote(spec_id=spec_id) + log(run, "eval done") + status[run] = "complete" + return status + + +@app.local_entrypoint() +def qwen_chain() -> None: + handle = qwen_next_chain.spawn() + print( + f"spawned {handle.object_id}; durable markers: chain_log.jsonl, " + f"final_eval.json in each run dir; relaunch is idempotent", + flush=True, + ) diff --git a/src/gpu/modal_export.py b/src/gpu/modal_export.py index 0a34c29..75d29c8 100644 --- a/src/gpu/modal_export.py +++ b/src/gpu/modal_export.py @@ -248,7 +248,13 @@ def parity_model(model_id: str, precisions: list[str], limit: int = 0) -> dict[s test_path = Path(spec["test_volume"]) / spec["test_data"] from imf.export import load_byte_seq2seq - from imf.parity import reference_decode, run_parity, write_parity + from imf.parity import ( + reference_decode, + run_margin_analysis, + run_parity, + write_margin_report, + write_parity, + ) model = load_byte_seq2seq(checkpoint) pairs = _load_pairs(test_path) @@ -272,6 +278,71 @@ def parity_model(model_id: str, precisions: list[str], limit: int = 0) -> dict[s if not report.passed: raise RuntimeError(f"parity gate FAILED for {zip_path.name}") write_parity(zip_path, report) + margins = run_margin_analysis(model, zip_path, pairs, max_len=128) + write_margin_report(margins, out_dir / f"{mid}-margins-{precision}.json") + reports[precision] += ( + f" | margin flips={margins.flip_rate:.4%} kld={margins.kld_mean:.2e} " + f"p10={margins.margin_p10} low-share={margins.flip_low_margin_share}" + ) + # margin release policy (E1): near-tie flips are inherent to flat + # byte models, but confident-position flips mean the artifact's + # decision surface moved. Pre-fix heb-diac int8 measured 7.5% + # confident flips; every head-fp32 artifact measures < 0.7%. + confident_flip_rate = margins.flip_rate * (1 - margins.flip_low_margin_share) + if confident_flip_rate > 0.01: + raise RuntimeError( + f"margin gate FAILED for {zip_path.name}: " + f"{confident_flip_rate:.2%} of positions flip argmax at " + f"confident margins (budget: 1%)" + ) + MODELS_VOLUME.commit() + return reports + + +@app.function( + cpu=8, + memory=32 * 1024, + timeout=5 * 3600, + volumes={**CHECKPOINT_VOLUMES, **DATASET_VOLUMES, "/outputs": MODELS_VOLUME}, +) +def margin_model(model_id: str, precisions: list[str], limit: int = 0) -> dict[str, str]: + """Margin analysis alone over already-exported zips — read-only for the + zips (diagnostic JSON only); validates published artifacts without + touching their metadata.""" + import sys + + sys.path.insert(0, "/root/interscript-ml/src") + + spec = MODELS[model_id] + checkpoint = Path(spec["volume"]) / spec["checkpoint"] + test_path = Path(spec["test_volume"]) / spec["test_data"] + + from imf.export import load_byte_seq2seq + from imf.parity import run_margin_analysis, write_margin_report + + model = load_byte_seq2seq(checkpoint) + pairs = _load_pairs(test_path) + if limit: + pairs = pairs[:limit] + + out_dir = Path("/outputs/imf") / model_id + meta_path = Path("/root/interscript-ml", spec["metadata"]) + mid = re.search(r"^id:\s*(\S+)", meta_path.read_text(encoding="utf-8"), re.M).group(1) + reports: dict[str, str] = {} + for precision in precisions: + zip_path = out_dir / f"{mid}-{precision}.zip" + if not zip_path.exists(): + reports[precision] = "zip not exported (skipped)" + continue + report = run_margin_analysis(model, zip_path, pairs, max_len=128) + write_margin_report(report, out_dir / f"{mid}-margins-{precision}.json") + reports[precision] = ( + f"samples={report.samples} tokens={report.tokens} " + f"flips={report.flipped_tokens} ({report.flip_rate:.4%}) " + f"kld={report.kld_mean:.2e} margins p1/p10/p50=" + f"{report.margin_p1}/{report.margin_p10}/{report.margin_p50} " + f"low-margin-flip-share={report.flip_low_margin_share}" + ) MODELS_VOLUME.commit() return reports @@ -288,3 +359,191 @@ def parity(model: str, precisions: str = "fp32,fp16,int8", limit: int = 0) -> No reports = parity_model.remote(model, precisions.split(","), limit) for precision, status in reports.items(): print(f"{model} [{precision}] {status}") + + +@app.local_entrypoint() +def margins(model: str, precisions: str = "fp32,fp16,int8", limit: int = 0) -> None: + reports = margin_model.remote(model, precisions.split(","), limit) + for precision, status in reports.items(): + print(f"{model} [{precision}] {status}") + + +@app.function( + cpu=8, + memory=32 * 1024, + timeout=5 * 3600, + volumes={**CHECKPOINT_VOLUMES, **DATASET_VOLUMES, "/outputs": MODELS_VOLUME}, +) +def int8_pc_probe(model_id: str = "heb-diac", limit: int = 300) -> dict: + """E1 follow-up: does per-channel int8 remove the confident-position + argmax flips? Rebuilds the int8 graphs from the fp32 zip with + per_channel=True, packages them as a probe zip (copy of the shipped + int8 zip with graphs swapped — NOT a release artifact), and compares + margin reports on the same pairs.""" + import sys + import tempfile + import zipfile + + sys.path.insert(0, "/root/interscript-ml/src") + + spec = MODELS[model_id] + checkpoint = Path(spec["volume"]) / spec["checkpoint"] + test_path = Path(spec["test_volume"]) / spec["test_data"] + + from imf.export import load_byte_seq2seq, quantize_int8 + from imf.parity import run_margin_analysis + + model = load_byte_seq2seq(checkpoint) + pairs = _load_pairs(test_path)[:limit] + + out_dir = Path("/outputs/imf") / model_id + meta_path = Path("/root/interscript-ml", spec["metadata"]) + mid = re.search(r"^id:\s*(\S+)", meta_path.read_text(encoding="utf-8"), re.M).group(1) + fp32_zip = out_dir / f"{mid}-fp32.zip" + int8_zip = out_dir / f"{mid}-int8.zip" + if not fp32_zip.exists() or not int8_zip.exists(): + raise RuntimeError(f"need both {fp32_zip.name} and {int8_zip.name} on the volume") + + shipped = run_margin_analysis(model, int8_zip, pairs, max_len=128) + + with tempfile.TemporaryDirectory() as tmp: + tmp = Path(tmp) + with zipfile.ZipFile(fp32_zip) as zf: + zf.extract("encoder.onnx", tmp) + dec = "decoder-kv.onnx" if "decoder-kv.onnx" in zf.namelist() else "decoder.onnx" + zf.extract(dec, tmp) + enc_pc = tmp / "encoder-pc.onnx" + dec_pc = tmp / dec.replace(".onnx", "-pc.onnx") + quantize_int8(tmp / "encoder.onnx", enc_pc, per_channel=True) + quantize_int8(tmp / dec, dec_pc, per_channel=True) + + probe_zip = tmp / f"{mid}-int8-pc-probe.zip" + with zipfile.ZipFile(int8_zip) as src, zipfile.ZipFile( + probe_zip, "w", zipfile.ZIP_DEFLATED + ) as dst: + for name in src.namelist(): + if name == "encoder.onnx": + dst.writestr(name, enc_pc.read_bytes()) + elif name == dec: + dst.writestr(name, dec_pc.read_bytes()) + else: + dst.writestr(name, src.read(name)) + + per_channel_report = run_margin_analysis(model, probe_zip, pairs, max_len=128) + size_shipped = int8_zip.stat().st_size + size_probe = probe_zip.stat().st_size + + def row(r): + return { + "flips": r.flipped_tokens, "tokens": r.tokens, + "flip_rate": r.flip_rate, "kld_mean": r.kld_mean, + "flip_low_margin_share": r.flip_low_margin_share, + } + + return { + "model": model_id, "pairs": len(pairs), + "shipped_int8": row(shipped), "per_channel_int8": row(per_channel_report), + "size_bytes": {"shipped": size_shipped, "per_channel": size_probe}, + } + + +@app.local_entrypoint() +def int8_pc(model: str = "heb-diac", limit: int = 300) -> None: + print(int8_pc_probe.remote(model, limit)) + + +@app.function( + cpu=8, + memory=32 * 1024, + timeout=5 * 3600, + volumes={**CHECKPOINT_VOLUMES, **DATASET_VOLUMES, "/outputs": MODELS_VOLUME}, +) +def int8_head_probe(model_id: str = "heb-diac", limit: int = 300) -> dict: + """E1 follow-up 2: per-channel alone did NOT fix heb-diac's 9.3% + flip rate (8.5% remaining, 78% still at confident positions). This + probe keeps the logits-producing MatMul (the tied head) in fp32 and + quantizes only the body — per-tensor and per-channel variants.""" + import sys + import tempfile + import zipfile + + sys.path.insert(0, "/root/interscript-ml/src") + + spec = MODELS[model_id] + checkpoint = Path(spec["volume"]) / spec["checkpoint"] + test_path = Path(spec["test_volume"]) / spec["test_data"] + + from imf.export import head_matmul_names, load_byte_seq2seq + from imf.parity import run_margin_analysis + + model = load_byte_seq2seq(checkpoint) + pairs = _load_pairs(test_path)[:limit] + + out_dir = Path("/outputs/imf") / model_id + meta_path = Path("/root/interscript-ml", spec["metadata"]) + mid = re.search(r"^id:\s*(\S+)", meta_path.read_text(encoding="utf-8"), re.M).group(1) + fp32_zip = out_dir / f"{mid}-fp32.zip" + int8_zip = out_dir / f"{mid}-int8.zip" + if not fp32_zip.exists() or not int8_zip.exists(): + raise RuntimeError(f"need both {fp32_zip.name} and {int8_zip.name} on the volume") + + shipped = run_margin_analysis(model, int8_zip, pairs, max_len=128) + + results: dict = { + "model": model_id, "pairs": len(pairs), + "shipped_int8": { + "flips": shipped.flipped_tokens, "tokens": shipped.tokens, + "flip_rate": shipped.flip_rate, "kld_mean": shipped.kld_mean, + "flip_low_margin_share": shipped.flip_low_margin_share, + }, + } + + with tempfile.TemporaryDirectory() as tmp: + tmp = Path(tmp) + with zipfile.ZipFile(fp32_zip) as zf: + zf.extract("encoder.onnx", tmp) + dec = "decoder-kv.onnx" if "decoder-kv.onnx" in zf.namelist() else "decoder.onnx" + zf.extract(dec, tmp) + + head_nodes = head_matmul_names(tmp / dec) + results["head_nodes_excluded"] = head_nodes + + from onnxruntime.quantization import QuantType, quantize_dynamic + + for variant, per_channel in (("head32", False), ("head32_pc", True)): + enc_q = tmp / f"encoder-{variant}.onnx" + dec_q = tmp / dec.replace(".onnx", f"-{variant}.onnx") + quantize_dynamic( + str(tmp / "encoder.onnx"), str(enc_q), + weight_type=QuantType.QInt8, op_types_to_quantize=["MatMul"], + per_channel=per_channel, + ) + quantize_dynamic( + str(tmp / dec), str(dec_q), + weight_type=QuantType.QInt8, op_types_to_quantize=["MatMul"], + per_channel=per_channel, nodes_to_exclude=head_nodes, + ) + probe_zip = tmp / f"{mid}-int8-{variant}-probe.zip" + with zipfile.ZipFile(int8_zip) as src, zipfile.ZipFile( + probe_zip, "w", zipfile.ZIP_DEFLATED + ) as dst: + for name in src.namelist(): + if name == "encoder.onnx": + dst.writestr(name, enc_q.read_bytes()) + elif name == dec: + dst.writestr(name, dec_q.read_bytes()) + else: + dst.writestr(name, src.read(name)) + report = run_margin_analysis(model, probe_zip, pairs, max_len=128) + results[f"int8_{variant}"] = { + "flips": report.flipped_tokens, "tokens": report.tokens, + "flip_rate": report.flip_rate, "kld_mean": report.kld_mean, + "flip_low_margin_share": report.flip_low_margin_share, + "size_bytes": probe_zip.stat().st_size, + } + return results + + +@app.local_entrypoint() +def int8_head(model: str = "heb-diac", limit: int = 300) -> None: + print(int8_head_probe.remote(model, limit)) diff --git a/src/gpu/muon.py b/src/gpu/muon.py new file mode 100644 index 0000000..4431dfc --- /dev/null +++ b/src/gpu/muon.py @@ -0,0 +1,119 @@ +"""Muon optimizer (single file): orthogonalized momentum via +Newton–Schulz for 2D weight matrices, with AdamW-fallback groups. + +From Keller Jordan's Muon (modded-nanogpt); the recipe Qwen3.8-Flash-Next +/ LongCat-Flash-Lite report training with. Rules of use: +- Matrices that are updated as WHOLE weights (attention/FFN projections) + get Newton–Schulz-orthogonalized momentum. +- Embedding-like tensors — byte embeddings, layer norms, the tied + lm_head, relative-attention bias, and memory-layer lookup tables + (random access per the quantization-class policy) — stay on AdamW + inside the same optimizer object, via groups flagged ``adamw=True``. + +State save/resume works through the standard ``torch.optim`` dict. +""" + +from __future__ import annotations + +import torch + + +def zeropower_via_newtonschulz5(g: torch.Tensor, steps: int = 5) -> torch.Tensor: + # quintic iteration coefficients from the modded-nanogpt lineage + a, b, c = 3.4445, -4.7750, 2.0315 + x = g.bfloat16() + x = x / (x.norm() + 1e-7) + transposed = g.size(-2) > g.size(-1) + if transposed: + x = x.mT + for _ in range(steps): + a_mat = x @ x.mT + b_mat = b * a_mat + c * (a_mat @ a_mat) + x = a * x + b_mat @ x + if transposed: + x = x.mT + return x.to(g.dtype) + + +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), + dict(lr=lr, momentum=momentum, nesterov=nesterov, + ns_steps=ns_steps, weight_decay=weight_decay, adamw=False), + ) + + def add_adamw_group(self, params, lr: float = 1e-4, betas=(0.9, 0.999), + weight_decay: float = 0.0) -> None: + """Embedding-like parameters: standard AdamW math, shared + scheduler (the cosine scales every group's lr).""" + self.add_param_group(dict(params=list(params), lr=lr, betas=tuple(betas), + weight_decay=weight_decay, adamw=True)) + + @torch.no_grad() + def step(self, closure=None): # noqa: ARG002 + for group in self.param_groups: + if group.get("adamw"): + self._adamw_step(group) + else: + self._muon_step(group) + + 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 + 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 _adamw_step(self, group) -> None: + beta1, beta2 = group["betas"] + for p in group["params"]: + if p.grad is None: + continue + state = self.state[p] + if "adamw_exp_avg" not in state: + state["adamw_exp_avg"] = torch.zeros_like(p.grad) + state["adamw_exp_avg_sq"] = torch.zeros_like(p.grad) + state["adamw_step"] = 0 + exp_avg, exp_avg_sq = state["adamw_exp_avg"], state["adamw_exp_avg_sq"] + state["adamw_step"] += 1 + exp_avg.lerp_(p.grad, 1 - beta1) + exp_avg_sq.mul_(beta2).addcmul_(p.grad, p.grad, value=1 - beta2) + bias_c1 = 1 - beta1 ** state["adamw_step"] + bias_c2 = 1 - beta2 ** state["adamw_step"] + denom = (exp_avg_sq / bias_c2).sqrt_().add_(1e-8) + if group["weight_decay"]: + p.mul_(1 - group["lr"] * group["weight_decay"]) + p.addcdiv_(exp_avg / bias_c1, denom, value=-group["lr"]) + + +def split_parameters(named_params): + """The standard split: orthogonalizable 2D hidden weights vs + embedding-like tensors (1D params, embeddings, tied head, relative + bias, memory lookup tables).""" + muon, adamw = [], [] + for name, p in named_params: + if not p.requires_grad: + continue + embedding_like = ( + p.ndim < 2 + or "embed_tokens" in name + or name == "shared.weight" # T5 tied byte embedding (transformers 5.x) + or "lm_head" in name + or "relative_attention" in name + or "memory.values" in name + or "memory.k1" in name + or "memory.k2" in name + ) + (adamw if embedding_like else muon).append(p) + return muon, adamw diff --git a/src/gpu/pkm.py b/src/gpu/pkm.py new file mode 100644 index 0000000..a5ce908 --- /dev/null +++ b/src/gpu/pkm.py @@ -0,0 +1,129 @@ +"""Product-key memory layers for byte-level students. + +The Qwen3.8-Flash-Next / LongCat-Flash-Lite capacity axis (arXiv +2601.21204): parameters and compute are separable — a lookup memory adds +knowledge capacity at near-zero FLOPs. Precedent at character level: +Lample et al., "Large Memory Layers with Product Keys", NeurIPS 2019. + +Design notes: +- Injected as a parallel residual branch on decoder FFNs: + ``y = FFN(x) + g * mem(LN(x))`` with the gate ``g`` zero-initialized, so + a pretrained backbone's function is preserved exactly at step 0 + (ReZero-style bootstrap; memory parameters receive gradient once the + gate moves). +- The value table is a random-access tensor: it belongs to the + embedding-like quantization class (see TODO.qwen-next/01) and stays on + AdamW in the Muon split (TODO.qwen-next/03). +""" + +from __future__ import annotations + +import torch +import torch.nn as nn + + +class ProductKeyMemory(nn.Module): + """Sparse memory read: two half-codebooks of ``n_keys`` keys span + ``n_keys**2`` slots; per position, top-k candidates from each half + combine into a k*k grid from which the final ``topk`` slots are + gathered and softmax-weighted.""" + + def __init__(self, d_model: int, n_keys: int = 128, topk: int = 32) -> None: + super().__init__() + if d_model % 2: + raise ValueError(f"d_model must be even, got {d_model}") + self.n_keys = n_keys + self.topk = topk + d_k = d_model // 2 + self.ln = nn.LayerNorm(d_model) + self.wq = nn.Linear(d_model, d_model, bias=False) + self.k1 = nn.Parameter(torch.randn(n_keys, d_k) / d_k**0.5) + self.k2 = nn.Parameter(torch.randn(n_keys, d_k) / d_k**0.5) + self.values = nn.Parameter(torch.randn(n_keys * n_keys, d_model) / d_model**0.5) + self.wo = nn.Linear(d_model, d_model, bias=False) + self.scale = d_k**-0.5 + + def forward(self, x): + h = self.ln(x) + q1, q2 = self.wq(h).chunk(2, dim=-1) + s1 = torch.einsum("btd,cd->btc", q1, self.k1) * self.scale + s2 = torch.einsum("btd,cd->btc", q2, self.k2) * self.scale + t1, i1 = s1.topk(self.topk, dim=-1) + t2, i2 = s2.topk(self.topk, dim=-1) + cand = t1[:, :, :, None] + t2[:, :, None, :] + w, idx = cand.flatten(-2).topk(self.topk, dim=-1) + a, b = idx // self.topk, idx % self.topk + slot = torch.gather(i1, -1, a) * self.n_keys + torch.gather(i2, -1, b) + v = self.values[slot] # (B, T, topk, d_model) + mem = torch.einsum("btk,btkd->btd", torch.softmax(w, dim=-1), v) + return self.wo(mem) + + +class _FFNWithMemory(nn.Module): + """Wraps a T5LayerFF: keeps its output contract, adds the gated + memory branch computed from the same (pre-FFN) hidden states.""" + + def __init__(self, ffn: nn.Module, memory: ProductKeyMemory) -> None: + super().__init__() + self.ffn = ffn + self.memory = memory + self.gate = nn.Parameter(torch.zeros(())) + + def forward(self, hidden_states, **kwargs): + out = self.ffn(hidden_states, **kwargs) + mem = self.gate * self.memory(hidden_states) + if isinstance(out, tuple): + return (out[0] + mem,) + out[1:] + return out + mem + + +def inject_pkm(model, layer_indices=(-2, -4, -6), n_keys: int = 128, topk: int = 32): + """Wrap decoder-block FFNs with memory branches in place. Negative + indices count from the output side — the late decoder blocks, where + lexical (table-friendly) decisions crystallize.""" + blocks = model.decoder.block + n = len(blocks) + for i in layer_indices: + idx = i if i >= 0 else n + i + if not 0 <= idx < n: + raise IndexError(f"layer index {i} out of range for {n} decoder blocks") + blocks[idx].layer[1] = _FFNWithMemory(blocks[idx].layer[1], ProductKeyMemory( + model.config.d_model, n_keys=n_keys, topk=topk)) + base = sum(p.numel() for p in model.parameters()) + mem = sum(p.numel() for m in model.modules() if isinstance(m, ProductKeyMemory) + for p in m.parameters()) + print(f"[pkm] injected {len(tuple(layer_indices))} memory layers: " + f"+{mem / 1e6:.1f}M params on a {base / 1e6:.0f}M model", flush=True) + return model + + +def load_student_with_pkm(path, pkm_cfg: dict): + """Load a PKM student saved with ``save_pretrained``: the vanilla + class ignores the injected parameters, so re-inject then pull them + from the checkpoint file. Raises if any PKM parameter is missing.""" + from transformers import AutoModelForSeq2SeqLM + + student = AutoModelForSeq2SeqLM.from_pretrained(path) + inject_pkm(student, **pkm_cfg) + sd = None + import glob + + for pattern in ("model.safetensors", "model*.safetensors", "pytorch_model.bin"): + hits = glob.glob(str(path / pattern)) + if hits: + if hits[0].endswith(".bin"): + sd = torch.load(hits[0], map_location="cpu", weights_only=True) + else: + from safetensors.torch import load_file + + sd = load_file(hits[0]) + break + if sd is None: + raise RuntimeError(f"no weight file found in {path}") + missing, unexpected = student.load_state_dict(sd, strict=False) + # tied embeddings (shared/lm_head) are intentionally absent from the + # checkpoint file — only the injected memory keys are load-bearing here + missing_pkm = [k for k in missing if "memory." in k or k.endswith(".gate")] + if missing_pkm: + raise RuntimeError(f"PKM parameters missing from checkpoint: {missing_pkm[:5]}") + return student diff --git a/src/imf/export.py b/src/imf/export.py index 12bf100..152a856 100644 --- a/src/imf/export.py +++ b/src/imf/export.py @@ -256,13 +256,24 @@ def convert_fp16(model): return copy.deepcopy(model).half() -def quantize_int8(src: Path | str, dst: Path | str) -> Path: +def quantize_int8( + src: Path | str, dst: Path | str, per_channel: bool = False, + nodes_to_exclude: list[str] | None = None, +) -> Path: """fp32 -> dynamically quantized int8 (MatMul weights QInt8). MatMul-only: quantizing other ops inserts precision casts that break ORT's session-time SimplifiedLayerNormFusion, and preprocessing the graph (quant_pre_process) pins concrete example shapes into DynamicQuantizeLinear buffers. + + nodes_to_exclude keeps the logits-producing MatMul (the tied head) in + fp32 — measured on heb-diac-1.1 (E1, docs/EXPERIMENTS.md): quantizing + the node that computes argmax moves the decision boundary directly — + 9.34% of positions flip (80% at confident margins); keeping the head + fp32 cuts that to 0.26%, all near-tie, at +0.4% artifact size. + per_channel=True is rejected: no additional benefit once the head is + excluded, and +25% size. """ from onnxruntime.quantization import QuantType, quantize_dynamic @@ -271,10 +282,42 @@ def quantize_int8(src: Path | str, dst: Path | str) -> Path: str(dst), weight_type=QuantType.QInt8, op_types_to_quantize=["MatMul"], + per_channel=per_channel, + nodes_to_exclude=nodes_to_exclude or [], ) return Path(dst) +def head_matmul_names(graph_path: Path | str) -> list[str]: + """Names of the MatMul node(s) producing the logits output — the + tied lm_head, which stays fp32 in int8 exports (see quantize_int8).""" + import onnx + + return _head_matmul_names(onnx.load(str(graph_path)).graph) + + +def _head_matmul_names(graph) -> list[str]: + out_names = {o.name for o in graph.output} + producers = {} + for node in graph.node: + for o in node.output: + producers[o] = node + + heads = [ + n.name for n in graph.node + if n.op_type == "MatMul" and set(n.output) & out_names + ] + if not heads: + # logits may sit behind an Identity/Transpose/Reshape producer + for name in out_names: + node = producers.get(name) + for inp in node.input if node else []: + p = producers.get(inp) + if p is not None and p.op_type == "MatMul": + heads.append(p.name) + return sorted(set(heads)) + + def quantize_int4(src: Path | str, dst: Path | str, block_size: int = 64) -> Path: """fp32 -> 4-bit blockwise MatMul (MatMulNBits, com.microsoft domain). @@ -401,7 +444,10 @@ def export_zips( if precision == "fp32" or precision == "fp16": dst.write_bytes(src.read_bytes()) elif precision == "int8": - quantize_int8(graphs[name], dst) + # the tied head stays fp32 (see quantize_int8) — + # decoder graphs only; the encoder has no head + exclude = head_matmul_names(src) if "decoder" in name else [] + quantize_int8(src, dst, nodes_to_exclude=exclude) elif precision == "int4": quantize_int4(graphs[name], dst) else: diff --git a/src/imf/parity.py b/src/imf/parity.py index b0567c2..62ff899 100644 --- a/src/imf/parity.py +++ b/src/imf/parity.py @@ -20,7 +20,14 @@ from pathlib import Path from framework.evaluator import char_error_rate -from imf.export import BYTE_OFFSET, EOS_ID, PAD_ID, encode_bytes, onnx_greedy_kv +from imf.export import ( + BYTE_OFFSET, + EOS_ID, + PAD_ID, + _zero_pasts, + encode_bytes, + onnx_greedy_kv, +) from imf.schema import ModelMetadata, Parity @@ -41,6 +48,29 @@ def passed(self) -> bool: ) +@dataclass(frozen=True) +class MarginReport: + """Teacher-forced fragility analysis: what the CER gate cannot see. + + Byte students have flat top-1 margins, so quantization noise can flip + near-tie argmaxes without moving CER on a golden set. This report + measures that directly: per-position top1−top2 margins of the torch + reference, argmax flip rate against the zip, KL divergence, and the + share of flips that land on near-tie positions (benign) versus + confident ones (dangerous).""" + + samples: int + tokens: int + flipped_tokens: int + flip_rate: float # fraction of teacher-forced positions with argmax disagreement + kld_mean: float # mean KL(reference || zip) over positions + margin_p1: float # reference top1−top2 margin quantiles, in logits + margin_p10: float + margin_p50: float + flip_low_margin_share: float # flips at margin < p10 / all flips (1.0 = benign) + precision: str = "fp32" + + def _torch_greedy_tokens(model, text: str, max_len: int) -> list[int]: import torch @@ -138,6 +168,127 @@ def run_parity( ) +def _margin_stats(ref, got): + """Per-position stats for one sequence. ``ref``/``got`` are (T, V) + float64 teacher-forced logits from the torch reference and the zip.""" + import numpy as np + + top2 = np.partition(ref, -2, axis=-1)[:, -2:] + margins = top2[:, 1] - top2[:, 0] + flips = ref.argmax(axis=-1) != got.argmax(axis=-1) + + ref_s = np.exp(ref - ref.max(axis=-1, keepdims=True)) + ref_s = ref_s / ref_s.sum(axis=-1, keepdims=True) + got_s = np.exp(got - got.max(axis=-1, keepdims=True)) + got_s = got_s / got_s.sum(axis=-1, keepdims=True) + kld = (ref_s * (np.log(ref_s + 1e-12) - np.log(got_s + 1e-12))).sum(axis=-1) + return margins, flips, kld + + +def _torch_forced_logits(model, source: str, target_ids: list[int]): + """Teacher-forced decoder logits — the exact math ``_torch_greedy_tokens`` + runs one step of, computed for every position at once.""" + import torch + + src = torch.tensor([encode_bytes(source)], dtype=torch.long) + if src.shape[1] == 1: + raise ValueError("source must be at least one byte") + dec_ids = torch.tensor([[PAD_ID] + target_ids[:-1]], dtype=torch.long) + enc = model.get_encoder()(input_ids=src)[0] + hidden = model.get_decoder()(input_ids=dec_ids, encoder_hidden_states=enc)[0] + return model.lm_head(hidden * (model.config.d_model**-0.5))[0] + + +def _onnx_forced_logits(enc_sess, dec_sess, source: str, target_ids: list[int]): + """Teacher-forced logits from the zip's decoder graph. Works for both + the plain decoder (one full-sequence call) and the KV decoder (zero + pasts + full input_ids is the same computation).""" + import numpy as np + + ids = np.array([encode_bytes(source)], dtype=np.int64) + hidden = enc_sess.run(None, {"input_ids": ids})[0] + dec_ids = np.array([[PAD_ID] + target_ids[:-1]], dtype=np.int64) + inputs = {i.name for i in dec_sess.get_inputs()} + if "encoder_hidden_states" in inputs and not any( + name.startswith("past_") for name in inputs + ): + return dec_sess.run( + None, {"input_ids": dec_ids, "encoder_hidden_states": hidden} + )[0][0] + out_names = [o.name for o in dec_sess.get_outputs()] + out = dec_sess.run( + None, + {"input_ids": dec_ids, "encoder_hidden_states": hidden, **_zero_pasts(dec_sess)}, + ) + return dict(zip(out_names, out, strict=True))["logits"][0] + + +def run_margin_analysis(model, zip_path, pairs, max_len: int = 256) -> MarginReport: + """pairs: iterable of (source_text, gold_target) — the same probe set + the CER parity gate uses. Teacher-forces both sides and measures the + argmax flip rate, reference top1−top2 margin quantiles, and KL + divergence. Complements ``run_parity``: CER measures what already + broke, margins measure how close the rest is to breaking.""" + import numpy as np + import yaml + + zip_path = Path(zip_path) + enc, dec = _sessions_from_zip(zip_path) + with zipfile.ZipFile(zip_path) as zf: + precision = yaml.safe_load(zf.read("metadata.yaml"))["precision"] + + samples = 0 + kld_sum = 0.0 + margin_chunks: list = [] + flip_chunks: list = [] + for source, target in pairs: + target_ids = encode_bytes(target)[:max_len] # trailing EOS included + if len(target_ids) < 2: + continue + samples += 1 + ref = _torch_forced_logits(model, source, target_ids) + ref = ref.detach().numpy().astype("float64") + got = np.asarray(_onnx_forced_logits(enc, dec, source, target_ids), dtype="float64") + margins, flips, kld = _margin_stats(ref, got) + margin_chunks.append(margins) + flip_chunks.append(flips) + kld_sum += float(kld.sum()) + + margins = np.concatenate(margin_chunks) + flips = np.concatenate(flip_chunks) + p1, p10, p50 = (float(np.quantile(margins, q)) for q in (0.01, 0.10, 0.50)) + n_flips = int(flips.sum()) + return MarginReport( + samples=samples, + tokens=int(margins.size), + flipped_tokens=n_flips, + flip_rate=round(n_flips / max(margins.size, 1), 6), + kld_mean=round(kld_sum / max(margins.size, 1), 8), + margin_p1=round(p1, 4), + margin_p10=round(p10, 4), + margin_p50=round(p50, 4), + flip_low_margin_share=round( + float((margins[flips] < p10).mean()) if n_flips else 0.0, 4 + ), + precision=precision, + ) + + +def write_margin_report(report: MarginReport, out_path: Path | str) -> Path: + """Emit the margin analysis as JSON next to a release zip (diagnostic + artifact; the release gate remains the CER parity block).""" + from dataclasses import asdict + + import json + + out_path = Path(out_path) + out_path.parent.mkdir(parents=True, exist_ok=True) + out_path.write_text( + json.dumps(asdict(report), indent=2) + "\n", encoding="utf-8" + ) + return out_path + + def write_parity(zip_path: Path | str, report: ParityReport) -> Path: """Write the parity block into the zip's metadata and enforce strict validation. Raises if the gate does not pass.""" diff --git a/tests/test_gpu_pkm_muon.py b/tests/test_gpu_pkm_muon.py new file mode 100644 index 0000000..ce275aa --- /dev/null +++ b/tests/test_gpu_pkm_muon.py @@ -0,0 +1,131 @@ +"""Smoke tests for ``gpu.pkm`` and ``gpu.muon`` (CPU, tiny models). + +The identity property matters most: a zero-initialized memory gate must +leave the pretrained backbone's function bit-identical at step 0 — that +is what makes injection safe on a pretrained student. +""" + +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") +transformers = pytest.importorskip("transformers") + +from gpu.muon import Muon, split_parameters # noqa: E402 +from gpu.pkm import inject_pkm, load_student_with_pkm # noqa: E402 + + +def _tiny_t5(): + from transformers import T5Config, T5ForConditionalGeneration + + config = T5Config( + vocab_size=259, d_model=32, d_ff=64, d_kv=16, + num_layers=4, num_decoder_layers=4, num_heads=2, + feed_forward_proj="relu", decoder_start_token_id=0, + ) + return T5ForConditionalGeneration(config) + + +def _gates(model): + return [b.layer[1].gate for b in model.decoder.block if hasattr(b.layer[1], "gate")] + + +def test_pkm_gate_zero_preserves_function() -> None: + model = _tiny_t5().eval() + ids = torch.randint(3, 259, (2, 7)) + labels = torch.randint(3, 259, (2, 5)) + with torch.no_grad(): + before = model(input_ids=ids, labels=labels).logits.clone() + inject_pkm(model, layer_indices=[-1, -3], n_keys=8, topk=4) + with torch.no_grad(): + after = model(input_ids=ids, labels=labels).logits + assert torch.equal(before, after) + assert len(_gates(model)) == 2 + + +def test_pkm_gradients_flow_once_gate_moves() -> None: + model = _tiny_t5().train() + inject_pkm(model, layer_indices=[-1], n_keys=8, topk=4) + with torch.no_grad(): + _gates(model)[0].fill_(1.0) + ids = torch.randint(3, 259, (2, 7)) + labels = torch.randint(3, 259, (2, 5)) + model(input_ids=ids, labels=labels).loss.backward() + memory = model.decoder.block[-1].layer[1].memory + for name in ("wq", "wo", "values", "k1", "k2"): + p = getattr(memory, name) + p = p.weight if not isinstance(p, torch.Tensor) else p + assert p.grad is not None, name + assert p.grad.abs().sum() > 0, name + + +def test_pkm_generate_smoke() -> None: + model = _tiny_t5().eval() + inject_pkm(model, layer_indices=[-2], n_keys=8, topk=4) + with torch.no_grad(): + _gates(model)[0].fill_(0.5) + ids = torch.randint(3, 259, (1, 6)) + out = model.generate(input_ids=ids, max_length=10, num_beams=1) + assert out.shape[0] == 1 + + +def test_pkm_checkpoint_roundtrip(tmp_path: Path) -> None: + cfg = {"layer_indices": [-1], "n_keys": 8, "topk": 4} + model = _tiny_t5().eval() + inject_pkm(model, **cfg) + with torch.no_grad(): + _gates(model)[0].fill_(0.5) + _gates(model)[0].add_(0.13) # non-trivial gate value must survive + ids = torch.randint(3, 259, (2, 7)) + labels = torch.randint(3, 259, (2, 5)) + with torch.no_grad(): + before = model(input_ids=ids, labels=labels).logits + out_dir = tmp_path / "best" + model.save_pretrained(str(out_dir)) + loaded = load_student_with_pkm(out_dir, cfg).eval() + with torch.no_grad(): + after = loaded(input_ids=ids, labels=labels).logits + assert torch.allclose(before, after, atol=1e-5) + + +def test_muon_split_routes_embedding_like_params() -> None: + model = _tiny_t5() + inject_pkm(model, layer_indices=[-1], n_keys=8, topk=4) + muon_params, adamw_params = split_parameters(model.named_parameters()) + assert muon_params and adamw_params + assert all(p.ndim == 2 for p in muon_params) + id_adamw = {id(p) for p in adamw_params} + names_adamw = [n for n, p in model.named_parameters() if id(p) in id_adamw] + names_muon = [n for n, p in model.named_parameters() if id(p) not in id_adamw] + assert any("embed_tokens" in n or n == "shared.weight" for n in names_adamw) + assert any("memory.values" in n for n in names_adamw) + assert any("decoder.block" in n for n in names_muon) + + +def test_muon_step_updates_params_and_saves_state() -> None: + model = _tiny_t5() + inject_pkm(model, layer_indices=[-1], n_keys=8, topk=4) + muon_params, adamw_params = split_parameters(model.named_parameters()) + opt = Muon(muon_params, lr=0.01) + opt.add_adamw_group(adamw_params, lr=1e-3) + p0 = muon_params[0].detach().clone() + q0 = adamw_params[0].detach().clone() + loss = model( + input_ids=torch.randint(3, 259, (2, 7)), + labels=torch.randint(3, 259, (2, 5)), + ).loss + loss.backward() + opt.step() + assert not torch.equal(p0, muon_params[0]) + assert not torch.equal(q0, adamw_params[0]) + state = opt.state_dict() + has_momentum = any( + "momentum_buffer" in v for s in state["state"].values() for v in [s] if isinstance(s, dict) + ) + assert has_momentum diff --git a/tests/test_imf_export.py b/tests/test_imf_export.py index cd15fa0..614c60e 100644 --- a/tests/test_imf_export.py +++ b/tests/test_imf_export.py @@ -134,3 +134,24 @@ def test_fp16_smaller_than_fp32(zips: dict[str, Path]) -> None: assert zips["fixture-1.0-fp16.zip"].stat().st_size < zips[ "fixture-1.0-fp32.zip" ].stat().st_size + + +def test_int8_keeps_head_matmul_in_fp32(zips: dict[str, Path]) -> None: + """The logits-producing MatMul (the tied head) must survive int8 + quantization unconverted — quantizing it moves argmax directly + (heb-diac-1.1: 9.34% confident-position flips; head-fp32 fixes it, + E1 in docs/EXPERIMENTS.md).""" + import onnx + + from imf.export import _head_matmul_names + + for graph_name in ("decoder.onnx", "decoder-kv.onnx"): + with zipfile.ZipFile(zips["fixture-1.0-fp32.zip"]) as zf: + heads = _head_matmul_names(onnx.load(zf.open(graph_name)).graph) + assert heads, f"head MatMul not found in fp32 {graph_name}" + with zipfile.ZipFile(zips["fixture-1.0-int8.zip"]) as zf: + nodes = {n.name: n for n in onnx.load(zf.open(graph_name)).graph.node} + for head in heads: + assert head in nodes, (graph_name, head) + # converted heads become MatMulInteger with the same name + assert nodes[head].op_type == "MatMul", (graph_name, head) diff --git a/tests/test_imf_parity.py b/tests/test_imf_parity.py index 36955a2..60c5f99 100644 --- a/tests/test_imf_parity.py +++ b/tests/test_imf_parity.py @@ -22,7 +22,16 @@ load_byte_seq2seq, make_fixture_checkpoint, ) -from imf.parity import ParityReport, run_parity, write_golden, write_parity # noqa: E402 +from imf.parity import ( # noqa: E402 + MarginReport, + ParityReport, + _margin_stats, + run_margin_analysis, + run_parity, + write_golden, + write_margin_report, + write_parity, +) from imf.validator import validate_zip # noqa: E402 METADATA = { @@ -116,3 +125,45 @@ def test_golden_jsonl_roundtrip(gated_zip: Path, tmp_path: Path) -> None: for row in rows: assert set(row) == {"input", "tokens", "output"} assert all(isinstance(t, int) for t in row["tokens"]) + + +@pytest.fixture(scope="module") +def fixture_model(tmp_path_factory: pytest.TempPathFactory): + ckpt = make_fixture_checkpoint(tmp_path_factory.mktemp("ckpt-margin") / "fixture") + return load_byte_seq2seq(ckpt) + + +def test_margin_stats_flags_flips_and_confidence() -> None: + import numpy as np + + ref = np.array([[10.0, 9.0, 0.0], [10.0, 0.0, 0.0]]) + got = np.vstack([np.array([[9.9, 10.0, 0.0]]), ref[1:]]) + margins, flips, kld = _margin_stats(ref, got) + assert flips.tolist() == [True, False] + assert margins.tolist() == [1.0, 10.0] + assert kld[0] > 0.0 + assert kld[1] == pytest.approx(0.0, abs=1e-9) + + +def test_margin_analysis_fp32_zip_is_clean(gated_zip: Path, fixture_model) -> None: + report = run_margin_analysis(fixture_model, gated_zip, PAIRS * 30, max_len=12) + assert report.precision == "fp32" + assert report.samples == 120 + assert report.flipped_tokens == 0 + assert report.flip_rate == 0.0 + assert report.kld_mean < 1e-6 + assert report.margin_p50 >= 0.0 + assert report.tokens == report.samples * 6 # "xxxxx" + EOS, teacher-forced + + +def test_margin_report_json_roundtrip(gated_zip: Path, fixture_model, tmp_path: Path) -> None: + report = run_margin_analysis(fixture_model, gated_zip, PAIRS, max_len=12) + out = write_margin_report(report, tmp_path / "fixture-1.0-margins-fp32.json") + data = json.loads(out.read_text(encoding="utf-8")) + assert set(data) == { + "samples", "tokens", "flipped_tokens", "flip_rate", "kld_mean", + "margin_p1", "margin_p10", "margin_p50", "flip_low_margin_share", + "precision", + } + assert data["samples"] == report.samples + assert data["flip_rate"] == report.flip_rate