diff --git a/TODO.sota-2026/01-r8-teacher-yallamorph.md b/TODO.sota-2026/01-r8-teacher-yallamorph.md index 2e5d812..484fbdb 100644 --- a/TODO.sota-2026/01-r8-teacher-yallamorph.md +++ b/TODO.sota-2026/01-r8-teacher-yallamorph.md @@ -1,6 +1,8 @@ # 01 — r8 teacher: run-009-yallamorph (YallaMorph/CamelMorph aux stream) -Status: SPECIFIED (2026-10-01) — not launched +Status: TRAINING IN FLIGHT (launched 2026-10-01 18:07, app ap-f1R8ChDikMKUGyBuJGO6; +26,289 steps on A100-80GB; supervisor /tmp/r9-supervisor.sh v4 relaunches on +kill-storm deaths — training is checkpoint-resumable and EVAL_DONE-idempotent) Literature basis: YallaMorph (arXiv 2609.10153, EMNLP 2026) — 663,804 controlled morphological-generation instances over 4,795 lemmas, constructed from CamelMorph MSA via CAMeL Tools. The GitHub repo ships @@ -44,16 +46,31 @@ Recipe = r7 verbatim (train_arabic_r7.py) + one new aux stream: ## Steps -1. [ ] Clone CAMeL-Lab/YallaMorph; extract xlsx samples as the +1. [x] Clone CAMeL-Lab/YallaMorph; extract xlsx samples as the validation set for our generated forms. -2. [ ] Data build: camel-tools + camel_data MSA; generate paradigm - pairs; validate forms against YallaMorph samples (match rate - reported); dedupe; cap 300k aux lines; volume put to - /datasets/yallamorph-aux/lines.txt. -3. [ ] TDD the line builder (pure function: feature dict → line pair). -4. [ ] train_arabic_r9_yallamorph.py (r7 copy + stream M + gates). -5. [ ] Launch `modal run --detach` (retry-loop supervisor per - [[modal-always-detach]]). +2. [x] Data build: camel-tools 1.5.7 + **Camel Morph MSA v1.0** + (LREC-COLING 2024, CC BY 4.0 — the resource YallaMorph was + constructed from; camel_data's calima-msa-r13 CANNOT generate + mood/command forms — use the camel_morph repo DB). Key + interface lesson: generation must be UNDERSPECIFIED (pos + + proclitic variant only); fully-specified requests silently + reject cells with unmarked features (1st person gen='u'). + 13,000 lemmas (YallaMorph sample in-DB: v 513 / n 1,862 / + adj 594 + inventory top-up, seed 42) → 4,735,166 raw pairs → + 300,000 lines (60/40 verb/nominal), 29MB, on volume + /datasets/yallamorph-aux/ (lines.txt + DONE + README). + Validation vs YallaMorph few-shot gold: **38/52 exact**; + residual divergence = proclitic-chain conventions + (hamzat-istifham أَلِـ, sin/lam stacking) + DB-version lemma + gaps — not incorrect forms. +3. [x] TDD the line builder (build_yallamorph_aux.py, 12 tests + green; rababa PR #104). +4. [x] train_arabic_r9_yallamorph.py (r7 copy + MORPH stream + gates; + rababa PR #104). Mix verified in launch logs: anchor=586,505, + news-mix=50,003, morph-used=197,169, **aux-share 25.00%**. +5. [x] Launch `modal run --detach` (+ supervisor; one double-launch + incident from `modal app list` name truncation — grep prefix + "rababa-ara", dupes stopped, volume verified clean). 6. [ ] ID gate: windowed zero-skip SadeedDiac-25 full 1,200-para DER ≤ 2.389 (r7 2.2864 + 0.1 tolerance). 7. [ ] OOD gate: eval_wikinews_multiref improves over 17.3794/11.8273. diff --git a/TODO.sota-2026/05-ride-sft-residual-extrapolation.md b/TODO.sota-2026/05-ride-sft-residual-extrapolation.md index d72805c..746343b 100644 --- a/TODO.sota-2026/05-ride-sft-residual-extrapolation.md +++ b/TODO.sota-2026/05-ride-sft-residual-extrapolation.md @@ -1,6 +1,6 @@ # 05 — RIDE-style SFT-residual extrapolation: probe-first arm -Status: SPECIFIED (2026-10-01) — probe only; training arm gated on probe +Status: PROBE PASSED (2026-10-01) — training arm spec'd below, launch pending owner call Literature basis: RIDE (arXiv 2609.36484) — extrapolate the teacher-over-base residual directly in representation space: student hidden states regressed toward @@ -41,13 +41,49 @@ residual extrapolation of an SFT delta. ## Steps -1. [ ] TDD pure computation: `residual_directions(h_base, h_teacher)` - and cosine sim on synthetic tensors (tests first, watch fail). -2. [ ] Modal probe script (two models × 200 units × 2 domains; - A100 minutes, not hours). -3. [ ] Run probe; write verdict + per-layer cosine table here. -4. [ ] Gate decision: close, or spec the training arm separately. +1. [x] TDD pure computation: `residual_directions(h_base, h_teacher)` + and cosine sim on synthetic tensors (7/7 green, + rababa/test_ride_probe_math.py). +2. [x] Modal probe script (rababa/probe_ride_direction.py; A10G, + ~8 min, app ap-LePMpScA3RA8EN6TLDP2vN). +3. [x] Run probe; verdict + per-layer cosine table below. +4. [x] Gate decision: **TRANSFER** — training arm spec'd below. -## Result +## Result (measured 2026-10-01) -(to be written only from measured numbers) +Per-layer cos(d_classical, d_news), 200 units per domain, mean-pooled +encoder hidden states, run-006-morph vs run-007-news: + +``` +L00 +0.9414 L01 +0.9513 L02 +0.9199 L03 +0.9002 L04 +0.8765 +L05 +0.8696 L06 +0.8593 L07 +0.7734 L08 +0.5931 L09 +0.4521 +L10 +0.2145 L11 -0.0015 L12 -0.1086 L13 -0.0842 L14 +0.0167 +L15 +0.1624 L16 +0.1971 L17 +0.2387 L18 +0.2775 +max +0.9513 >= 0.5 -> TRANSFER +``` + +Reading: the r7-over-r6 SFT residual is strongly domain-general in +early/mid encoder layers (L0–L8: 0.59–0.95) and idiosyncratic deep +(L11+ ≈ noise). The news-mix fine-tune moved surface/orthographic +processing in a direction that transfers to classical text — the +RIDE displacement premise holds where representations are shared. +Artefact: rababa-checkpoints:/ride_probe_r6_r7.json. + +## Training-arm spec (gated on this probe; launch = owner decision) + +The closure rule (TODO 03) permits this arm: it is a mechanism novel +to the ledger (representation-space displacement; all prior student +levers were loss/data/optimizer-side) and now has a measured transfer +premise. + +- Infra: feature-regression aux loss in modal_distill — teacher/base + hidden states must be cached per layer subset. Restrict to L0–L8 + (probe: only these transfer; deep-layer displacement would inject + domain idiosyncrasy). +- Target: h_t' = h_teacher + λ(h_teacher − h_base), λ ∈ {0.5, 1.0}. +- Loss: L = CE(labels) + β·MSE(h_student[L0-8], h_t'[L0-8]), + β tuned so MSE term ≈ 10% of total at start (pre-registration). +- Single-variable off the 2.1 recipe (run-007 data, teacher labels, + Muon, seed 42); adopt gate ≥ 0.3pp DER improvement (E4-style bar). +- Est. build: teacher/base hidden-state dump (one-off Modal job, + ~1h A100) + trainer loss path + spec; run cost ≈ one 2.1-recipe arm. diff --git a/src/gpu/distill_specs.yaml b/src/gpu/distill_specs.yaml index 76ac5f0..bf5c3f8 100644 --- a/src/gpu/distill_specs.yaml +++ b/src/gpu/distill_specs.yaml @@ -657,3 +657,45 @@ heb-diac-small-s46-layerdrop: val: hebrew-v4/val.jsonl out: rababa_hebrew_distill_small/run-003-s46-layerdrop note: Hebrew lite rung; cross-lingual depth-cut replication +ara-diac-small-2-1-ride: + # TODO.sota-2026/05 arm: RIDE displacement. Probe passed 2026-10-01 + # (r7-r6 residual direction domain-general in encoder L0-L8, + # max cos 0.9513). Single-variable off the 2.1 recipe (2-6ep + # verbatim): encoder-hidden regression toward ridge-projected + # h_t + lam*(h_t - h_b), frozen r7 teacher + r6 base, layers 0-8 + # only (probe: deeper layers are domain-idiosyncratic). + # Gate: adopt iff >= 0.3pp better than 4.5701 (E4-style bar). + teacher: rababa_arabic_byt5/run-007-news/best + teacher_volume: rababa + out_volume: rababa + student_init: google/byt5-small + 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-016-ride + labels_file: teacher_labels_r7.jsonl + labels_complete: 'true' + mode: sequence + optimizer: muon + ride_aux: + base: rababa_arabic_byt5/run-006-morph/best + layers: + - 0 + - 1 + - 2 + - 3 + - 4 + - 5 + - 6 + - 7 + - 8 + lam: '1.0' + beta_frac: '0.1' + sub: '4' + fit_batches: '16' + note: 'TODO.sota-2026/05; gate vs 4.5701 (2.1) - adopt >= 0.3pp better' diff --git a/src/gpu/modal_distill.py b/src/gpu/modal_distill.py index 543e977..dfc90d6 100644 --- a/src/gpu/modal_distill.py +++ b/src/gpu/modal_distill.py @@ -789,6 +789,46 @@ def distill_sequence(spec_id: str, epochs: int = 3) -> dict: f"sub={gkd_cfg.get('sample_sub', 2)} cap={gkd_cfg.get('sample_cap', 1024)}", flush=True, ) + ride_cfg = spec.get("ride_aux") + ride = None + if ride_cfg: + # RIDE displacement arm (TODO.sota-2026/05; probe passed + # 2026-10-01: r7-r6 residual direction domain-general L0-L8). + # Regress student encoder hiddens toward ridge-projected + # h_t + lam*(h_t - h_b) from frozen teacher (r7) and base (r6). + _ensure_src_path() + from gpu.ride import displaced_targets, fit_ridge, masked_mse + + base_path = str( + Path(VOLUME_MOUNTS[ride_cfg.get("base_volume", teacher_vol)]) / ride_cfg["base"] + ) + base_tok = AutoTokenizer.from_pretrained(base_path) + base = ( + AutoModelForSeq2SeqLM.from_pretrained(base_path) + .to("cuda", dtype=torch.float16) + .eval() + ) + for p in base.parameters(): + p.requires_grad_(False) + sample = "الْعَرَبِيَّةُ byte-parity 123" + ids_s = student_tok(sample).input_ids + for name, tok in (("teacher", teacher_tok), ("base", base_tok)): + if tok(sample).input_ids != ids_s: + raise RuntimeError(f"ride_aux: {name} tokenizer differs from student byte table") + ride = { + "base": base, + "layers": [int(x) for x in ride_cfg.get("layers", list(range(9)))], + "lam": float(ride_cfg.get("lam", 1.0)), + "beta_frac": float(ride_cfg.get("beta_frac", 0.1)), + "sub": int(ride_cfg.get("sub", 4)), + "beta": None, + "W": None, + } + print( + f"[{spec_id}] ride_aux: base={ride_cfg['base']} " + f"layers={ride['layers']} lam={ride['lam']} sub={ride['sub']}", + flush=True, + ) student.train() class Pairs(Dataset): @@ -1205,12 +1245,96 @@ def _usable(ck: Path) -> bool: scheduler.step() print(f"[{spec_id}] resume training from step-{step}", flush=True) + if ride is not None: + # displacement needs the teacher resident on cuda for both the + # ridge fit and every training step (labeling offloads it) + teacher.to("cuda") + teacher.eval() + # Ridge map per selected layer: teacher hidden space -> the + # student's INITIAL hidden space (linearity preserves the + # displacement through the projection). Fit on the first + # fit_batches batches, no grad, before any training step. + from gpu.ride import fit_ridge as _fit_ridge + + ride_state_path = ckpts[-1] / "ride.pt" if ckpts else None + if ride_state_path is not None and ride_state_path.exists(): + st = torch.load(ride_state_path, map_location="cpu", weights_only=True) + ride["W"] = {int(L): w.to("cuda") for L, w in st["W"].items()} + ride["beta"] = st["beta"] + print(f"[{spec_id}] ride resumed (beta={ride['beta']:.3e})", flush=True) + else: + it = iter(train_loader) + Ht = {L: [] for L in ride["layers"]} + Hs = {L: [] for L in ride["layers"]} + student.eval() + with torch.no_grad(): + for _ in range(int(ride_cfg.get("fit_batches", 16))): + try: + b_ids, b_am, _lbl = next(it) + except StopIteration: + break + b_ids, b_am = b_ids.to("cuda"), b_am.to("cuda") + t_hs = teacher.encoder( + input_ids=b_ids, attention_mask=b_am, output_hidden_states=True + ).hidden_states + s_hs = student.encoder( + input_ids=b_ids, attention_mask=b_am, output_hidden_states=True + ).hidden_states + # mask-flatten to real positions: batches have + # variable T, so keep (N_real, D) rows only + keep = b_am.bool() + for L in ride["layers"]: + Ht[L].append(t_hs[L][keep].float().cpu()) + Hs[L].append(s_hs[L][keep].float().cpu()) + student.train() + alpha = float(ride_cfg.get("ridge_alpha", 1.0)) + ride["W"] = { + L: _fit_ridge(torch.cat(Ht[L]), torch.cat(Hs[L]), alpha=alpha).to("cuda") + for L in ride["layers"] + } + print(f"[{spec_id}] ride ridge fit on {len(Ht[ride['layers'][0]])} batches", flush=True) + for _ in range(epochs): for ids, am, labels in train_loader: if step >= total_steps: break ids, am, labels = ids.to("cuda"), am.to("cuda"), labels.to("cuda") - if mtp_head is not None: + if ride is not None: + s_out = student( + input_ids=ids, attention_mask=am, labels=labels, + output_hidden_states=True, + ) + loss = s_out.loss + sub = ride["sub"] + with torch.no_grad(): + t_hs = teacher.encoder( + input_ids=ids[:sub], attention_mask=am[:sub], + output_hidden_states=True, + ).hidden_states + b_hs = ride["base"].encoder( + input_ids=ids[:sub], attention_mask=am[:sub], + output_hidden_states=True, + ).hidden_states + mse_total = None + for L in ride["layers"]: + tgt = displaced_targets( + t_hs[L].float(), b_hs[L].float(), ride["lam"] + ) @ ride["W"][L].T + m = masked_mse( + s_out.encoder_hidden_states[L][:sub].float(), tgt, am[:sub] + ) + mse_total = m if mse_total is None else mse_total + m + if ride["beta"] is None: + ride["beta"] = ( + ride["beta_frac"] * loss.detach() / (mse_total.detach() + 1e-12) + ).item() + print( + f"[{spec_id}] ride beta calibrated: {ride['beta']:.3e} " + f"(mse0={float(mse_total):.3f})", + flush=True, + ) + loss = loss + ride["beta"] * mse_total + elif 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, @@ -1270,6 +1394,12 @@ def _usable(ck: Path) -> bool: (ck / "labels.sha").write_text(labels_digest) torch.save(student.state_dict(), ck / "student.pt") torch.save(optimizer.state_dict(), ck / "optim.pt") + if ride is not None and ride["beta"] is not None: + torch.save( + {"W": {L: w.cpu() for L, w in ride["W"].items()}, + "beta": ride["beta"]}, + ck / "ride.pt", + ) CHECKPOINTS.commit() SECRYST_CHECKPOINTS.commit() PERSIAN_CHECKPOINTS.commit() diff --git a/src/gpu/ride.py b/src/gpu/ride.py new file mode 100644 index 0000000..5eef8c6 --- /dev/null +++ b/src/gpu/ride.py @@ -0,0 +1,47 @@ +"""RIDE displacement-arm math (TODO.sota-2026/05). + +The probe (rababa/probe_ride_direction.py, 2026-10-01) measured the +r7-over-r6 SFT residual direction as domain-general in encoder layers +0-8 (cos 0.59-0.95 across classical/news text) and idiosyncratic +deeper. These helpers implement the training arm: regress the +student's encoder hidden states toward ridge-projected, extrapolated +teacher targets h_t' = h_teacher + lam*(h_teacher - h_base). + +The ridge map is fit at arm start (teacher hiddens -> the student's +initial hiddens on the same inputs); by linearity the displacement +survives the projection: W(h_t + lam*(h_t - h_b)) = W h_t + +lam*(W h_t - W h_b). +""" + +from __future__ import annotations + +import torch + + +def displaced_targets(h_teacher: torch.Tensor, h_base: torch.Tensor, lam: float) -> torch.Tensor: + return h_teacher + lam * (h_teacher - h_base) + + +def masked_mse(pred: torch.Tensor, target: torch.Tensor, attention_mask: torch.Tensor) -> torch.Tensor: + """Mean squared error over kept positions and all dims. + + pred/target: (B, T, D); attention_mask: (B, T), 1 = kept. + Returns 0.0 when nothing is kept. + """ + mask = attention_mask.to(pred.dtype).unsqueeze(-1) + diff2 = (pred - target).pow(2) * mask + n = mask.sum() * pred.shape[-1] + if n == 0: + return pred.new_tensor(0.0) + return diff2.sum() / n + + +def fit_ridge(H_teacher: torch.Tensor, H_student: torch.Tensor, alpha: float) -> torch.Tensor: + """Closed-form ridge W minimizing ||H_teacher @ W.T - H_student||^2 + + alpha*||W||^2. Returns W of shape (d_student, d_teacher).""" + Ht = H_teacher.double() + Hs = H_student.double() + d_t = Ht.shape[-1] + A = Ht.T @ Ht + alpha * torch.eye(d_t, dtype=Ht.dtype, device=Ht.device) + W = torch.linalg.solve(A, Ht.T @ Hs).T + return W.to(H_teacher.dtype) diff --git a/tests/test_ride.py b/tests/test_ride.py new file mode 100644 index 0000000..14a0c0b --- /dev/null +++ b/tests/test_ride.py @@ -0,0 +1,93 @@ +"""Tests for the RIDE displacement-arm math (TODO.sota-2026/05).""" + +import unittest +from pathlib import Path +import sys + +import torch + +sys.path.insert(0, str(Path(__file__).resolve().parent.parent / "src")) + +from gpu.ride import displaced_targets, fit_ridge, masked_mse + + +class TestDisplacedTargets(unittest.TestCase): + def test_lambda_zero_is_teacher(self): + h_t = torch.tensor([[1.0, 2.0], [0.0, -1.0]]) + h_b = torch.tensor([[3.0, 6.0], [0.0, -3.0]]) + self.assertTrue(torch.equal(displaced_targets(h_t, h_b, lam=0.0), h_t)) + + def test_lambda_one_extrapolates(self): + h_t = torch.tensor([[2.0, 4.0]]) + h_b = torch.tensor([[0.0, 0.0]]) + self.assertTrue(torch.equal(displaced_targets(h_t, h_b, lam=1.0), torch.tensor([[4.0, 8.0]]))) + + def test_fractional_lambda(self): + h_t = torch.tensor([[3.0]]) + h_b = torch.tensor([[1.0]]) + self.assertTrue(torch.equal(displaced_targets(h_t, h_b, lam=0.5), torch.tensor([[4.0]]))) + + +class TestMaskedMse(unittest.TestCase): + def test_exact_match_is_zero(self): + h = torch.tensor([[1.0, 2.0]]) + self.assertEqual(masked_mse(h, h.clone(), torch.tensor([1])), 0.0) + + def test_masked_out_positions_ignored(self): + pred = torch.tensor([[[0.0], [9.0]]]) + tgt = torch.tensor([[[1.0], [0.0]]]) + am = torch.tensor([[1, 0]]) # second position padded/masked + self.assertAlmostEqual( + masked_mse(pred, tgt, am).item(), 1.0) + + def test_averages_over_kept_positions_and_dims(self): + pred = torch.tensor([[0.0, 2.0], [4.0, 4.0]]) + tgt = torch.tensor([[2.0, 2.0], [2.0, 2.0]]) + # diffs: 2,0,2,2 -> mse 12/4 = 3 + self.assertAlmostEqual( + masked_mse(pred, tgt, torch.tensor([1, 1])).item(), 3.0) + + def test_zero_kept_is_zero(self): + self.assertEqual( + masked_mse(torch.tensor([[1.0]]), torch.tensor([[0.0]]), torch.tensor([0])).item(), 0.0) + + def test_batch_shape_with_padding(self): + # (B=2, T=2, D=1); row 1 fully padded + pred = torch.tensor([[[0.0], [2.0]], [[9.0], [9.0]]]) + tgt = torch.tensor([[[2.0], [2.0]], [[0.0], [0.0]]]) + am = torch.tensor([[1, 1], [0, 0]]) + # kept row 0: diffs 4, 0 -> sum 4 / (2 kept x D=1) = 2 + self.assertAlmostEqual(masked_mse(pred, tgt, am).item(), 2.0) + + +class TestFitRidge(unittest.TestCase): + def test_identity_when_spaces_align(self): + H = torch.tensor([[1.0, 0.0], [0.0, 1.0], [1.0, 1.0]]) + W = fit_ridge(H, H, alpha=1e-8) + got = H @ W.T + self.assertTrue(torch.allclose(got, H, atol=1e-4)) + + def test_scales_output(self): + H = torch.tensor([[2.0], [4.0]]) + S = torch.tensor([[4.0], [8.0]]) # student = 2x teacher + W = fit_ridge(H, S, alpha=1e-8) + self.assertAlmostEqual(W[0][0].item(), 2.0, places=4) + + def test_alpha_regularizes(self): + H = torch.tensor([[1.0], [0.0]]) + W0 = fit_ridge(H, H, alpha=0.0) + W1 = fit_ridge(H, H, alpha=10.0) + self.assertLess(abs(W1[0][0].item()), abs(W0[0][0].item()) + 1e-9) + + def test_maps_width(self): + # teacher d=1 -> student d=2 + H = torch.tensor([[1.0], [2.0], [3.0]]) + S = torch.tensor([[1.0, -1.0], [2.0, -2.0], [3.0, -3.0]]) + W = fit_ridge(H, S, alpha=1e-8) # (2, 1) + self.assertEqual(W.shape, (2, 1)) + got = H @ W.T + self.assertTrue(torch.allclose(got, S, atol=1e-4)) + + +if __name__ == "__main__": + unittest.main()