feat: GKD with external teacher backends + on-policy + JSD mix

[marksverdhei]

Summary

Adds Generalized Knowledge Distillation (GKD) on top of the context-baking machinery from #35. The student can now distill from a separate teacher — a different HuggingFace model, a vLLM server, or any OpenAI-compatible API with logprobs — while continuing to bake an arbitrary prefix context into its weights in the same training sweep.

LoRA remains the default student adapter (it regularizes against capability degradation under the combined objective), but the trainer is no longer LoRA-gated — separate-teacher modes work with full FT too.

What's in

Teacher backend abstraction (src/bakery/teachers/)

• TeacherBackend ABC returning TopKLogprobs (sparse top-k view of the teacher's distribution, renormalized over those K)
• HFTeacher — in-process Transformers model (dev/test path)
• VLLMTeacher — vLLM as OpenAI-compat server (production, TRL-style)
• OpenAIAPITeacher — any OpenAI-compat endpoint that exposes logprobs (Together, Fireworks, vLLM itself)

The same TopKLogprobs shape works for both dense (K=vocab) and sparse (K=20) teachers; the trainer's KL math is identical across backends.

Sparse KL math (src/bakery/kl.py)

• topk_forward_kl — forward KL over the teacher's top-k support. At K=V, identical to dense KL.
• topk_jsd — mixed forward/reverse KL: loss = (1-β)·KL(P_t||P_s) + β·KL(P_s||P_t). β=0 is forward (default), β=1 is reverse, β=0.5 is symmetric. Same convention as TRL's GKDTrainer.

Trainer integration (src/bakery/trainer.py)

• teacher_backend + teacher_top_k constructor args
• gkd_on_policy_fraction — probability of sampling the trajectory from the student (mode-seeking, on-policy GKD)
• gkd_jsd_beta — routes loss to topk_jsd instead of topk_forward_kl
• Unified _align_and_slice helper shared by dense and sparse KL paths
• _sample_from_student for the on-policy path (adapters enabled, student's trimmed view)

Config (src/bakery/config.py)

New TeacherConfig exposed as a 5th HfArgumentParser dataclass. All fields prefixed with teacher_ to avoid CLI flag collisions with student DataConfig:

• teacher_backend: local-toggle (default), hf, vllm, openai
• teacher_model_name_or_path, teacher_api_base, teacher_api_key, teacher_api_model
• teacher_top_k, teacher_torch_dtype, teacher_device, teacher_attn_implementation
• gkd_on_policy_fraction, gkd_jsd_beta

Example

examples/gkd_gemma3.yaml — Gemma 3 1B teacher ↦ 270M student with prefix context baked in one sweep, LoRA default.

Backward compatibility

• teacher_backend=local-toggle (the default) keeps every prior bakery code path bit-identical: same dense KL, same adapter-toggle teacher, same trajectory generation.
• gkd_jsd_beta=0 (default) reduces topk_jsd exactly to topk_forward_kl.
• gkd_on_policy_fraction=0 (default) keeps trajectories sampled from the teacher.
• Older examples (basic.yaml, multi_turn_prefix.yaml, etc.) are unchanged.

Test plan

• 50 new tests under test_topk_kl.py (incl. JSD identities), test_teacher_backends.py, test_openai_teacher_score.py (mocked vLLM payloads — runs offline), test_trainer_gkd.py (compute_loss + prediction_step end-to-end with external teacher + on-policy routing + guardrails)
• Full CPU suite: 315 passing (pytest -m "not gpu and not benchmark")
• GPU smoke (RTX 3090, Gemma 3 270M ← Gemma 3 1B HF teacher):
  • off-policy forward KL (default): loss 6.108
  • on-policy gkd_on_policy_fraction=1.0 + gkd_jsd_beta=0.5: loss 4.099
• Backward compat: local-toggle path verified unchanged via test_local_toggle_path_unchanged

Known limitations / follow-ups

• OpenAIAPITeacher.score() assumes the server's tokenizer matches the student's (Gemma family → Gemma student, Qwen → Qwen, etc.). Cross-tokenizer distillation (e.g. Llama→Qwen) is out of scope.
• VLLMInProcessTeacher is stubbed (raises NotImplementedError). Use VLLMTeacher (HTTP) or HFTeacher for now.
• True GKD with on-policy + sequence-level reward is partial: the student sampler is wired, but trajectory selection isn't yet tied to teacher-score quality.

[marksverdhei]

Independent verification during hivemind triage:

• Cloned PR head locally, uv sync clean
• uv run pytest -q -x --ignore=tests/test_benchmark.py → 315 passed, 3 skipped in 119.22s
• No failures, no errors
• Warnings: peft fan_in_fan_out for Conv1D (upstream noise) + the new PromptBakingTrainer is deprecated; use ContextBakingTrainer warning that this PR introduces by design

PR state: MERGEABLE/CLEAN, both CI checks green (prek + pytest), main hasn't moved since #35 landed on the same day this branch was opened. Ready to squash-merge whenever you sign off — flagging because it's been open since May 11 and might just need a final eye.

[marksverdhei]

Reviewed the distillation core (kl.py math, the teacher backends, and the trainer alignment). The math in kl.py is solid, but I think there's a position off-by-one in the OpenAI/vLLM API teacher that would silently corrupt GKD against an API-served teacher (the headline path — gkd_olivia_*.yaml). Flagging for verification before merge.

The alignment convention is inconsistent between teacher backends

Dense HFTeacher (logits convention): score() does topk_from_logits(self.model(...).logits). A causal LM's logits[t] = P(x[t+1] | x[0..t]), so its TopKLogprobs.values[t] is the distribution predicting x[t+1].

OpenAIAPITeacher (predicted-position convention): it stores top_logprobs[i] at t_full = real_start + i (openai_compat.py:137-140). But for an echo completion, top_logprobs[i] = P(x[i] | x[0..i-1]), i.e. it predicts x[i], not x[i+1].

Put precisely: top_logprobs[i] ≡ dense logits[i-1]. So the API teacher ends up with values[t] = logits[t-1], while the dense teacher has values[t] = logits[t] — the two backends are shifted one position apart.

Why that breaks training

_kl_from_logits / the top-k path apply one shared shift — _align_and_slice uses teacher_data[t_start-1:-1] with the comment "Logits predict next token → shift by 1" (trainer.py:364-366, and t_vals = teacher_logprobs[t_start-1:-1] at ~457). That shift is correct for the dense teacher but, applied to the API teacher's predicted-position tensor, aligns the student's distribution for x[t+1] against the teacher's distribution for x[t-1]. Every GKD step against a vLLM/OpenAI teacher then distills mode-by-mode against the wrong position.

Why CI is green anyway

test_openai_teacher_score.py mocks httpx.Client and asserts shape + renormalization (exp(values).sum≈1) and the null-first-position skip — none of which pin the cross-position alignment. test_trainer_gkd.py mocks the teacher backend, so it never exercises the real echo→position mapping. (Same shape of issue as a mocked API call passing while the real contract is wrong.)

Suggested fix

Store the echoed distribution one slot earlier so the API teacher matches the dense/logits convention:

# top_logprobs[i] == logits[i-1]; store at i-1 to match HFTeacher (values[t] predicts x[t+1])
t_full = real_start + t_local - 1
if t_full < real_start:        # t_local==0 is null anyway; guard servers that emit it
    continue

and a regression test at the alignment level: feed a mocked echo response whose top token at each position is a known id, and assert indices[t] lands where the dense teacher would put logits[t] (e.g. compare against topk_from_logits on synthetic logits for the same sequence).

Smaller note (not blocking)

topk_jsd computes (1-β)·KL(t‖s) + β·KL(s‖t) — a direct forward/reverse-KL interpolation. That's a valid objective and matches the β=0→forward, β=1→reverse endpoints, but it isn't TRL GKDTrainer's loss: TRL uses the mixture-distribution generalized JSD (M = β·s + (1-β)·t, then β·KL(s‖M)+(1-β)·KL(t‖M)). The docstring's "same convention TRL's GKDTrainer uses" overstates it slightly — worth softening or switching to the true generalized-JSD if TRL parity is intended.

The rest (topk_forward_kl gather+renormalize, the HPC trust_env=False proxy guard, the top-k renormalization) looks correct to me. 👍

[marksverdhei]

Scope follow-up on the alignment off-by-one above: it isn't isolated to OpenAIAPITeacher — VLLMTeacher inherits OpenAIAPITeacher.score() verbatim (it only flips echo_supported=True), so the same shift applies to the vLLM HTTP path too, which vllm.py's own docstring calls "the production path that mirrors TRL's pattern."

So across the four backends:

• HFTeacher (dense, topk_from_logits on raw logits) — logits convention, correct.
• OpenAIAPITeacher — predicted-position convention → off by one.
• VLLMTeacher(OpenAIAPITeacher) — inherits score() → off by one (production path).
• VLLMInProcessTeacher — NotImplementedError stub.

Net: every implemented HTTP teacher is affected and only the dense in-process teacher is aligned. Fixing it once in OpenAIAPITeacher.score() (store top_logprobs[i] at i-1) covers both HTTP backends. Worth a same-tokenizer end-to-end check too: score a short sequence through HFTeacher and through VLLMTeacher against the same model and assert the per-position top-1 ids match — that would have caught this and pins the two conventions together going forward.

[marksverdhei]

Finished reviewing the rest of the PR (config, the teacher factory, the dispatch, and the on-policy path) so you know the blast radius — the alignment off-by-one is the only blocking issue I found; everything else looks sound:

• config.TeacherConfig — clean; teacher_-prefixed fields avoid the DataConfig collision, defaults are sensible (local-toggle, top_k=64, on_policy_fraction=0.0, jsd_beta=0.0), help text is clear.
• make_teacher (the config→backend factory) — correct: normalizes the backend string, local-toggle/local/self/none/"" → None, hf/vllm/openai wired with the right fields (model falls back to teacher_model_name_or_path), unknown → a clear ValueError.
• _generate_trajectory on-policy routing — correct probabilistic split (random.random() < on_policy_fraction → student sample, else teacher/disable_adapters); the gkd_on_policy_fraction bounds check in __init__ is good. (nit: the import random is inside the method — fine, just unusual.)
• _teacher_forward / _kl_dispatch / compute_loss / prediction_step — clean dispatch (dense logits for local-toggle, top-k for external), empty-batch + None-loss guards, and the sequential_eval CPU-offload path all look right.

So: solid architecture overall — just the API/vLLM teacher position shift (and the smaller topk_jsd vs TRL-JSD doc nuance) to resolve. Nice PR once that's in. 👍