From 464ad763f6a3eeba6893aebe87c83043861e91fe Mon Sep 17 00:00:00 2001 From: Ronald Tse Date: Mon, 31 Aug 2026 20:34:24 +0200 Subject: [PATCH] =?UTF-8?q?feat(e5):=20MTP-aux=20distillation=20rung=20?= =?UTF-8?q?=E2=80=94=20registered=20+=20implemented?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit EXPERIMENTS.md E5: control = run-006-r7-muon verbatim, single delta = a 3-step multi-token-prediction head (~1.1M params, 0.4%) on the student decoder, aux CE at beta 0.15, head discarded at inference so the shipped artifact stays vanilla and size-identical. Source: Hy4's native MTP mapped as a TRAINING auxiliary (TODO 07-hy4) — serving-side speculation stays parked, decode measured non-binding. gpu/mtp.py (MTPHead + build/mtp_named), distill_sequence wiring (head before optimizer so Muon sees its matrices; mtp_head.pt saved beside student.pt for resume; provenance-only copy under best/), spec ara-diac-small-2-mtp -> run-007-r7-muon-mtp. 3 unit tests (shift/mask semantics, uniform-logit baseline, naming). Gate (pre-agreed): adopt at <=4.5218 full-set windowed DER; honest report in [4.5218, 4.8218). Prediction: 4.5-4.75. --- docs/EXPERIMENTS.md | 25 +++++++++++++++++++ src/gpu/distill_specs.yaml | 23 +++++++++++++++++ src/gpu/modal_distill.py | 42 +++++++++++++++++++++++++++++-- src/gpu/mtp.py | 51 ++++++++++++++++++++++++++++++++++++++ tests/test_gpu_mtp.py | 51 ++++++++++++++++++++++++++++++++++++++ 5 files changed, 190 insertions(+), 2 deletions(-) create mode 100644 src/gpu/mtp.py create mode 100644 tests/test_gpu_mtp.py diff --git a/docs/EXPERIMENTS.md b/docs/EXPERIMENTS.md index e084b30..83c3a66 100644 --- a/docs/EXPERIMENTS.md +++ b/docs/EXPERIMENTS.md @@ -195,6 +195,31 @@ All rows passed the CER parity gate at release. Readings: teacher+0.5pp budget remains the disclosed north star. - **Prediction (registered):** 4.3–5.0, by E3's 4.829 on weaker labels. +## E5 — MTP-aux distillation rung (run-007-r7-muon-mtp) + +- **Status:** REGISTERED 2026-09-01, launching. +- **Source:** Tencent Hy4-preview carries a native multi-token + prediction layer; mapped to our stack as a TRAINING auxiliary (TODO + 07-hy4) — per-position multi-step heads densify supervision on the + decode path; serving-side speculation stays parked (decode measured + non-binding at our sizes). +- **Hypothesis:** the student's residual errors concentrate where the + single-step objective leaves the byte decision underconstrained; + forcing each decoder position to also predict t+1..t+3 regularizes + the hidden state toward the local sequence structure that + diacritization output exhibits (letter + haraqat pattern). +- **Design:** control = run-006-r7-muon verbatim (same teacher labels, + corpus, limits, schedule, Muon groups); single delta = MTPHead + attached to the student decoder (3 steps, byte-vocab 259, ~1.1M + params ≈ 0.4%), auxiliary CE weighted β=0.15, head DISCARDED at + inference (zero serving cost/size delta in the shipped artifact). +- **Pre-agreed gate:** adopt if full-set windowed DER ≤ 4.5218 + (≥0.3pp over 4.8218, the E3-style bar); report honestly in + [4.5218, 4.8218); investigate if worse. +- **Prediction (registered):** 4.5–4.75 — denser supervision helps + the tail, but the factorial attributes most of the remaining gap to + domain coverage, so the effect should be second-order. + ## Parked - **Speculative decoding** (LongCat converts sparsity→speed): revisit diff --git a/src/gpu/distill_specs.yaml b/src/gpu/distill_specs.yaml index 68ca28a..25171cb 100644 --- a/src/gpu/distill_specs.yaml +++ b/src/gpu/distill_specs.yaml @@ -116,6 +116,29 @@ ara-diac-small-muon: labels_complete: 'true' mode: sequence note: vanilla ByT5-small + Muon (factorial cell 4) +ara-diac-small-2-mtp: + # E5 (EXPERIMENTS.md): control = ara-diac-small-2 verbatim; single + # delta = MTP-aux head (3 steps, beta 0.15), discarded at inference + teacher: rababa_arabic_byt5/run-007-news/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-007-r7-muon-mtp + labels_file: teacher_labels_r7.jsonl + mtp_aux: + steps: 3 + beta: 0.15 + mode: sequence + ara-diac-small-2: teacher: rababa_arabic_byt5/run-007-news/best teacher_volume: rababa diff --git a/src/gpu/modal_distill.py b/src/gpu/modal_distill.py index 99f2320..ac5a14b 100644 --- a/src/gpu/modal_distill.py +++ b/src/gpu/modal_distill.py @@ -429,6 +429,8 @@ def val_loss() -> float: ck.mkdir(exist_ok=True) torch.save(student.state_dict(), ck / "student.pt") torch.save(optimizer.state_dict(), ck / "optim.pt") + if mtp_head is not None: + torch.save(mtp_head.state_dict(), ck / "mtp_head.pt") CHECKPOINTS.commit() vl = val_loss() @@ -742,6 +744,14 @@ def distill_sequence(spec_id: str, epochs: int = 3) -> dict: from gpu.pkm import inject_pkm inject_pkm(student, **spec["pkm"]) + mtp_head = None + if spec.get("mtp_aux"): + _ensure_src_path() + from gpu.mtp import build_mtp + + mtp_head = build_mtp(student, **spec["mtp_aux"]) + n = sum(q.numel() for q in mtp_head.parameters()) / 1e6 + print(f"[{spec_id}] mtp_aux head: {n:.2f}M params", flush=True) student.train() class Pairs(Dataset): @@ -1043,7 +1053,12 @@ def __getitem__(self, i): _ensure_src_path() from gpu.muon import Muon, split_parameters - muon_params, adamw_params = split_parameters(student.named_parameters()) + named = list(student.named_parameters()) + if mtp_head is not None: + from gpu.mtp import mtp_named + + named += list(mtp_named(mtp_head)) + muon_params, adamw_params = split_parameters(named) optimizer = Muon( muon_params, lr=float(spec.get("muon_lr", 0.01)), momentum=0.95, weight_decay=0.01, @@ -1083,6 +1098,16 @@ def _usable(ck: Path) -> bool: optimizer.load_state_dict( torch.load(ckpts[-1] / "optim.pt", map_location="cpu", weights_only=True) ) + if mtp_head is not None and (ckpts[-1] / "mtp_head.pt").exists(): + mtp_head.load_state_dict( + torch.load(ckpts[-1] / "mtp_head.pt", map_location="cpu", weights_only=True) + ) + elif mtp_head is not None: + print( + f"[{spec_id}] WARNING: no mtp_head.pt at resume — " + "fresh head, aux dynamics reset", + flush=True, + ) step = int(ckpts[-1].name.split("-")[1]) for _ in range(step): scheduler.step() @@ -1093,7 +1118,18 @@ def _usable(ck: Path) -> bool: if step >= total_steps: break ids, am, labels = ids.to("cuda"), am.to("cuda"), labels.to("cuda") - loss = student(input_ids=ids, attention_mask=am, labels=labels).loss + if mtp_head is not None: + beta = float(spec["mtp_aux"].get("beta", 0.15)) + s_out = student( + input_ids=ids, attention_mask=am, labels=labels, + output_hidden_states=True, + ) + aux = mtp_head.aux_loss( + s_out.decoder_hidden_states[-1], labels + ) + loss = s_out.loss + beta * aux + else: + loss = student(input_ids=ids, attention_mask=am, labels=labels).loss loss.backward() torch.nn.utils.clip_grad_norm_(student.parameters(), 1.0) optimizer.step() @@ -1119,6 +1155,8 @@ def _usable(ck: Path) -> bool: best.mkdir(exist_ok=True) student.save_pretrained(str(best)) student_tok.save_pretrained(str(best)) + if mtp_head is not None: # provenance only; never in the shipped artifact + torch.save(mtp_head.state_dict(), best / "mtp_head.pt") CHECKPOINTS.commit() SECRYST_CHECKPOINTS.commit() PERSIAN_CHECKPOINTS.commit() diff --git a/src/gpu/mtp.py b/src/gpu/mtp.py new file mode 100644 index 0000000..53ae66b --- /dev/null +++ b/src/gpu/mtp.py @@ -0,0 +1,51 @@ +"""Multi-token-prediction auxiliary head (E5, TODO 07-hy4). + +Per-position multi-step prediction as a TRAINING auxiliary: each +decoder position's hidden state additionally predicts the target +tokens at t+1..t+k, densifying supervision on the decode path. The +head is discarded at inference — the shipped student stays vanilla +and size-identical; only the training run carries it (saved as +mtp_head.pt beside student.pt for resume, never exported). +""" + +from __future__ import annotations + +import torch +from torch import nn +import torch.nn.functional as F + + +class MTPHead(nn.Module): + def __init__(self, d_model: int, vocab: int, steps: int = 3): + super().__init__() + self.steps = steps + self.heads = nn.ModuleList( + nn.Linear(d_model, vocab) for _ in range(steps) + ) + + def aux_loss(self, hidden: torch.Tensor, labels: torch.Tensor) -> torch.Tensor: + """hidden [B, T, d] is the decoder's final hidden state whose + position t already predicts labels[t] through the main head; + step-k heads predict labels[t+k] from the same position.""" + total = 0.0 + n = 0 + for k, head in enumerate(self.heads, start=1): + tgt = labels[:, k:] + m = tgt != -100 + if m.any(): + total = total + F.cross_entropy( + head(hidden[:, :-k])[m].float(), tgt[m] + ) + n += 1 + return total / max(n, 1) + + +def build_mtp(student, steps: int = 3) -> MTPHead: + return MTPHead( + student.config.d_model, student.config.vocab_size, steps=steps + ).to(student.device) + + +def mtp_named(head: MTPHead): + for name, p in head.named_parameters(): + yield f"mtp.{name}", p diff --git a/tests/test_gpu_mtp.py b/tests/test_gpu_mtp.py new file mode 100644 index 0000000..acb15e2 --- /dev/null +++ b/tests/test_gpu_mtp.py @@ -0,0 +1,51 @@ +import torch + +from gpu.mtp import MTPHead, mtp_named + + +def _labels(): + t = torch.full((2, 6), -100) + t[0, :5] = torch.tensor([10, 11, 12, 13, 14]) + t[1, :3] = torch.tensor([4, 5, 6]) + return t + + +def test_mtp_head_shapes_and_shift(): + torch.manual_seed(0) + head = MTPHead(d_model=8, vocab=16, steps=3) + hidden = torch.randn(2, 6, 8) + loss = head.aux_loss(hidden, _labels()) + assert loss.ndim == 0 and float(loss) > 0 + + +def test_mtp_shift_targets_not_inputs(): + torch.manual_seed(0) + head = MTPHead(d_model=8, vocab=16, steps=1) + head.heads[0].weight.data.zero_() + head.heads[0].bias.data.zero_() + + labels = _labels() + hidden = torch.randn(2, 6, 8) + # uniform-logit head: CE is log(vocab) wherever any target is valid + loss = head.aux_loss(hidden, labels) + import math + + assert math.isclose(float(loss), math.log(16), rel_tol=1e-4) + + # a step-1 head must see labels[:, 1:] as targets: with all -100 + # beyond position 0 there is nothing to predict + only_first = torch.full((1, 4), -100) + only_first[0, 0] = 7.0 + assert float(head.aux_loss(torch.randn(1, 4, 8), only_first)) == 0.0 + + +def test_mtp_learnable_and_named(): + head = MTPHead(d_model=8, vocab=16, steps=2) + named = dict(mtp_named(head)) + assert set(named) == { + "mtp.heads.0.weight", "mtp.heads.0.bias", + "mtp.heads.1.weight", "mtp.heads.1.bias", + } + out = head.aux_loss(torch.randn(2, 6, 8), _labels()) + out.backward() + assert head.heads[0].weight.grad is not None