diff --git a/scripts/figures/frontier.py b/scripts/figures/frontier.py index cab8720..71207b9 100644 --- a/scripts/figures/frontier.py +++ b/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") diff --git a/src/gpu/distill_specs.yaml b/src/gpu/distill_specs.yaml index ea0e65e..f954c50 100644 --- a/src/gpu/distill_specs.yaml +++ b/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 diff --git a/src/gpu/gkd.py b/src/gpu/gkd.py new file mode 100644 index 0000000..3acb14e --- /dev/null +++ b/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() diff --git a/src/gpu/modal_distill.py b/src/gpu/modal_distill.py index e76ccd8..7f70aa6 100644 --- a/src/gpu/modal_distill.py +++ b/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() diff --git a/tests/test_gkd.py b/tests/test_gkd.py new file mode 100644 index 0000000..4bdc63d --- /dev/null +++ b/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