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

examples/gkd_gemma3.yaml

@@ -0,0 +1,63 @@
+# GKD + context baking: distill Gemma 3 1B (teacher) into Gemma 3 270M (student)
+# in a single sweep, with a prefix context baked into the student weights.
+#
+# Both models share the same tokenizer, so we can score student-tokenized
+# sequences against the teacher directly. The student picks up:
+#   - the prefix context (baked away from inference-time tokens)
+#   - the teacher's stronger behavior (top-k forward KL)
+#
+# LoRA stays on by default as a regularizer to avoid degradation under the
+# combined objective.
+#
+# Usage:
+#   bakery --config examples/gkd_gemma3.yaml
+
+# --- Standard TrainingArguments ---
+output_dir: "./outputs/gkd_gemma3"
+num_train_epochs: 1
+learning_rate: 1e-4
+per_device_train_batch_size: 1
+gradient_accumulation_steps: 4
+logging_steps: 5
+save_strategy: "epoch"
+bf16: true
+report_to: "none"
+seed: 42
+
+# --- Student model (the one we're training) ---
+model_name_or_path: "google/gemma-3-270m-it"
+torch_dtype: "bfloat16"
+attn_implementation: "sdpa"
+
+# --- Teacher: a separate Gemma 3 1B, in-process HF ---
+teacher_backend: "hf"
+teacher_model_name_or_path: "google/gemma-3-1b-it"
+teacher_torch_dtype: "bfloat16"
+teacher_top_k: 64
+
+# --- Context to bake (in addition to teacher capability) ---
+prefix_messages:
+  - role: system
+    content: "Answer in one short sentence."
+student_retained_turns: 0
+target_roles:
+  - assistant
+
+# --- LoRA (default, regularizes student under the combined objective) ---
+r: 16
+lora_alpha: 32
+target_modules:
+  - q_proj
+  - k_proj
+  - v_proj
+  - o_proj
+lora_dropout: 0.05
+
+# --- Data: small prompt set for a smoke run ---
+training_prompts:
+  - "What is the capital of France?"
+  - "What color is the sky?"
+  - "Name a prime number greater than ten."
+  - "What is two plus two?"
+num_trajectories: 1
+trajectory_length: 32

examples/gkd_olivia_12b.yaml

@@ -0,0 +1,70 @@
+# GKD on Olivia: distill Gemma 3 27B aurora-SFT (teacher) into Gemma 3 12B (student).
+#
+# Teacher is served by vLLM (TP=2) on the same node — `olivia/run_gkd_node.sh`
+# brings it up on GPU 0+1 before launching this trainer. The trainer connects
+# over localhost via the OpenAI-compatible /v1/completions endpoint with
+# echo=True, logprobs=K=64 (sparse top-k forward path in src/bakery/kl.py).
+#
+# Student is the matching aurora-SFT 12B checkpoint, so we are distilling the
+# 27B teacher's behavior *on top of* the same SFT initialization, not from
+# scratch. LoRA r=16 keeps the student trainable in a single GH200 (96GB),
+# and lets us co-locate it with the 27B-LoRA student on the same GPU per
+# Markus's call.
+#
+# Usage on Olivia:
+#   sbatch olivia/run_gkd_node.sh examples/gkd_olivia_12b.yaml
+
+# --- Standard TrainingArguments ---
+output_dir: "/cluster/projects/nn30001k/markuhei/bakery/outputs/gkd_aurora_27b_to_12b"
+num_train_epochs: 1
+learning_rate: 1e-4
+per_device_train_batch_size: 1
+gradient_accumulation_steps: 8
+logging_steps: 10
+save_strategy: "epoch"
+bf16: true
+gradient_checkpointing: true
+report_to: "none"
+seed: 42
+max_seq_length: 2048
+
+# --- Student model (the one we're training) ---
+model_name_or_path: "NbAiLab/nb-gpt-gemma3-12b-instruct-epoch-3-aurora-sft-2603-posttrain"
+torch_dtype: "bfloat16"
+attn_implementation: "sdpa"
+
+# --- Teacher: 27B aurora-SFT, served by vLLM TP=2 on localhost ---
+teacher_backend: "vllm"
+teacher_api_base: "http://localhost:8765/v1"
+teacher_api_model: "NbAiLab/nb-gpt-gemma3-27b-instruct-epoch-3-aurora-sft-2603-posttrain"
+teacher_top_k: 64
+
+# --- GKD recipe ---
+gkd_on_policy_fraction: 0.5
+gkd_jsd_beta: 0.5
+temperature: 1.0
+
+# --- No global prefix: aurora rows carry their own system + chat history ---
+# Both teacher and student were SFT-trained on this exact distribution; injecting
+# a synthetic system message would only push us off it. The trainer sees each
+# row's full `messages` list as-is. `student_retained_turns` is moot when the
+# prefix is empty (it only trims prefix_messages, never the row's turns).
+prefix_messages: []
+target_roles:
+  - assistant
+
+# --- LoRA ---
+r: 16
+lora_alpha: 32
+target_modules:
+  - q_proj
+  - k_proj
+  - v_proj
+  - o_proj
+lora_dropout: 0.05
+
+# --- Data: aurora SFT 2603 ---
+dataset: "NbAiLab/aurora-sft-2603"
+dataset_split: "train"
+trajectory_length: 256
+num_trajectories: 1

examples/gkd_olivia_1b.yaml

@@ -0,0 +1,52 @@
+# GKD on Olivia: distill Gemma 3 27B aurora-SFT (teacher) into Gemma 3 1B (student).
+# Sibling of gkd_olivia_12b.yaml — same recipe, smaller student.
+#
+# Usage on Olivia:
+#   sbatch olivia/run_gkd_node.sh examples/gkd_olivia_1b.yaml
+
+output_dir: "/cluster/projects/nn30001k/markuhei/bakery/outputs/gkd_aurora_27b_to_1b"
+num_train_epochs: 1
+learning_rate: 2e-4
+per_device_train_batch_size: 4
+gradient_accumulation_steps: 2
+logging_steps: 10
+save_strategy: "epoch"
+bf16: true
+gradient_checkpointing: true
+report_to: "none"
+seed: 42
+max_seq_length: 2048
+
+model_name_or_path: "NbAiLab/nb-gpt-gemma3-1b-instruct-epoch-3-aurora-sft-2603-posttrain"
+torch_dtype: "bfloat16"
+attn_implementation: "sdpa"
+
+teacher_backend: "vllm"
+teacher_api_base: "http://localhost:8765/v1"
+teacher_api_model: "NbAiLab/nb-gpt-gemma3-27b-instruct-epoch-3-aurora-sft-2603-posttrain"
+teacher_top_k: 64
+
+gkd_on_policy_fraction: 0.5
+gkd_jsd_beta: 0.5
+temperature: 1.0
+
+prefix_messages: []
+target_roles:
+  - assistant
+
+r: 16
+lora_alpha: 32
+target_modules:
+  - q_proj
+  - k_proj
+  - v_proj
+  - o_proj
+  - gate_proj
+  - up_proj
+  - down_proj
+lora_dropout: 0.05
+
+dataset: "NbAiLab/aurora-sft-2603"
+dataset_split: "train"
+trajectory_length: 256
+num_trajectories: 1

examples/gkd_olivia_270m.yaml

@@ -0,0 +1,52 @@
+# GKD on Olivia: distill Gemma 3 27B aurora-SFT (teacher) into Gemma 3 270M (student).
+# Sibling of gkd_olivia_12b.yaml — same recipe, smallest student.
+#
+# Usage on Olivia:
+#   sbatch olivia/run_gkd_node.sh examples/gkd_olivia_270m.yaml
+
+output_dir: "/cluster/projects/nn30001k/markuhei/bakery/outputs/gkd_aurora_27b_to_270m"
+num_train_epochs: 1
+learning_rate: 3e-4
+per_device_train_batch_size: 8
+gradient_accumulation_steps: 1
+logging_steps: 10
+save_strategy: "epoch"
+bf16: true
+gradient_checkpointing: true
+report_to: "none"
+seed: 42
+max_seq_length: 2048
+
+model_name_or_path: "NbAiLab/nb-gpt-gemma3-270m-instruct-epoch-3-aurora-sft-2603-posttrain"
+torch_dtype: "bfloat16"
+attn_implementation: "sdpa"
+
+teacher_backend: "vllm"
+teacher_api_base: "http://localhost:8765/v1"
+teacher_api_model: "NbAiLab/nb-gpt-gemma3-27b-instruct-epoch-3-aurora-sft-2603-posttrain"
+teacher_top_k: 64
+
+gkd_on_policy_fraction: 0.5
+gkd_jsd_beta: 0.5
+temperature: 1.0
+
+prefix_messages: []
+target_roles:
+  - assistant
+
+r: 32
+lora_alpha: 64
+target_modules:
+  - q_proj
+  - k_proj
+  - v_proj
+  - o_proj
+  - gate_proj
+  - up_proj
+  - down_proj
+lora_dropout: 0.05
+
+dataset: "NbAiLab/aurora-sft-2603"
+dataset_split: "train"
+trajectory_length: 256
+num_trajectories: 1

examples/gkd_olivia_27b.yaml

@@ -0,0 +1,57 @@
+# GKD on Olivia: distill Gemma 3 27B aurora-SFT into itself (LoRA self-distillation).
+#
+# DEFERRED — wired but NOT in scope for the first round per olivia-lead's call.
+# Greenlight after the 12B run produces a sane loss curve.
+#
+# Sibling of gkd_olivia_12b.yaml — same recipe, but the student is the same
+# 27B SFT checkpoint that the teacher serves. This is self-distillation under
+# the GKD objective: useful only if there is structure in the JSD-beta=0.5
+# (mix of forward + reverse KL) that a LoRA on top of the SFT base can learn.
+# Memory budget is tight on a single GH200 — 54GB base + activations + LoRA.
+#
+# Usage on Olivia (when greenlit):
+#   sbatch olivia/run_gkd_node.sh examples/gkd_olivia_27b.yaml
+
+output_dir: "/cluster/projects/nn30001k/markuhei/bakery/outputs/gkd_aurora_27b_self"
+num_train_epochs: 1
+learning_rate: 5e-5
+per_device_train_batch_size: 1
+gradient_accumulation_steps: 16
+logging_steps: 10
+save_strategy: "epoch"
+bf16: true
+gradient_checkpointing: true
+report_to: "none"
+seed: 42
+max_seq_length: 2048
+
+model_name_or_path: "NbAiLab/nb-gpt-gemma3-27b-instruct-epoch-3-aurora-sft-2603-posttrain"
+torch_dtype: "bfloat16"
+attn_implementation: "sdpa"
+
+teacher_backend: "vllm"
+teacher_api_base: "http://localhost:8765/v1"
+teacher_api_model: "NbAiLab/nb-gpt-gemma3-27b-instruct-epoch-3-aurora-sft-2603-posttrain"
+teacher_top_k: 64
+
+gkd_on_policy_fraction: 0.5
+gkd_jsd_beta: 0.5
+temperature: 1.0
+
+prefix_messages: []
+target_roles:
+  - assistant
+
+r: 16
+lora_alpha: 32
+target_modules:
+  - q_proj
+  - k_proj
+  - v_proj
+  - o_proj
+lora_dropout: 0.05
+
+dataset: "NbAiLab/aurora-sft-2603"
+dataset_split: "train"
+trajectory_length: 256
+num_trajectories: 1

examples/gkd_olivia_4b.yaml

@@ -0,0 +1,53 @@
+# GKD on Olivia: distill Gemma 3 27B aurora-SFT (teacher) into Gemma 3 4B (student).
+# Sibling of gkd_olivia_12b.yaml — same recipe, smaller student. See that file
+# for the full rationale.
+#
+# Usage on Olivia:
+#   sbatch olivia/run_gkd_node.sh examples/gkd_olivia_4b.yaml
+
+output_dir: "/cluster/projects/nn30001k/markuhei/bakery/outputs/gkd_aurora_27b_to_4b"
+num_train_epochs: 1
+learning_rate: 2e-4
+per_device_train_batch_size: 2
+gradient_accumulation_steps: 4
+logging_steps: 10
+save_strategy: "epoch"
+bf16: true
+gradient_checkpointing: true
+report_to: "none"
+seed: 42
+max_seq_length: 2048
+
+model_name_or_path: "NbAiLab/nb-gpt-gemma3-4b-instruct-epoch-3-aurora-sft-2603-posttrain"
+torch_dtype: "bfloat16"
+attn_implementation: "sdpa"
+
+teacher_backend: "vllm"
+teacher_api_base: "http://localhost:8765/v1"
+teacher_api_model: "NbAiLab/nb-gpt-gemma3-27b-instruct-epoch-3-aurora-sft-2603-posttrain"
+teacher_top_k: 64
+
+gkd_on_policy_fraction: 0.5
+gkd_jsd_beta: 0.5
+temperature: 1.0
+
+prefix_messages: []
+target_roles:
+  - assistant
+
+r: 16
+lora_alpha: 32
+target_modules:
+  - q_proj
+  - k_proj
+  - v_proj
+  - o_proj
+  - gate_proj
+  - up_proj
+  - down_proj
+lora_dropout: 0.05
+
+dataset: "NbAiLab/aurora-sft-2603"
+dataset_split: "train"
+trajectory_length: 256
+num_trajectories: 1

src/bakery/__init__.py

@@ -3,7 +3,13 @@
 Context baking (prefix-context distillation) via KL divergence with LoRA.
 """
 
-from bakery.config import BakeryConfig, ContextConfig, DataConfig, LoraConfig
+from bakery.config import (
+    BakeryConfig,
+    ContextConfig,
+    DataConfig,
+    LoraConfig,
+    TeacherConfig,
+)
 from bakery.trainer import ContextBakingTrainer, PromptBakingTrainer
 from bakery.data import (
     create_conversational_dataset,
@@ -12,14 +18,16 @@
     load_dataset,
     prompt_baking_collator,
 )
-from bakery.kl import compute_kl_divergence
+from bakery.kl import compute_kl_divergence, topk_forward_kl
 from bakery.masking import build_target_mask
+from bakery.teachers import HFTeacher, TeacherBackend, TopKLogprobs, make_teacher
 
 __all__ = [
     "BakeryConfig",
     "ContextConfig",
     "DataConfig",
     "LoraConfig",
+    "TeacherConfig",
     "ContextBakingTrainer",
     "PromptBakingTrainer",
     "create_conversational_dataset",
@@ -28,5 +36,10 @@
     "load_dataset",
     "prompt_baking_collator",
     "compute_kl_divergence",
+    "topk_forward_kl",
     "build_target_mask",
+    "HFTeacher",
+    "TeacherBackend",
+    "TopKLogprobs",
+    "make_teacher",
 ]