Skip to content

feat(gkd): on-policy distillation wiring + spec (armed) - #148

Merged
ronaldtse merged 1 commit into
mainfrom
feat/gkd-rung
Sep 3, 2026
Merged

feat(gkd): on-policy distillation wiring + spec (armed)#148
ronaldtse merged 1 commit into
mainfrom
feat/gkd-rung

Conversation

@ronaldtse

Copy link
Copy Markdown
Contributor

GKD rung implemented per its EXPERIMENTS.md registration (armed in #147):

  • src/gpu/gkd.py — modal-free pure-torch helpers: gkd_weight (full ratio to 2/3 of training, linear to zero over the final third), token_logprobs, reverse_kl (student-only gradients; teacher scored under no_grad). 7 unit tests (importorskip torch, the MTP-test pattern).
  • distill_sequence wiring — teacher moved back to cuda for scoring; every sample_every=4 steps a sub-batch of 2 student-sampled sequences (temp 1.0, cap 1024) is scored by both models on (prompt -> sample) and the reverse-KL enters the loss at the annealed weight. Control-verbatim otherwise: same units, schedule, Muon, 30k units.
  • Spec ara-diac-small-2-gkd (run-011-r7-muon-gkd) copied field-for-field from run-006 with only the gkd block added.
  • Fixed pre-existing B905 in scripts/figures/frontier.py (was making main red on ruff).

Launch remains gated on a GPU slot (G2b + one other app currently hold the budget; GKD launches when a slot frees).

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).
@ronaldtse
ronaldtse merged commit a120384 into main Sep 3, 2026
11 checks passed
@ronaldtse
ronaldtse deleted the feat/gkd-rung branch September 3, 2026 10:08
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant