feat(gkd): on-policy distillation wiring + spec (armed, launch gated)

gpu/gkd.py: pure-torch helpers (anneal schedule, token logprobs,
reverse-KL with student-only gradients; 7 unit tests). distill_sequence
wiring: teacher kept resident for scoring, student samples every
sample_every steps (sub-batch, capped), reverse-KL at the annealed
weight — control-verbatim otherwise. Spec ara-diac-small-2-gkd per the
EXPERIMENTS registration. Fixes a pre-existing B905 in
scripts/figures/frontier.py (main-red).

Author: Ronald Tse

scripts/figures/frontier.py

@@ -62,7 +62,7 @@ def main() -> None:
     pts = [(r[1], r[2]) for r in RUNS if r[2] is not None]
     labels = [r[0] for r in RUNS if r[2] is not None]
     ax.plot([p[0] for p in pts], [p[1] for p in pts], "o-")
-    for (x, y), lab in zip(pts, labels):
+    for (x, y), lab in zip(pts, labels, strict=True):
         ax.annotate(lab, (x, y), fontsize=7, xytext=(4, 4),
                     textcoords="offset points")
     ax.set_xscale("log")

src/gpu/distill_specs.yaml

@@ -180,6 +180,33 @@ ara-diac-small-2:
   labels_file: teacher_labels_r7.jsonl
   mode: sequence
   note: 'r7 teacher + Muon; E4 gate: beat the shipped 8.259 by >= 2pp'
+ara-diac-small-2-gkd:
+  # GKD rung (EXPERIMENTS.md "GKD — on-policy distillation rung"):
+  # control run-006 verbatim; delta = reverse-KL on student-sampled
+  # sequences scored by the frozen r7 teacher, ratio 0.3, annealed to
+  # zero over the final third. Gate <= 4.5218; prediction 4.30-4.65.
+  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-011-r7-muon-gkd
+  labels_file: teacher_labels_r7.jsonl
+  mode: sequence
+  gkd:
+    ratio: 0.3
+    temperature: 1.0
+    sample_every: 4
+    sample_sub: 2
+    sample_cap: 1024
 ara-diac-small-layerdrop:
   teacher: rababa_arabic_byt5/run-006-morph/best
   teacher_volume: rababa

src/gpu/gkd.py

@@ -0,0 +1,42 @@
+"""On-policy GKD helpers (EXPERIMENTS.md "GKD — on-policy distillation
+rung"). Pure torch, no modal import — unit-testable outside CI's GPU
+image. The training-side wiring lives in gpu.modal_distill."""
+
+from __future__ import annotations
+
+import torch
+import torch.nn.functional as F
+
+
+def gkd_weight(step: int, total_steps: int, ratio: float,
+               anneal_from: float = 2 / 3) -> float:
+    """Full ratio until `anneal_from` of training, then linear to zero
+    at total_steps (the registered anneal over the final third)."""
+    if total_steps <= 0 or step >= total_steps:
+        return 0.0
+    frac = step / total_steps
+    if frac < anneal_from:
+        return ratio
+    return ratio * max(0.0, 1.0 - (frac - anneal_from) / (1.0 - anneal_from))
+
+
+def token_logprobs(logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
+    """Per-token log-probs of `targets` under `logits` (teacher-forced
+    shift): logits [B, L, V], targets [B, L] -> [B, L-1]."""
+    lp = F.log_softmax(logits[:, :-1].float(), dim=-1)
+    return lp.gather(-1, targets[:, 1:].unsqueeze(-1)).squeeze(-1)
+
+
+def reverse_kl(student_logits: torch.Tensor, teacher_logits: torch.Tensor,
+               targets: torch.Tensor, mask: torch.Tensor | None = None) -> torch.Tensor:
+    """KL(student || teacher) on the student's own sampled tokens:
+    mean over (masked) positions of logp_s - logp_t. Gradient flows
+    through the student term only — call under no_grad for teacher."""
+    s_lp = token_logprobs(student_logits, targets)
+    with torch.no_grad():
+        t_lp = token_logprobs(teacher_logits, targets)
+    diff = s_lp - t_lp
+    if mask is not None:
+        m = mask[:, 1:].to(diff.dtype)
+        return (diff * m).sum() / m.sum().clamp(min=1.0)
+    return diff.mean()

src/gpu/modal_distill.py

@@ -751,6 +751,15 @@ def distill_sequence(spec_id: str, epochs: int = 3) -> dict:
         mtp_head = build_mtp(student, steps=int(mtp_cfg.get("steps", 3)))
         n = sum(q.numel() for q in mtp_head.parameters()) / 1e6
         print(f"[{spec_id}] mtp_aux head: {n:.2f}M params", flush=True)
+    gkd_cfg = spec.get("gkd")
+    if gkd_cfg:
+        print(
+            f"[{spec_id}] gkd: ratio={gkd_cfg.get('ratio', 0.3)} "
+            f"temp={gkd_cfg.get('temperature', 1.0)} "
+            f"every={gkd_cfg.get('sample_every', 4)} "
+            f"sub={gkd_cfg.get('sample_sub', 2)} cap={gkd_cfg.get('sample_cap', 1024)}",
+            flush=True,
+        )
     student.train()
 
     class Pairs(Dataset):
@@ -1033,6 +1042,12 @@ def accept_label(src: str, label: str) -> None:
     student.gradient_checkpointing_enable()
     if mtp_head is not None:  # built while student was still on cpu
         mtp_head.to("cuda")
+    if gkd_cfg is not None:  # on-policy scoring needs the teacher resident
+        teacher.to("cuda")
+        teacher.eval()
+        _ensure_src_path()
+        from gpu.gkd import gkd_weight, reverse_kl
+        print(f"[{spec_id}] gkd: teacher resident on cuda for scoring", flush=True)
 
     class TeacherPairs(Dataset):
         def __len__(self):
@@ -1131,6 +1146,32 @@ def _usable(ck: Path) -> bool:
                 loss = s_out.loss + beta * aux
             else:
                 loss = student(input_ids=ids, attention_mask=am, labels=labels).loss
+            if (
+                gkd_cfg is not None
+                and step % int(gkd_cfg.get("sample_every", 4)) == 0
+            ):
+                w = gkd_weight(step, total_steps, float(gkd_cfg.get("ratio", 0.3)))
+                if w > 0:
+                    sub = int(gkd_cfg.get("sample_sub", 2))
+                    with torch.no_grad():
+                        gen = student.generate(
+                            input_ids=ids[:sub], attention_mask=am[:sub],
+                            do_sample=True,
+                            temperature=float(gkd_cfg.get("temperature", 1.0)),
+                            max_new_tokens=int(gkd_cfg.get("sample_cap", 1024)),
+                        )
+                    # gen is the decoder-side sample; score both models on
+                    # (prompt -> sample). Skip the decoder-start token.
+                    cont_mask = torch.zeros_like(gen, dtype=torch.bool)
+                    cont_mask[:, 1:] = True
+                    s_logits = student(
+                        input_ids=ids[:sub], attention_mask=am[:sub], labels=gen
+                    ).logits
+                    with torch.no_grad():
+                        t_logits = teacher(
+                            input_ids=ids[:sub], attention_mask=am[:sub], labels=gen
+                        ).logits
+                    loss = loss + w * reverse_kl(s_logits, t_logits, gen, cont_mask)
             loss.backward()
             torch.nn.utils.clip_grad_norm_(student.parameters(), 1.0)
             optimizer.step()

tests/test_gkd.py

@@ -0,0 +1,60 @@
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from gpu.gkd import gkd_weight, reverse_kl, token_logprobs  # noqa: E402
+
+
+class TestGkdWeight:
+    def test_full_ratio_before_anneal(self):
+        assert gkd_weight(0, 300, 0.3) == 0.3
+        assert gkd_weight(199, 300, 0.3) == 0.3
+
+    def test_linear_anneal_over_final_third(self):
+        assert gkd_weight(200, 300, 0.3) == pytest.approx(0.3)  # anneal starts at full
+        assert gkd_weight(250, 300, 0.3) == pytest.approx(0.15)
+        assert gkd_weight(299, 300, 0.3) == pytest.approx(0.003)
+
+    def test_zero_at_end_and_degenerate(self):
+        assert gkd_weight(300, 300, 0.3) == 0.0
+        assert gkd_weight(0, 0, 0.3) == 0.0
+
+
+class TestTokenLogprobs:
+    def test_shift_and_gather(self):
+        torch.manual_seed(0)
+        logits = torch.randn(2, 5, 7)
+        targets = torch.randint(0, 7, (2, 5))
+        lp = token_logprobs(logits, targets)
+        assert lp.shape == (2, 4)
+        ref = torch.log_softmax(logits[0, :-1].float(), -1)
+        expect = ref.gather(-1, targets[0, 1:].unsqueeze(-1)).squeeze(-1)
+        assert torch.allclose(lp[0], expect)
+
+
+class TestReverseKl:
+    def test_identical_distributions_zero(self):
+        torch.manual_seed(0)
+        logits = torch.randn(2, 6, 7)
+        targets = torch.randint(0, 7, (2, 6))
+        assert reverse_kl(logits, logits.clone(), targets).abs() < 1e-6
+
+    def test_mask_restricts_positions(self):
+        torch.manual_seed(0)
+        s = torch.randn(1, 6, 7)
+        t = torch.randn(1, 6, 7)
+        targets = torch.randint(0, 7, (1, 6))
+        mask = torch.zeros(1, 6, dtype=torch.bool)
+        mask[:, 4:] = True  # only the continuation region (positions >= 4)
+        masked = reverse_kl(s, t, targets, mask)
+        full = reverse_kl(s, t, targets)
+        assert not torch.allclose(masked, full)
+
+    def test_gradient_flows_through_student_only(self):
+        s = torch.randn(1, 4, 5, requires_grad=True)
+        t = torch.randn(1, 4, 5, requires_grad=True)
+        targets = torch.randint(0, 5, (1, 4))
+        loss = reverse_kl(s, t, targets)
+        loss.backward()
+        assert s.grad is not None
+        assert t.grad is None