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.