diff --git a/examples/gkd_gemma3.yaml b/examples/gkd_gemma3.yaml new file mode 100644 index 0000000..94043ca --- /dev/null +++ b/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 diff --git a/examples/gkd_olivia_12b.yaml b/examples/gkd_olivia_12b.yaml new file mode 100644 index 0000000..2cf2aa4 --- /dev/null +++ b/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 diff --git a/examples/gkd_olivia_1b.yaml b/examples/gkd_olivia_1b.yaml new file mode 100644 index 0000000..2e892ff --- /dev/null +++ b/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 diff --git a/examples/gkd_olivia_270m.yaml b/examples/gkd_olivia_270m.yaml new file mode 100644 index 0000000..0e967d3 --- /dev/null +++ b/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 diff --git a/examples/gkd_olivia_27b.yaml b/examples/gkd_olivia_27b.yaml new file mode 100644 index 0000000..bf9defb --- /dev/null +++ b/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 diff --git a/examples/gkd_olivia_4b.yaml b/examples/gkd_olivia_4b.yaml new file mode 100644 index 0000000..a05bfa8 --- /dev/null +++ b/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 diff --git a/src/bakery/__init__.py b/src/bakery/__init__.py index 066bcd2..749aa9b 100644 --- a/src/bakery/__init__.py +++ b/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", ] diff --git a/src/bakery/cli.py b/src/bakery/cli.py index 06fefb6..215f371 100644 --- a/src/bakery/cli.py +++ b/src/bakery/cli.py @@ -19,7 +19,14 @@ ) from peft import LoraConfig as PeftLoraConfig, get_peft_model -from bakery.config import BakeryConfig, ContextConfig, DataConfig, LoraConfig +from bakery.config import ( + BakeryConfig, + ContextConfig, + DataConfig, + LoraConfig, + TeacherConfig, +) +from bakery.teachers import make_teacher from bakery.trainer import ContextBakingTrainer from bakery.data import ( create_conversational_dataset, @@ -78,11 +85,17 @@ def main(): pre_args, remaining_args = pre_parser.parse_known_args() config_file = pre_args.config - parser = HfArgumentParser((BakeryConfig, DataConfig, LoraConfig, ContextConfig)) - - baking_config, data_config, lora_config, context_config = parser.parse_yaml_file( - config_file, allow_extra_keys=True + parser = HfArgumentParser( + (BakeryConfig, DataConfig, LoraConfig, ContextConfig, TeacherConfig) ) + + ( + baking_config, + data_config, + lora_config, + context_config, + teacher_config, + ) = parser.parse_yaml_file(config_file, allow_extra_keys=True) # Apply CLI overrides on top of YAML config. # We parse the remaining CLI args into fresh dataclasses, then detect # which fields were explicitly set by comparing against a baseline @@ -95,15 +108,15 @@ def main(): explicit_keys.add(arg.lstrip("-").replace("-", "_")) override_parser = HfArgumentParser( - (BakeryConfig, DataConfig, LoraConfig, ContextConfig) + (BakeryConfig, DataConfig, LoraConfig, ContextConfig, TeacherConfig) ) overrides = override_parser.parse_args_into_dataclasses( args=["--output_dir", baking_config.output_dir] + remaining_args, return_remaining_strings=True, ) for override_cfg, base_cfg in zip( - overrides[:4], - (baking_config, data_config, lora_config, context_config), + overrides[:5], + (baking_config, data_config, lora_config, context_config, teacher_config), ): for k, v in vars(override_cfg).items(): if k in explicit_keys: @@ -141,12 +154,6 @@ def main(): {"role": "system", "content": baking_config.system_prompt} ] - if not context_config.prefix_messages: - raise ValueError( - "No prefix context configured. Set prefix_messages (inline or via " - "prefix_messages_file), or the deprecated system_prompt / corpus_file." - ) - # Load data. Use conversational loader when the source preserves multi-turn # history (HF `messages` column or JSON rows with `messages`/`prefix_messages`); # otherwise use the simple (prompts, responses) path. @@ -162,6 +169,14 @@ def main(): conversational_rows = None if conversational_rows is None: training_prompts, precomputed_responses = load_data(data_config) + + if not context_config.prefix_messages and conversational_rows is None: + raise ValueError( + "No prefix context configured. Set prefix_messages (inline or via " + "prefix_messages_file), or the deprecated system_prompt / corpus_file. " + "Empty prefix is only valid when conversational rows carry their own " + "messages (e.g. an SFT chat dataset)." + ) eval_qa = load_eval_data(data_config.eval_file) heldout_qa = load_eval_data(data_config.heldout_file) @@ -437,10 +452,22 @@ def forward(self, x): eval_dataset = create_dataset(eval_prompts, eval_responses) print(f" Eval samples: {len(eval_dataset)}") + teacher_backend = make_teacher(teacher_config) + if teacher_backend is not None: + print( + f" Teacher backend: {teacher_config.teacher_backend} " + f"({teacher_config.teacher_model_name_or_path or teacher_config.teacher_api_model}); " + f"top_k={teacher_config.teacher_top_k}" + ) + trainer = ContextBakingTrainer( model=model, args=baking_config, context_config=context_config, + teacher_backend=teacher_backend, + teacher_top_k=teacher_config.teacher_top_k, + gkd_on_policy_fraction=teacher_config.gkd_on_policy_fraction, + gkd_jsd_beta=teacher_config.gkd_jsd_beta, train_dataset=train_dataset, eval_dataset=eval_dataset, processing_class=tokenizer, diff --git a/src/bakery/config.py b/src/bakery/config.py index 6233115..4e13b80 100644 --- a/src/bakery/config.py +++ b/src/bakery/config.py @@ -248,6 +248,100 @@ def __post_init__(self): bias: str = field(default="none", metadata={"help": "LoRA bias mode."}) +@dataclass +class TeacherConfig: + """External teacher backend configuration for GKD / context distillation. + + Default `teacher_backend="local-toggle"` keeps the original bakery behavior: + the student model with LoRA adapters serves as its own teacher when the + adapters are disabled. Set to "hf", "vllm", or "openai" to use a separate + teacher model — required for distilling a stronger model into a smaller one + ("elevated skill" + prefix context baked in one sweep). + + All field names are prefixed with `teacher_` to avoid CLI flag collisions + with DataConfig (which controls the student model). + """ + + teacher_backend: str = field( + default="local-toggle", + metadata={ + "help": "Teacher source: 'local-toggle' (default; same model, " + "adapters off), 'hf' (separate HuggingFace model), 'vllm' (vLLM " + "OpenAI-compatible server), or 'openai' (any OpenAI-compat API " + "with logprobs)." + }, + ) + teacher_model_name_or_path: Optional[str] = field( + default=None, + metadata={ + "help": "Teacher model name (required for 'hf'). For 'vllm'/'openai' " + "this is the model identifier the server expects (falls back to teacher_api_model)." + }, + ) + teacher_api_base: Optional[str] = field( + default=None, + metadata={ + "help": "Base URL for OpenAI-compat / vLLM API (e.g. http://localhost:8000/v1)." + }, + ) + teacher_api_key: Optional[str] = field( + default=None, + metadata={"help": "API key; falls back to OPENAI_API_KEY env var."}, + ) + teacher_api_model: Optional[str] = field( + default=None, + metadata={ + "help": "Model name to send in API requests (when different from " + "teacher_model_name_or_path)." + }, + ) + teacher_torch_dtype: str = field( + default="bfloat16", + metadata={"help": "Dtype for HF teacher (float32, float16, bfloat16)."}, + ) + teacher_device: Optional[str] = field( + default=None, + metadata={ + "help": "Device for HF teacher: 'cuda', 'cuda:1', 'cpu', or None for auto." + }, + ) + teacher_trust_remote_code: bool = field( + default=False, + metadata={"help": "Trust remote code when loading HF teacher."}, + ) + teacher_attn_implementation: Optional[str] = field( + default=None, + metadata={ + "help": "Attention impl for HF teacher: flash_attention_2, sdpa, eager." + }, + ) + teacher_top_k: int = field( + default=64, + metadata={ + "help": "How many tokens of the teacher's distribution to keep per " + "position when computing KL. K covering the full vocab → identical to " + "dense KL. Lower K → cheaper, sparser. API backends often cap at 20." + }, + ) + gkd_on_policy_fraction: float = field( + default=0.0, + metadata={ + "help": "Fraction of training trajectories sampled on-policy from the " + "student (vs. from the teacher). 0.0 = pure off-policy distillation " + "(default). 1.0 = pure on-policy (the GKD recipe). Mixed values yield " + "stochastic per-batch routing." + }, + ) + gkd_jsd_beta: float = field( + default=0.0, + metadata={ + "help": "Mix between forward and reverse KL. 0.0 = forward KL " + "(teacher → student, mode-covering, default). 1.0 = reverse KL " + "(mode-seeking, used in on-policy GKD). 0.5 = symmetric JSD." + }, + ) + + @dataclass class ContextConfig: """Prefix context and target-mask configuration for context baking. diff --git a/src/bakery/kl.py b/src/bakery/kl.py index 6e3d64d..809e65b 100644 --- a/src/bakery/kl.py +++ b/src/bakery/kl.py @@ -46,6 +46,111 @@ def compute_kl_divergence( return masked_kl.sum() / num_tokens if num_tokens > 0 else masked_kl.sum() +def topk_forward_kl( + student_logits: torch.Tensor, + teacher_topk_indices: torch.Tensor, + teacher_topk_logprobs: torch.Tensor, + mask: torch.Tensor, + temperature: float = 1.0, + per_sample: bool = False, +) -> torch.Tensor: + """Forward KL over the teacher's top-k support. + + Works for both dense (K = vocab) and sparse (K = 20) teachers. The teacher's + top-k logprobs are assumed to be renormalized over those K (so exp(t).sum=1 + along the last dim) — this is what `topk_from_logits` produces and what we + require API backends to deliver. + + KL is computed as: sum_k p_teacher * (log p_teacher - log p_student_at_topk), + where p_student_at_topk is the student's softmax restricted (via gather) to + the teacher's top-k indices and re-normalized over those K so the comparison + is on the same support. + + Args: + student_logits: (B, T, V) student logits over full vocab. + teacher_topk_indices: (B, T, K) vocab ids of teacher's top-k. + teacher_topk_logprobs: (B, T, K) renormalized log-probs over those K. + mask: (B, T) target mask (1 where loss applies). + temperature: softmax temperature applied to student logits. + per_sample: if True, return [B] averaged per-sample. + + Returns: + Scalar KL, or [B] if per_sample=True. + """ + student_logprobs_full = F.log_softmax(student_logits / temperature, dim=-1) + student_logprobs_at_topk = torch.gather( + student_logprobs_full, dim=-1, index=teacher_topk_indices + ) # (B, T, K) + # Renormalize student over the same top-k support so we compare apples to + # apples (matches the teacher's renormalization). + student_logprobs_renorm = ( + student_logprobs_at_topk + - student_logprobs_at_topk.logsumexp(dim=-1, keepdim=True) + ) + + p_teacher = teacher_topk_logprobs.exp() + kl_per_pos = (p_teacher * (teacher_topk_logprobs - student_logprobs_renorm)).sum( + dim=-1 + ) # (B, T) + + mask_f = mask.float() + masked_kl = kl_per_pos * mask_f + if per_sample: + num_tokens = mask_f.sum(dim=-1).clamp(min=1.0) + return masked_kl.sum(dim=-1) / num_tokens + num_tokens = mask_f.sum() + return masked_kl.sum() / num_tokens if num_tokens > 0 else masked_kl.sum() + + +def topk_jsd( + student_logits: torch.Tensor, + teacher_topk_indices: torch.Tensor, + teacher_topk_logprobs: torch.Tensor, + mask: torch.Tensor, + beta: float = 0.5, + temperature: float = 1.0, + per_sample: bool = False, +) -> torch.Tensor: + """Mixed forward/reverse KL on the teacher's top-k support. + + loss = (1 - β) · D_KL(P_teacher || P_student) + + β · D_KL(P_student || P_teacher) + + β = 0 → pure forward KL (teacher → student); mode-covering; default bakery + β = 1 → pure reverse KL (student → teacher); mode-seeking; on-policy GKD + β = 0.5 → symmetric (JSD-style) average + + This is the same convention TRL's `GKDTrainer` uses for its `beta` knob — + monotonic in β, β=0 reduces exactly to forward KL, β=1 to reverse KL. + + Both KL terms are computed on the teacher's top-k support, with student and + teacher each renormalized over those K so the comparison is on the same + support. Computed in log space. + """ + if beta < 0.0 or beta > 1.0: + raise ValueError(f"beta must be in [0, 1], got {beta}") + + s_lp_full = F.log_softmax(student_logits / temperature, dim=-1) + s_lp_topk = torch.gather(s_lp_full, dim=-1, index=teacher_topk_indices) + s_lp_topk = s_lp_topk - s_lp_topk.logsumexp(dim=-1, keepdim=True) + + t_lp = teacher_topk_logprobs # already renormalized over top-k + + # Forward KL: sum p_t · (log p_t - log p_s) + kl_forward = (t_lp.exp() * (t_lp - s_lp_topk)).sum(dim=-1) # (B, T) + # Reverse KL: sum p_s · (log p_s - log p_t) + kl_reverse = (s_lp_topk.exp() * (s_lp_topk - t_lp)).sum(dim=-1) + mixed = (1.0 - beta) * kl_forward + beta * kl_reverse # (B, T) + + mask_f = mask.float() + masked = mixed * mask_f + if per_sample: + n = mask_f.sum(dim=-1).clamp(min=1.0) + return masked.sum(dim=-1) / n + n = mask_f.sum() + return masked.sum() / n if n > 0 else masked.sum() + + @contextmanager def disable_adapters(model): """Context manager to temporarily disable LoRA adapters.""" diff --git a/src/bakery/teachers/__init__.py b/src/bakery/teachers/__init__.py new file mode 100644 index 0000000..c7bf76f --- /dev/null +++ b/src/bakery/teachers/__init__.py @@ -0,0 +1,21 @@ +"""Teacher backends for GKD / context distillation. + +The trainer asks a `TeacherBackend` to score token sequences. Three backends: + +- `HFTeacher` — load a separate HuggingFace model in-process (fallback / dev path). +- `VLLMTeacher` — query a vLLM server via its OpenAI-compatible HTTP API (production). +- `OpenAIAPITeacher` — query any OpenAI-compatible endpoint that exposes `top_logprobs`. + +All backends return the same shape: `TopKLogprobs`, a sparse top-k view of the teacher's +distribution at each scored position. The trainer's KL loss consumes that uniformly, +regardless of whether K is the full vocab (HF dense) or K=20 (API sparse). + +A *student-side same-tokenizer assumption* applies: token ids passed to a teacher backend +are interpreted in the student's tokenizer. Mixing tokenizers (e.g. Llama→Qwen) is out of +scope for now. +""" + +from bakery.teachers.base import TeacherBackend, TopKLogprobs, make_teacher +from bakery.teachers.hf import HFTeacher + +__all__ = ["TeacherBackend", "TopKLogprobs", "HFTeacher", "make_teacher"] diff --git a/src/bakery/teachers/base.py b/src/bakery/teachers/base.py new file mode 100644 index 0000000..8ca83a6 --- /dev/null +++ b/src/bakery/teachers/base.py @@ -0,0 +1,130 @@ +"""TeacherBackend abstract interface + TopKLogprobs sparse view.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from dataclasses import dataclass +from typing import List, Optional + +import torch + + +@dataclass +class TopKLogprobs: + """Sparse top-k view of a teacher's distribution over tokens. + + `indices[b, t, k]` is a vocab id; `values[b, t, k]` is its logprob under the + teacher *renormalized over the top-k support* (so `exp(values).sum(-1) ≈ 1`). + Renormalizing on the teacher side keeps the KL math identical between dense + and sparse teachers — the student is scored only on those K positions. + + `attention_mask[b, t]` is 1 at positions the teacher actually scored. + """ + + indices: torch.Tensor # (B, T, K) long + values: torch.Tensor # (B, T, K) float (log-probabilities, renormalized over K) + attention_mask: torch.Tensor # (B, T) long/bool + + @property + def is_dense(self) -> bool: + """True if K covers the full vocab (no truncation).""" + return ( + self.indices.shape[-1] == self.values.shape[-1] + and self.indices.shape[-1] >= 32000 + ) # heuristic; cheap enough not to matter + + +class TeacherBackend(ABC): + """A model that scores token sequences and returns top-k logprobs. + + Backends MAY also implement `generate` for on-the-fly trajectory sampling, + but the trainer's loss path only needs `score`. + """ + + @abstractmethod + def score( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + top_k: int, + ) -> TopKLogprobs: + """Score a batch of student-tokenized sequences. + + Args: + input_ids: (B, T) student-tokenizer ids. + attention_mask: (B, T) 1=real, 0=padding. + top_k: how many tokens to keep per position. Backends that can only + provide ≤K (e.g. an API with `top_logprobs=20`) should return + whatever they have and the trainer will reduce K to match. + + Returns: + TopKLogprobs with the teacher's top-k at each position. Values are + renormalized over those K so exp(values).sum(-1) ≈ 1. + """ + + def generate(self, messages: List[dict], max_new_tokens: int = 256) -> str: + """Sample a completion conditioned on `messages`. Optional. + + Default: NotImplemented. Backends that support generation override this. + """ + raise NotImplementedError( + f"{type(self).__name__} does not implement generate()" + ) + + @property + def name(self) -> str: + return type(self).__name__ + + +def topk_from_logits( + logits: torch.Tensor, top_k: int +) -> tuple[torch.Tensor, torch.Tensor]: + """Compute top-k indices and renormalized logprobs from dense logits. + + Helper for backends that have full logits available (HF, vLLM with raw output). + Returns `(indices, values)` where values are renormalized over the top-k. + """ + k = min(top_k, logits.shape[-1]) + top_values, top_indices = torch.topk(logits, k=k, dim=-1) # (..., K) + # Convert raw logit values to logprobs renormalized over the top-k support. + top_logprobs = top_values - top_values.logsumexp(dim=-1, keepdim=True) + return top_indices, top_logprobs + + +def make_teacher(teacher_config) -> Optional[TeacherBackend]: + """Build a TeacherBackend from a TeacherConfig, or None for local-toggle mode.""" + backend = (teacher_config.teacher_backend or "local-toggle").lower() + if backend in ("local-toggle", "local", "self", "none", ""): + return None + if backend == "hf": + from bakery.teachers.hf import HFTeacher + + return HFTeacher( + model_name_or_path=teacher_config.teacher_model_name_or_path, + torch_dtype=teacher_config.teacher_torch_dtype, + device=teacher_config.teacher_device, + trust_remote_code=teacher_config.teacher_trust_remote_code, + attn_implementation=teacher_config.teacher_attn_implementation, + ) + if backend == "vllm": + from bakery.teachers.vllm import VLLMTeacher + + return VLLMTeacher( + api_base=teacher_config.teacher_api_base, + api_key=teacher_config.teacher_api_key, + model=teacher_config.teacher_api_model + or teacher_config.teacher_model_name_or_path, + ) + if backend in ("openai", "openai-compat", "openai_compat"): + from bakery.teachers.openai_compat import OpenAIAPITeacher + + return OpenAIAPITeacher( + api_base=teacher_config.teacher_api_base, + api_key=teacher_config.teacher_api_key, + model=teacher_config.teacher_api_model + or teacher_config.teacher_model_name_or_path, + ) + raise ValueError( + f"Unknown teacher backend: {backend!r}. " + "Expected one of: local-toggle, hf, vllm, openai." + ) diff --git a/src/bakery/teachers/hf.py b/src/bakery/teachers/hf.py new file mode 100644 index 0000000..f63ff53 --- /dev/null +++ b/src/bakery/teachers/hf.py @@ -0,0 +1,124 @@ +"""HuggingFace in-process teacher backend. + +Loads a separate Transformers model and scores student-tokenized sequences. +Same-tokenizer-as-student assumption — caller is responsible for that. +""" + +from __future__ import annotations + +import logging +from typing import List, Optional + +import torch +from transformers import AutoModelForCausalLM, AutoTokenizer, GenerationConfig + +from bakery.teachers.base import TeacherBackend, TopKLogprobs, topk_from_logits + +logger = logging.getLogger(__name__) + + +_DTYPE_MAP = { + "float32": torch.float32, + "fp32": torch.float32, + "float16": torch.float16, + "fp16": torch.float16, + "bfloat16": torch.bfloat16, + "bf16": torch.bfloat16, +} + + +class HFTeacher(TeacherBackend): + """Score and generate using a HuggingFace Transformers model in-process. + + Designed to share GPU memory with the student where possible (the student is + typically a smaller model). For larger teachers, consider running them on a + separate device or via vLLM. + """ + + def __init__( + self, + model_name_or_path: str, + torch_dtype: str = "bfloat16", + device: Optional[str] = None, + trust_remote_code: bool = False, + attn_implementation: Optional[str] = None, + ): + if not model_name_or_path: + raise ValueError("HFTeacher requires model_name_or_path") + + self.model_name_or_path = model_name_or_path + dtype = _DTYPE_MAP.get(torch_dtype, torch.bfloat16) + + load_kwargs = dict( + dtype=dtype, + trust_remote_code=trust_remote_code, + ) + if attn_implementation: + load_kwargs["attn_implementation"] = attn_implementation + if device: + load_kwargs["device_map"] = device + + logger.info("Loading HF teacher: %s (dtype=%s)", model_name_or_path, dtype) + self.model = AutoModelForCausalLM.from_pretrained( + model_name_or_path, **load_kwargs + ) + self.model.eval() + for p in self.model.parameters(): + p.requires_grad_(False) + + self.tokenizer = AutoTokenizer.from_pretrained( + model_name_or_path, trust_remote_code=trust_remote_code + ) + if self.tokenizer.pad_token_id is None: + self.tokenizer.pad_token = self.tokenizer.eos_token + + self._gen_config = GenerationConfig( + max_new_tokens=256, + do_sample=True, + top_p=0.9, + temperature=0.8, + pad_token_id=self.tokenizer.pad_token_id, + ) + + @torch.no_grad() + def score( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + top_k: int, + ) -> TopKLogprobs: + device = next(self.model.parameters()).device + input_ids = input_ids.to(device) + attention_mask = attention_mask.to(device) + + fwd = dict(input_ids=input_ids, attention_mask=attention_mask) + mtype = getattr(self.model.config, "model_type", None) + if mtype in ("gemma3", "gemma4", "gemma4_text"): + fwd["token_type_ids"] = torch.zeros_like(input_ids) + if mtype in ("gemma4", "gemma4_text"): + fwd["mm_token_type_ids"] = torch.zeros_like(input_ids) + + logits = self.model(**fwd).logits # (B, T, V) + indices, values = topk_from_logits(logits, top_k) + return TopKLogprobs( + indices=indices, + values=values, + attention_mask=attention_mask, + ) + + @torch.no_grad() + def generate(self, messages: List[dict], max_new_tokens: int = 256) -> str: + prompt = self.tokenizer.apply_chat_template( + messages, tokenize=False, add_generation_prompt=True + ) + inputs = self.tokenizer(prompt, return_tensors="pt", add_special_tokens=False) + device = next(self.model.parameters()).device + inputs = {k: v.to(device) for k, v in inputs.items()} + gen_config = GenerationConfig( + **{**self._gen_config.to_dict(), "max_new_tokens": max_new_tokens} + ) + out = self.model.generate(**inputs, generation_config=gen_config) + text = self.tokenizer.decode( + out[0][inputs["input_ids"].shape[1] :], skip_special_tokens=True + ) + return text.strip() diff --git a/src/bakery/teachers/openai_compat.py b/src/bakery/teachers/openai_compat.py new file mode 100644 index 0000000..3691939 --- /dev/null +++ b/src/bakery/teachers/openai_compat.py @@ -0,0 +1,207 @@ +"""OpenAI-compatible HTTP teacher backend. + +Targets any server implementing `/v1/completions` with `echo=True, +max_tokens=0, logprobs=K`. The canonical implementation is vLLM, but the +same plumbing works against Together, Fireworks, etc. as long as they expose +those flags. + +Same-tokenizer assumption: token ids passed in by the trainer are interpreted +in the *student's* tokenizer. The teacher server must use a compatible tokenizer +(in practice this means same model family — Gemma family, Qwen family, Llama +family, etc.). The trainer registers the student tokenizer via +`set_student_tokenizer` once at construction time. + +Cross-tokenizer scoring (Llama→Qwen) is out of scope and will raise on +re-encode mismatch. +""" + +from __future__ import annotations + +import logging +import os +from typing import List, Optional + +import torch + +from bakery.teachers.base import TeacherBackend, TopKLogprobs + +logger = logging.getLogger(__name__) + + +class OpenAIAPITeacher(TeacherBackend): + """OpenAI-compatible chat/completions teacher with sparse top-k scoring.""" + + def __init__( + self, + api_base: str, + api_key: Optional[str] = None, + model: Optional[str] = None, + echo_supported: bool = True, + timeout: float = 60.0, + ): + try: + import httpx # noqa: F401 + except ImportError as e: + raise ImportError( + "OpenAIAPITeacher requires `httpx`. Install with: pip install httpx" + ) from e + + if not api_base: + raise ValueError("OpenAIAPITeacher requires api_base") + if not model: + raise ValueError("OpenAIAPITeacher requires model") + + self.api_base = api_base.rstrip("/") + self.api_key = api_key or os.environ.get("OPENAI_API_KEY") or "" + self.model = model + self.echo_supported = echo_supported + self.timeout = timeout + self._student_tokenizer = None + + def set_student_tokenizer(self, tokenizer) -> None: + """Register the student tokenizer (used to re-encode API token strings).""" + self._student_tokenizer = tokenizer + + def _client(self): + import httpx + + headers = {"Content-Type": "application/json"} + if self.api_key: + headers["Authorization"] = f"Bearer {self.api_key}" + # trust_env=False so HTTP_PROXY / HTTPS_PROXY in the environment do + # not hijack the request. On HPC nodes (e.g. Sigma2 Olivia) a transparent + # Squid proxy will otherwise rewrite localhost POSTs to /v1/completions + # into a proxy error page, silently breaking GKD when the teacher is + # served on the same node. The teacher endpoint is always explicitly + # configured via api_base, so respecting proxy env vars buys nothing. + return httpx.Client( + base_url=self.api_base, + headers=headers, + timeout=self.timeout, + trust_env=False, + ) + + def score( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + top_k: int, + ) -> TopKLogprobs: + """Score sequences via /v1/completions with echo=True, max_tokens=0. + + Per-row: send the student's input_ids (vLLM accepts a list of ints as + `prompt`), echo back, request top-k logprobs at every position. Re-encode + each returned top-k token *string* into a student vocab id. + """ + if not self.echo_supported: + raise NotImplementedError( + "Server doesn't advertise echo support; sparse scoring unavailable." + ) + if self._student_tokenizer is None: + raise RuntimeError( + "OpenAIAPITeacher.score requires a student tokenizer — call " + "set_student_tokenizer(...) before training." + ) + + B, T = input_ids.shape + vocab_size = len(self._student_tokenizer) + indices_out = torch.zeros(B, T, top_k, dtype=torch.long) + values_out = torch.full((B, T, top_k), -1e9, dtype=torch.float32) + mask_out = attention_mask.clone() + + client = self._client() + try: + for b in range(B): + row_attn = attention_mask[b] + # Strip left-padding before sending — vLLM doesn't need our padding. + real_start = int((row_attn > 0).nonzero(as_tuple=False)[0].item()) + real_ids = input_ids[b, real_start:].tolist() + resp = client.post( + "/completions", + json={ + "model": self.model, + "prompt": real_ids, + "max_tokens": 0, + "echo": True, + "logprobs": top_k, + "temperature": 0, + }, + ) + resp.raise_for_status() + choice = resp.json()["choices"][0] + lp_block = choice.get("logprobs") or {} + top_lps = lp_block.get("top_logprobs") or [] + # vLLM returns top_logprobs as a list (one per echoed token). + # Position 0 is the very first token, which has no prior context + # and is conventionally null. Fill from position 1 onward. + for t_local, pos_top in enumerate(top_lps): + if pos_top is None: + continue + t_full = real_start + t_local + if t_full >= T: + break + # pos_top: dict-like {token_string: logprob}. + items = list(pos_top.items()) + if not items: + continue + # Sort by logprob desc to be safe — most servers already do this. + items.sort(key=lambda kv: -kv[1]) + raw_logprobs = [] + raw_indices = [] + for tok_str, lp in items[:top_k]: + ids = self._student_tokenizer.encode( + tok_str, add_special_tokens=False + ) + if not ids: + continue + # Pick the single id that represents this token. If the + # API token is a multi-piece string under the student + # tokenizer, we approximate by the first piece. + sid = ids[0] + if sid >= vocab_size: + continue + raw_indices.append(sid) + raw_logprobs.append(lp) + if not raw_indices: + continue + n = min(len(raw_indices), top_k) + idx_t = torch.tensor(raw_indices[:n], dtype=torch.long) + lp_t = torch.tensor(raw_logprobs[:n], dtype=torch.float32) + # Renormalize over the K we actually captured so the trainer's + # `exp(values).sum ≈ 1` invariant holds. + lp_t = lp_t - torch.logsumexp(lp_t, dim=0) + indices_out[b, t_full, :n] = idx_t + values_out[b, t_full, :n] = lp_t + if n < top_k: + # Pad remaining slots by reusing the first index with + # very negative logprob — keeps tensors well-formed + # and the renormalization above keeps mass on the real K. + indices_out[b, t_full, n:] = idx_t[0] + values_out[b, t_full, n:] = -1e9 + finally: + client.close() + + return TopKLogprobs( + indices=indices_out, + values=values_out, + attention_mask=mask_out, + ) + + def generate(self, messages: List[dict], max_new_tokens: int = 256) -> str: + client = self._client() + try: + resp = client.post( + "/chat/completions", + json={ + "model": self.model, + "messages": messages, + "max_tokens": max_new_tokens, + "temperature": 0.8, + "top_p": 0.9, + }, + ) + resp.raise_for_status() + data = resp.json() + return data["choices"][0]["message"]["content"] + finally: + client.close() diff --git a/src/bakery/teachers/vllm.py b/src/bakery/teachers/vllm.py new file mode 100644 index 0000000..510d71c --- /dev/null +++ b/src/bakery/teachers/vllm.py @@ -0,0 +1,68 @@ +"""vLLM teacher backend. + +Two flavors share this file: + + 1. vLLM-as-OpenAI-server: hit `/v1/completions` with `echo=True, max_tokens=0, + logprobs=K` to score arbitrary student-tokenized sequences. This is the + production path that mirrors TRL's pattern. + + 2. vLLM-as-LLM (in-process): when the user prefers, instantiate `vllm.LLM` + directly. Returns full dense logprobs without HTTP overhead. + +For now the OpenAI-compat variant inherits from OpenAIAPITeacher and the +in-process variant is a thin wrapper. Full implementations land alongside the +TRL-style vLLM integration in a follow-up commit. +""" + +from __future__ import annotations + +from typing import List, Optional + +import torch + +from bakery.teachers.base import TeacherBackend, TopKLogprobs +from bakery.teachers.openai_compat import OpenAIAPITeacher + + +class VLLMTeacher(OpenAIAPITeacher): + """vLLM-as-OpenAI-server teacher. + + Inherits HTTP plumbing from OpenAIAPITeacher; assumes echo+logprobs are + supported (vLLM's default). + """ + + def __init__( + self, + api_base: str, + api_key: Optional[str] = None, + model: Optional[str] = None, + timeout: float = 60.0, + ): + super().__init__( + api_base=api_base, + api_key=api_key, + model=model, + echo_supported=True, + timeout=timeout, + ) + + +class VLLMInProcessTeacher(TeacherBackend): + """In-process vLLM teacher — to be filled in alongside TRL-style integration.""" + + def __init__(self, *args, **kwargs): + raise NotImplementedError( + "VLLMInProcessTeacher: not yet implemented. Use VLLMTeacher (HTTP) or " + "HFTeacher for now." + ) + + def score( + self, + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + top_k: int, + ) -> TopKLogprobs: + raise NotImplementedError + + def generate(self, messages: List[dict], max_new_tokens: int = 256) -> str: + raise NotImplementedError diff --git a/src/bakery/trainer.py b/src/bakery/trainer.py index d143159..dfa060f 100644 --- a/src/bakery/trainer.py +++ b/src/bakery/trainer.py @@ -28,8 +28,14 @@ ) from transformers.trainer_utils import EvalPrediction -from bakery.kl import compute_kl_divergence, disable_adapters +from bakery.kl import ( + compute_kl_divergence, + disable_adapters, + topk_forward_kl, + topk_jsd, +) from bakery.masking import build_target_mask +from bakery.teachers.base import TeacherBackend logger = logging.getLogger(__name__) @@ -50,6 +56,10 @@ def __init__( model: PreTrainedModel | str | None = None, args=None, context_config=None, + teacher_backend: Optional[TeacherBackend] = None, + teacher_top_k: int = 64, + gkd_on_policy_fraction: float = 0.0, + gkd_jsd_beta: float = 0.0, train_dataset: Dataset | None = None, eval_dataset: Dataset | None = None, processing_class: PreTrainedTokenizerBase | None = None, @@ -65,6 +75,16 @@ def __init__( self.trajectory_length = args.trajectory_length self.sampling_temperature = args.sampling_temperature self.kl_temperature = args.temperature + self.teacher_backend = teacher_backend + self.teacher_top_k = teacher_top_k + if not (0.0 <= gkd_on_policy_fraction <= 1.0): + raise ValueError( + f"gkd_on_policy_fraction must be in [0, 1], got {gkd_on_policy_fraction}" + ) + if not (0.0 <= gkd_jsd_beta <= 1.0): + raise ValueError(f"gkd_jsd_beta must be in [0, 1], got {gkd_jsd_beta}") + self.gkd_on_policy_fraction = gkd_on_policy_fraction + self.gkd_jsd_beta = gkd_jsd_beta # Back-compat: if no ContextConfig is passed but args has system_prompt, # auto-desugar into a minimal ContextConfig. Keeps direct-instantiation @@ -116,6 +136,13 @@ def __init__( pad_token_id=self.processing_class.pad_token_id, ) + # API teachers re-encode their returned token strings against the + # student tokenizer; register it now so they can. + if self.teacher_backend is not None and hasattr( + self.teacher_backend, "set_student_tokenizer" + ): + self.teacher_backend.set_student_tokenizer(self.processing_class) + # -- Tokenization -- def _tokenize(self, text, **kwargs) -> dict: @@ -307,57 +334,79 @@ def _inject_gemma_token_types(model, fwd: dict, input_ids: torch.Tensor) -> None # -- KL loss from aligned logits + masks -- - def _kl_from_logits( + def _align_and_slice( self, - teacher_logits: torch.Tensor, + i: int, + teacher_data: torch.Tensor, student_logits: torch.Tensor, teacher_mask_padded: torch.BoolTensor, student_mask_padded: torch.BoolTensor, - ) -> Optional[torch.Tensor]: - """Compute mean per-example KL over aligned target tokens. + ): + """Find the aligned target region for sample i and return slices. - Logits at position t predict token t+1, so we shift by 1. Teacher and - student may have different prefix lengths, so we align on the trainable - region by finding the first target token in each and taking the common - length. + Returns (t_slice, s_logits_slice, combined_mask) or None if no overlap. + `teacher_data` is the teacher tensor we want to align — full logits in + local-toggle mode, top-k indices/values in external-teacher mode + (slicing dim is the same: position). """ - losses = [] - B = teacher_logits.shape[0] device = student_logits.device + t_mask = teacher_mask_padded[i].to(device) + s_mask = student_mask_padded[i].to(device) - for i in range(B): - t_mask = teacher_mask_padded[i].to(device) - s_mask = student_mask_padded[i].to(device) + t_pos = t_mask.nonzero(as_tuple=False).squeeze(-1) + s_pos = s_mask.nonzero(as_tuple=False).squeeze(-1) + if t_pos.numel() == 0 or s_pos.numel() == 0: + return None - t_target_positions = t_mask.nonzero(as_tuple=False).squeeze(-1) - s_target_positions = s_mask.nonzero(as_tuple=False).squeeze(-1) - if t_target_positions.numel() == 0 or s_target_positions.numel() == 0: - continue + t_start = int(t_pos[0].item()) + s_start = int(s_pos[0].item()) - t_start = int(t_target_positions[0].item()) - s_start = int(s_target_positions[0].item()) + # Logits predict next token → shift by 1. + t_slice = teacher_data[i : i + 1, t_start - 1 : -1] + s_slice = student_logits[i : i + 1, s_start - 1 : -1, :] + t_tail_mask = t_mask[t_start:].float().unsqueeze(0) + s_tail_mask = s_mask[s_start:].float().unsqueeze(0) - # Logits predict next token → shift by 1. - t_logits = teacher_logits[i : i + 1, t_start - 1 : -1, :] - s_logits = student_logits[i : i + 1, s_start - 1 : -1, :] - t_tail_mask = t_mask[t_start:].float().unsqueeze(0) - s_tail_mask = s_mask[s_start:].float().unsqueeze(0) + min_len = min( + t_slice.shape[1], + s_slice.shape[1], + t_tail_mask.shape[1], + s_tail_mask.shape[1], + ) + if min_len == 0: + return None - min_len = min( - t_logits.shape[1], - s_logits.shape[1], - t_tail_mask.shape[1], - s_tail_mask.shape[1], - ) - if min_len == 0: - continue + combined_mask = t_tail_mask[:, :min_len] * s_tail_mask[:, :min_len] + if combined_mask.sum() == 0: + return None + return ( + t_slice[:, :min_len], + s_slice[:, :min_len, :], + combined_mask, + ) - t_logits = t_logits[:, :min_len, :] - s_logits = s_logits[:, :min_len, :] - combined_mask = t_tail_mask[:, :min_len] * s_tail_mask[:, :min_len] - if combined_mask.sum() == 0: - continue + def _kl_from_logits( + self, + teacher_logits: torch.Tensor, + student_logits: torch.Tensor, + teacher_mask_padded: torch.BoolTensor, + student_mask_padded: torch.BoolTensor, + ) -> Optional[torch.Tensor]: + """Dense KL over the full vocab (local-toggle mode).""" + losses = [] + B = teacher_logits.shape[0] + for i in range(B): + aligned = self._align_and_slice( + i, + teacher_logits, + student_logits, + teacher_mask_padded, + student_mask_padded, + ) + if aligned is None: + continue + t_logits, s_logits, combined_mask = aligned loss = compute_kl_divergence( t_logits.detach() if t_logits.requires_grad else t_logits, s_logits, @@ -370,13 +419,122 @@ def _kl_from_logits( return None return torch.stack(losses).mean() + def _kl_from_topk( + self, + teacher_indices: torch.Tensor, + teacher_logprobs: torch.Tensor, + student_logits: torch.Tensor, + teacher_mask_padded: torch.BoolTensor, + student_mask_padded: torch.BoolTensor, + ) -> Optional[torch.Tensor]: + """Top-k KL — external-teacher mode. + + Same alignment + shift logic as `_kl_from_logits`, but teacher data is + sparse (per-position top-k indices + logprobs renormalized over those K). + """ + losses = [] + B = teacher_indices.shape[0] + device = student_logits.device + + for i in range(B): + aligned_idx = self._align_and_slice( + i, + teacher_indices, + student_logits, + teacher_mask_padded, + student_mask_padded, + ) + if aligned_idx is None: + continue + t_idx, s_logits, combined_mask = aligned_idx + # Align values to the same window. + min_len = t_idx.shape[1] + t_mask = teacher_mask_padded[i].to(device) + t_pos = t_mask.nonzero(as_tuple=False).squeeze(-1) + if t_pos.numel() == 0: + continue + t_start = int(t_pos[0].item()) + t_vals = teacher_logprobs[i : i + 1, t_start - 1 : -1][:, :min_len] + t_idx = t_idx.to(device) + t_vals = t_vals.to(device) + + if self.gkd_jsd_beta > 0.0: + loss = topk_jsd( + student_logits=s_logits, + teacher_topk_indices=t_idx, + teacher_topk_logprobs=t_vals, + mask=combined_mask, + beta=self.gkd_jsd_beta, + temperature=self.kl_temperature, + ) + else: + loss = topk_forward_kl( + student_logits=s_logits, + teacher_topk_indices=t_idx, + teacher_topk_logprobs=t_vals, + mask=combined_mask, + temperature=self.kl_temperature, + ) + losses.append(loss) + + if not losses: + return None + return torch.stack(losses).mean() + # -- Trajectory generation -- + def _sample_from_student(self, user_message: str) -> str: + """On-policy GKD: sample a response from the student (adapters ENABLED). + + The student sees what it would see at inference time — its own trimmed + view (last `student_retained_turns` of the prefix), not the teacher's + full context. This is what gives on-policy GKD its mode-seeking pull. + """ + prefix = self._student_prefix(self.prefix_messages) + student_messages = list(prefix) + [{"role": "user", "content": user_message}] + prompt = self.processing_class.apply_chat_template( + student_messages, tokenize=False, add_generation_prompt=True + ) + inputs = self._tokenize(prompt, return_tensors="pt").to(self.model.device) + + was_training = self.model.training + self.model.eval() + with torch.no_grad(): + outputs = self.model.generate( + **inputs, generation_config=self.generation_config + ) + if was_training: + self.model.train() + response = self.processing_class.decode( + outputs[0][inputs["input_ids"].shape[1] :], skip_special_tokens=True + ) + return response.strip() + def _generate_trajectory(self, user_message: str) -> str: - """Generate a response from the teacher (adapters disabled, full prefix visible).""" + """Generate a response. + + Routing: + - With probability `gkd_on_policy_fraction`, sample from the student + itself (on-policy GKD). The teacher will then score those samples + during the loss step. + - Otherwise sample from the teacher: external `teacher_backend.generate` + if configured, else the same model with adapters disabled. + """ + if self.gkd_on_policy_fraction > 0.0: + import random + + if random.random() < self.gkd_on_policy_fraction: + return self._sample_from_student(user_message) + teacher_messages = list(self.prefix_messages) + [ {"role": "user", "content": user_message} ] + + if self.teacher_backend is not None: + return self.teacher_backend.generate( + teacher_messages, max_new_tokens=self.trajectory_length + ) + prompt = self.processing_class.apply_chat_template( teacher_messages, tokenize=False, add_generation_prompt=True ) @@ -404,27 +562,59 @@ def _generate_trajectory(self, user_message: str) -> str: def _zero_loss(self) -> torch.Tensor: return torch.tensor(0.0, device=self.args.device, requires_grad=True) + def _teacher_forward(self, model, batch): + """Run the teacher forward and return either (full_logits, None) or + (topk_indices, topk_logprobs) — caller decides which KL path to use. + + External teacher: returns top-k. Local-toggle: returns full logits. + """ + if self.teacher_backend is not None: + tk = self.teacher_backend.score( + input_ids=batch["teacher_fwd"]["input_ids"], + attention_mask=batch["teacher_fwd"]["attention_mask"], + top_k=self.teacher_top_k, + ) + return ("topk", tk.indices, tk.values) + with torch.no_grad(): + with disable_adapters(model): + out = model(**batch["teacher_fwd"]) + return ("dense", out.logits, None) + + def _kl_dispatch(self, teacher_payload, student_logits, batch): + mode = teacher_payload[0] + if mode == "dense": + return self._kl_from_logits( + teacher_payload[1], + student_logits, + batch["teacher_mask_padded"], + batch["student_mask_padded"], + ) + return self._kl_from_topk( + teacher_indices=teacher_payload[1], + teacher_logprobs=teacher_payload[2], + student_logits=student_logits, + teacher_mask_padded=batch["teacher_mask_padded"], + student_mask_padded=batch["student_mask_padded"], + ) + def compute_loss( self, model, inputs, return_outputs=False, num_items_in_batch=None ): - """Compute KL divergence loss with batched forward passes.""" + """Compute KL divergence loss with batched forward passes. + + Dispatches between local-toggle (dense KL on full logits) and external- + teacher (top-k KL on sparse logprobs) based on `teacher_backend`. + """ batch = self._build_batch(inputs, model) if batch is None: logger.warning("Empty batch after building — returning zero loss") zero = self._zero_loss() return (zero, None) if return_outputs else zero - with torch.no_grad(): - with disable_adapters(model): - teacher_outputs = model(**batch["teacher_fwd"]) + teacher_payload = self._teacher_forward(model, batch) student_outputs = model(**batch["student_fwd"]) - loss = self._kl_from_logits( - teacher_outputs.logits, - student_outputs.logits, - batch["teacher_mask_padded"], - batch["student_mask_padded"], - ) + loss = self._kl_dispatch(teacher_payload, student_outputs.logits, batch) if loss is None: logger.warning("No aligned logit pairs after slicing — returning zero loss") zero = self._zero_loss() @@ -450,18 +640,38 @@ def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None) base = model.module if hasattr(model, "module") else model fwd_fn = type(base).forward - with torch.no_grad(): - with disable_adapters(base): - teacher_logits = fwd_fn(base, **batch["teacher_fwd"]).logits.cpu() + if self.teacher_backend is not None: + tk = self.teacher_backend.score( + input_ids=batch["teacher_fwd"]["input_ids"], + attention_mask=batch["teacher_fwd"]["attention_mask"], + top_k=self.teacher_top_k, + ) + teacher_payload = ("topk", tk.indices.cpu(), tk.values.cpu()) torch.cuda.empty_cache() - student_outputs = fwd_fn(base, **batch["student_fwd"]) + with torch.no_grad(): + student_outputs = fwd_fn(base, **batch["student_fwd"]) + else: + with torch.no_grad(): + with disable_adapters(base): + teacher_logits = fwd_fn(base, **batch["teacher_fwd"]).logits.cpu() + torch.cuda.empty_cache() + student_outputs = fwd_fn(base, **batch["student_fwd"]) + teacher_payload = ( + "dense", + teacher_logits.to(student_outputs.logits.device), + None, + ) - loss = self._kl_from_logits( - teacher_logits.to(student_outputs.logits.device), - student_outputs.logits, - batch["teacher_mask_padded"], - batch["student_mask_padded"], - ) + # Re-hydrate topk tensors onto student device for the KL step. + if teacher_payload[0] == "topk": + dev = student_outputs.logits.device + teacher_payload = ( + "topk", + teacher_payload[1].to(dev), + teacher_payload[2].to(dev), + ) + + loss = self._kl_dispatch(teacher_payload, student_outputs.logits, batch) if loss is None: return (self._zero_loss(), None, None) return (loss.detach(), None, None) diff --git a/tests/test_openai_teacher_score.py b/tests/test_openai_teacher_score.py new file mode 100644 index 0000000..2954b02 --- /dev/null +++ b/tests/test_openai_teacher_score.py @@ -0,0 +1,214 @@ +"""Tests for OpenAIAPITeacher.score against a mocked vLLM-style server. + +We swap httpx.Client out for a stub that returns hand-crafted logprobs +payloads, so this runs offline / in CI. +""" + +from unittest.mock import MagicMock, patch + +import pytest +import torch +from transformers import AutoTokenizer + +from bakery.teachers.openai_compat import OpenAIAPITeacher + + +def _mock_response(payload): + resp = MagicMock() + resp.raise_for_status = MagicMock() + resp.json = MagicMock(return_value=payload) + return resp + + +def _mock_client(payload): + client = MagicMock() + client.post = MagicMock(return_value=_mock_response(payload)) + client.close = MagicMock() + return client + + +@pytest.fixture(scope="module") +def gpt2_tokenizer(): + tok = AutoTokenizer.from_pretrained("gpt2") + tok.pad_token = tok.eos_token + return tok + + +# ---------- guardrails ---------- + + +def test_score_requires_student_tokenizer(): + teacher = OpenAIAPITeacher(api_base="http://x", model="m") + with pytest.raises(RuntimeError, match="student tokenizer"): + teacher.score( + input_ids=torch.tensor([[1, 2, 3]]), + attention_mask=torch.ones(1, 3), + top_k=4, + ) + + +def test_score_requires_echo_support(): + teacher = OpenAIAPITeacher(api_base="http://x", model="m", echo_supported=False) + with pytest.raises(NotImplementedError, match="echo"): + teacher.score( + input_ids=torch.tensor([[1]]), + attention_mask=torch.ones(1, 1), + top_k=4, + ) + + +# ---------- score happy path ---------- + + +def _payload_for(tokens_and_topks): + """Helper: build a vLLM-style payload from [(token_string, {tok: lp})].""" + return { + "choices": [ + { + "logprobs": { + "tokens": [t for t, _ in tokens_and_topks], + "top_logprobs": [tp for _, tp in tokens_and_topks], + } + } + ] + } + + +def test_score_returns_topklogprobs_shape(gpt2_tokenizer): + teacher = OpenAIAPITeacher(api_base="http://x", model="m") + teacher.set_student_tokenizer(gpt2_tokenizer) + + # Three echoed positions; position 0 has no top_logprobs (typical), 1+2 do. + payload = _payload_for( + [ + ("Hello", None), + (",", {",": -0.1, " world": -2.0, "!": -3.5}), + (" world", {" world": -0.2, "!": -1.5, ".": -3.0}), + ] + ) + with patch.object(teacher, "_client", return_value=_mock_client(payload)): + out = teacher.score( + input_ids=torch.tensor([[15496, 11, 995]]), + attention_mask=torch.ones(1, 3), + top_k=3, + ) + assert out.indices.shape == (1, 3, 3) + assert out.values.shape == (1, 3, 3) + + +def test_score_values_renormalized(gpt2_tokenizer): + """After re-encoding + renorm, exp(values) at scored positions sum ≈ 1.""" + teacher = OpenAIAPITeacher(api_base="http://x", model="m") + teacher.set_student_tokenizer(gpt2_tokenizer) + payload = _payload_for( + [ + (",", None), + ( + " world", + {" world": -0.2, "!": -1.5, ".": -3.0, "Hello": -4.0}, + ), + ] + ) + with patch.object(teacher, "_client", return_value=_mock_client(payload)): + out = teacher.score( + input_ids=torch.tensor([[11, 995]]), + attention_mask=torch.ones(1, 2), + top_k=4, + ) + # Position 1 should have real renormalized logprobs. + s = out.values[0, 1].exp().sum().item() + assert 0.99 <= s <= 1.01 + + +def test_score_skips_null_first_position(gpt2_tokenizer): + """Echo's first position usually has null top_logprobs — we skip it.""" + teacher = OpenAIAPITeacher(api_base="http://x", model="m") + teacher.set_student_tokenizer(gpt2_tokenizer) + payload = _payload_for( + [ + ("Hello", None), + (",", {",": -0.1, "!": -2.0}), + ] + ) + with patch.object(teacher, "_client", return_value=_mock_client(payload)): + out = teacher.score( + input_ids=torch.tensor([[15496, 11]]), + attention_mask=torch.ones(1, 2), + top_k=2, + ) + # Position 0 stays at default placeholder values (-1e9). + assert (out.values[0, 0] <= -1e8).all() + # Position 1 has real values. + assert out.values[0, 1].max().item() > -1e3 + + +def test_score_strips_left_padding(gpt2_tokenizer): + """Padded prefix tokens should not be sent to the server.""" + teacher = OpenAIAPITeacher(api_base="http://x", model="m") + teacher.set_student_tokenizer(gpt2_tokenizer) + payload = _payload_for([(",", None), (" world", {" world": -0.2, "!": -1.0})]) + captured = {} + + def fake_post(*args, **kwargs): + captured["json"] = kwargs.get("json") + return _mock_response(payload) + + client = MagicMock() + client.post = fake_post + client.close = MagicMock() + with patch.object(teacher, "_client", return_value=client): + out = teacher.score( + input_ids=torch.tensor([[0, 0, 11, 995]]), + attention_mask=torch.tensor([[0, 0, 1, 1]]), + top_k=2, + ) + # Only 2 real tokens should have been sent (left padding stripped). + assert captured["json"]["prompt"] == [11, 995] + # The output preserves the original shape, padding positions left untouched. + assert out.indices.shape == (1, 4, 2) + + +def test_score_empty_top_logprobs_position(gpt2_tokenizer): + """A position whose top_logprobs is {} should not blow up.""" + teacher = OpenAIAPITeacher(api_base="http://x", model="m") + teacher.set_student_tokenizer(gpt2_tokenizer) + payload = _payload_for( + [ + ("Hello", None), + (",", {}), + (" world", {" world": -0.2, "!": -1.0}), + ] + ) + with patch.object(teacher, "_client", return_value=_mock_client(payload)): + out = teacher.score( + input_ids=torch.tensor([[15496, 11, 995]]), + attention_mask=torch.ones(1, 3), + top_k=3, + ) + # No crash, and the empty position keeps placeholder values. + assert (out.values[0, 1] <= -1e8).all() + assert out.values[0, 2].max().item() > -1e3 + + +# ---------- generate ---------- + + +def test_generate_sends_messages(): + teacher = OpenAIAPITeacher(api_base="http://x", model="m") + captured = {} + + def fake_post(*args, **kwargs): + captured["json"] = kwargs.get("json") + return _mock_response({"choices": [{"message": {"content": "the answer"}}]}) + + client = MagicMock() + client.post = fake_post + client.close = MagicMock() + with patch.object(teacher, "_client", return_value=client): + out = teacher.generate( + [{"role": "user", "content": "What is the answer?"}], + max_new_tokens=8, + ) + assert out == "the answer" + assert captured["json"]["messages"][0]["content"] == "What is the answer?" + assert captured["json"]["max_tokens"] == 8 diff --git a/tests/test_teacher_backends.py b/tests/test_teacher_backends.py new file mode 100644 index 0000000..f7284c9 --- /dev/null +++ b/tests/test_teacher_backends.py @@ -0,0 +1,135 @@ +"""Tests for TeacherBackend implementations and the make_teacher factory.""" + +import pytest +import torch + +from bakery.config import TeacherConfig +from bakery.teachers import HFTeacher, TopKLogprobs, make_teacher + + +# ---------- TopKLogprobs / topk_from_logits ---------- + + +def test_topklogprobs_is_dense_heuristic(): + """is_dense triggers only when K is large (gpt2 vocab-ish).""" + small = TopKLogprobs( + indices=torch.zeros(1, 1, 8, dtype=torch.long), + values=torch.zeros(1, 1, 8), + attention_mask=torch.ones(1, 1), + ) + assert not small.is_dense + + big = TopKLogprobs( + indices=torch.zeros(1, 1, 50000, dtype=torch.long), + values=torch.zeros(1, 1, 50000), + attention_mask=torch.ones(1, 1), + ) + assert big.is_dense + + +# ---------- make_teacher dispatch ---------- + + +def test_make_teacher_local_toggle_returns_none(): + """local-toggle (the default) means: no external teacher.""" + cfg = TeacherConfig() + assert make_teacher(cfg) is None + + +def test_make_teacher_none_aliases(): + """Several aliases all resolve to local-toggle.""" + for alias in ("local-toggle", "local", "self", "none", ""): + cfg = TeacherConfig(teacher_backend=alias) + assert make_teacher(cfg) is None + + +def test_make_teacher_unknown_raises(): + cfg = TeacherConfig(teacher_backend="not-a-real-backend") + with pytest.raises(ValueError, match="Unknown teacher backend"): + make_teacher(cfg) + + +def test_make_teacher_hf_requires_model_name(): + cfg = TeacherConfig(teacher_backend="hf") + with pytest.raises(ValueError, match="model_name_or_path"): + make_teacher(cfg) + + +def test_make_teacher_openai_requires_model_and_base(): + cfg = TeacherConfig(teacher_backend="openai") + with pytest.raises((ValueError, ImportError)): + make_teacher(cfg) + + +# ---------- HFTeacher (uses tiny GPT-2) ---------- + + +@pytest.fixture(scope="module") +def hf_teacher(): + """A real HFTeacher backed by gpt2 — tiny enough to run on CPU.""" + return HFTeacher(model_name_or_path="gpt2", torch_dtype="float32") + + +def test_hf_teacher_score_returns_topklogprobs(hf_teacher): + """score() returns top-k indices, renormalized logprobs, attention mask.""" + input_ids = torch.tensor([[15496, 11, 995]]) # "Hello, world" + attn = torch.ones_like(input_ids) + out = hf_teacher.score(input_ids, attn, top_k=10) + assert isinstance(out, TopKLogprobs) + assert out.indices.shape == (1, 3, 10) + assert out.values.shape == (1, 3, 10) + # Renormalized over top-k → exp().sum ≈ 1. + sums = out.values.exp().sum(dim=-1) + assert torch.allclose(sums, torch.ones_like(sums), atol=1e-4) + + +def test_hf_teacher_score_respects_top_k(hf_teacher): + """top_k smaller than vocab → that's what we get.""" + input_ids = torch.tensor([[15496, 11]]) + attn = torch.ones_like(input_ids) + out = hf_teacher.score(input_ids, attn, top_k=5) + assert out.indices.shape[-1] == 5 + + +def test_hf_teacher_score_full_vocab(hf_teacher): + """Asking for K > vocab caps at vocab size.""" + input_ids = torch.tensor([[15496]]) + attn = torch.ones_like(input_ids) + out = hf_teacher.score(input_ids, attn, top_k=999_999) + # GPT-2 vocab is 50257. + assert out.indices.shape[-1] == 50257 + + +def test_hf_teacher_generate_returns_string(hf_teacher): + """generate() produces a string, even from a trivial chat message.""" + # gpt2 has no chat template by default; set a minimal one for this test. + hf_teacher.tokenizer.chat_template = ( + "{% for m in messages %}{{ m['role'] }}: {{ m['content'] }}\n{% endfor %}" + ) + out = hf_teacher.generate([{"role": "user", "content": "Hello"}], max_new_tokens=3) + assert isinstance(out, str) + + +def test_hf_teacher_score_no_grad(hf_teacher): + """Teacher params are frozen; scoring should not produce gradients.""" + input_ids = torch.tensor([[15496, 11]]) + attn = torch.ones_like(input_ids) + out = hf_teacher.score(input_ids, attn, top_k=4) + # Values are detached (no_grad path). + assert not out.values.requires_grad + + +def test_hf_teacher_score_padding_mask_propagated(hf_teacher): + """attention_mask is returned on the TopKLogprobs.""" + input_ids = torch.tensor([[15496, 11, 0, 0]]) + attn = torch.tensor([[1, 1, 0, 0]]) + out = hf_teacher.score(input_ids, attn, top_k=4) + assert out.attention_mask.tolist() == [[1, 1, 0, 0]] + + +# ---------- HFTeacher init guardrails ---------- + + +def test_hf_teacher_requires_model_name(): + with pytest.raises(ValueError, match="model_name_or_path"): + HFTeacher(model_name_or_path="") diff --git a/tests/test_topk_kl.py b/tests/test_topk_kl.py new file mode 100644 index 0000000..f03201c --- /dev/null +++ b/tests/test_topk_kl.py @@ -0,0 +1,203 @@ +"""Tests for top-k forward KL — the sparse-friendly distillation loss.""" + +import math + +import pytest +import torch + +from bakery.kl import compute_kl_divergence, topk_forward_kl, topk_jsd +from bakery.teachers.base import topk_from_logits + + +# ---------- topk_from_logits ---------- + + +def test_topk_from_logits_returns_renormalized_logprobs(): + logits = torch.randn(2, 3, 100) + idx, lp = topk_from_logits(logits, top_k=10) + assert idx.shape == (2, 3, 10) + assert lp.shape == (2, 3, 10) + # Renormalized: exp sums to 1 along last dim. + sums = lp.exp().sum(dim=-1) + assert torch.allclose(sums, torch.ones_like(sums), atol=1e-5) + + +def test_topk_from_logits_k_capped_at_vocab(): + logits = torch.randn(1, 1, 5) + idx, lp = topk_from_logits(logits, top_k=10) + assert idx.shape[-1] == 5 + + +def test_topk_from_logits_picks_actual_topk(): + logits = torch.zeros(1, 1, 5) + logits[0, 0, 2] = 10.0 + logits[0, 0, 4] = 5.0 + idx, lp = topk_from_logits(logits, top_k=2) + assert set(idx[0, 0].tolist()) == {2, 4} + + +# ---------- topk_forward_kl ---------- + + +def test_topk_kl_zero_when_distributions_match(): + """If teacher == student logits, top-k KL should be ~0.""" + torch.manual_seed(0) + logits = torch.randn(1, 4, 50, requires_grad=True) + idx, lp = topk_from_logits(logits.detach(), top_k=20) + mask = torch.ones(1, 4) + loss = topk_forward_kl( + student_logits=logits, + teacher_topk_indices=idx, + teacher_topk_logprobs=lp, + mask=mask, + ) + assert loss.item() < 1e-4 + + +def test_topk_kl_dense_matches_full_kl(): + """When K = vocab, top-k KL should equal full-vocab forward KL.""" + torch.manual_seed(0) + V = 32 + teacher_logits = torch.randn(2, 5, V) + student_logits = torch.randn(2, 5, V) + mask = torch.ones(2, 5) + + full = compute_kl_divergence(teacher_logits, student_logits, mask) + + idx, lp = topk_from_logits(teacher_logits, top_k=V) + sparse = topk_forward_kl( + student_logits=student_logits, + teacher_topk_indices=idx, + teacher_topk_logprobs=lp, + mask=mask, + ) + assert torch.allclose(full, sparse, atol=1e-4) + + +def test_topk_kl_per_sample_shape(): + torch.manual_seed(1) + logits = torch.randn(3, 4, 20) + idx, lp = topk_from_logits(logits, top_k=8) + mask = torch.ones(3, 4) + out = topk_forward_kl( + student_logits=torch.randn(3, 4, 20), + teacher_topk_indices=idx, + teacher_topk_logprobs=lp, + mask=mask, + per_sample=True, + ) + assert out.shape == (3,) + + +def test_topk_kl_mask_zero_returns_zero(): + logits = torch.randn(1, 3, 20) + idx, lp = topk_from_logits(logits, top_k=4) + mask = torch.zeros(1, 3) + out = topk_forward_kl( + student_logits=torch.randn(1, 3, 20), + teacher_topk_indices=idx, + teacher_topk_logprobs=lp, + mask=mask, + ) + assert out.item() == 0.0 + + +def test_topk_kl_differentiable_through_student(): + torch.manual_seed(2) + student_logits = torch.randn(1, 3, 20, requires_grad=True) + teacher_logits = torch.randn(1, 3, 20) + idx, lp = topk_from_logits(teacher_logits, top_k=5) + mask = torch.ones(1, 3) + loss = topk_forward_kl(student_logits, idx, lp, mask) + loss.backward() + assert student_logits.grad is not None + assert student_logits.grad.abs().sum() > 0 + + +def test_topk_kl_temperature_scales_loss(): + """Higher temperature softens distributions → smaller KL.""" + torch.manual_seed(3) + V = 50 + teacher_logits = torch.randn(1, 4, V) * 3.0 + student_logits = torch.randn(1, 4, V) * 3.0 + idx, lp_t1 = topk_from_logits(teacher_logits, top_k=20) + idx_t2, lp_t2 = topk_from_logits(teacher_logits / 2.0, top_k=20) + mask = torch.ones(1, 4) + loss_t1 = topk_forward_kl(student_logits, idx, lp_t1, mask, temperature=1.0) + loss_t2 = topk_forward_kl(student_logits, idx_t2, lp_t2, mask, temperature=2.0) + # T=2 should produce a different (typically smaller for sharp distros) loss. + assert not math.isclose(loss_t1.item(), loss_t2.item(), abs_tol=1e-5) + + +def test_topk_kl_handles_partial_mask(): + torch.manual_seed(4) + student_logits = torch.randn(2, 5, 30) + teacher_logits = torch.randn(2, 5, 30) + idx, lp = topk_from_logits(teacher_logits, top_k=10) + mask = torch.tensor([[1, 1, 0, 0, 0], [0, 1, 1, 1, 0]], dtype=torch.float) + loss = topk_forward_kl(student_logits, idx, lp, mask) + # 5 active positions out of 10. + assert loss.item() > 0 + + +# ---------- topk_jsd ---------- + + +def test_topk_jsd_zero_when_distributions_match(): + """JSD between identical distributions is 0 at any β.""" + torch.manual_seed(10) + logits = torch.randn(1, 4, 30, requires_grad=True) + idx, lp = topk_from_logits(logits.detach(), top_k=12) + mask = torch.ones(1, 4) + for beta in (0.0, 0.25, 0.5, 0.75, 1.0): + loss = topk_jsd(logits, idx, lp, mask, beta=beta) + assert loss.item() < 1e-4 + + +def test_topk_jsd_beta_zero_equals_forward_kl(): + """β=0 reduces to forward KL.""" + torch.manual_seed(11) + V = 30 + teacher_logits = torch.randn(2, 4, V) + student_logits = torch.randn(2, 4, V) + idx, lp = topk_from_logits(teacher_logits, top_k=10) + mask = torch.ones(2, 4) + forward = topk_forward_kl(student_logits, idx, lp, mask) + jsd_b0 = topk_jsd(student_logits, idx, lp, mask, beta=0.0) + assert torch.allclose(forward, jsd_b0, atol=1e-4) + + +def test_topk_jsd_symmetric_at_half(): + """β=0.5 should equal (forward + reverse) / 2.""" + torch.manual_seed(12) + teacher_logits = torch.randn(1, 4, 40) + student_logits = torch.randn(1, 4, 40) + idx, lp = topk_from_logits(teacher_logits, top_k=10) + mask = torch.ones(1, 4) + forward = topk_jsd(student_logits, idx, lp, mask, beta=0.0) + reverse = topk_jsd(student_logits, idx, lp, mask, beta=1.0) + mid = topk_jsd(student_logits, idx, lp, mask, beta=0.5) + expected = 0.5 * (forward + reverse) + assert torch.allclose(mid, expected, atol=1e-5) + + +def test_topk_jsd_differentiable(): + torch.manual_seed(13) + student_logits = torch.randn(1, 3, 20, requires_grad=True) + teacher_logits = torch.randn(1, 3, 20) + idx, lp = topk_from_logits(teacher_logits, top_k=5) + mask = torch.ones(1, 3) + loss = topk_jsd(student_logits, idx, lp, mask, beta=0.5) + loss.backward() + assert student_logits.grad is not None + assert student_logits.grad.abs().sum() > 0 + + +def test_topk_jsd_rejects_invalid_beta(): + student_logits = torch.randn(1, 1, 5) + idx = torch.zeros(1, 1, 2, dtype=torch.long) + lp = torch.full((1, 1, 2), -math.log(2.0)) + mask = torch.ones(1, 1) + for bad in (-0.1, 1.1, 2.0): + with pytest.raises(ValueError, match="beta"): + topk_jsd(student_logits, idx, lp, mask, beta=bad) diff --git a/tests/test_trainer_gkd.py b/tests/test_trainer_gkd.py new file mode 100644 index 0000000..6543db6 --- /dev/null +++ b/tests/test_trainer_gkd.py @@ -0,0 +1,349 @@ +"""Integration tests: ContextBakingTrainer with an external teacher backend. + +Uses a tiny shared-tokenizer pair (gpt2 student + a fresh gpt2 teacher) to +validate the GKD code path end-to-end on CPU. +""" + +import pytest +import torch +from peft import LoraConfig as PeftLoraConfig, get_peft_model +from transformers import AutoModelForCausalLM, AutoTokenizer + +from bakery.config import BakeryConfig, ContextConfig +from bakery.data import prompt_baking_collator +from bakery.masking import clear_mask_cache +from bakery.teachers import HFTeacher +from bakery.trainer import ContextBakingTrainer + +CHAT_TEMPLATE = ( + "{% for m in messages %}" + "{{ m['role'] }}: {{ m['content'] }}\n" + "{% endfor %}" + "{% if add_generation_prompt %}assistant: {% endif %}" +) + + +def _mk_tokenizer(): + tok = AutoTokenizer.from_pretrained("gpt2") + tok.pad_token = tok.eos_token + tok.chat_template = CHAT_TEMPLATE + return tok + + +@pytest.fixture(scope="module") +def gkd_trainer(): + """Tiny student (gpt2 + LoRA) + tiny teacher (separate gpt2) on CPU.""" + tokenizer = _mk_tokenizer() + student = AutoModelForCausalLM.from_pretrained("gpt2") + student = get_peft_model( + student, + PeftLoraConfig( + r=4, lora_alpha=8, target_modules=["c_attn"], task_type="CAUSAL_LM" + ), + ) + + args = BakeryConfig( + output_dir="/tmp/bakery_gkd", + num_trajectories=1, + trajectory_length=8, + per_device_train_batch_size=1, + num_train_epochs=1, + logging_steps=1, + report_to="none", + use_cpu=True, + ) + context_config = ContextConfig( + prefix_messages=[{"role": "system", "content": "Be brief."}] + ) + # A fresh gpt2 as the "teacher" — same tokenizer, different params (no LoRA). + teacher = HFTeacher(model_name_or_path="gpt2", torch_dtype="float32") + + return ContextBakingTrainer( + model=student, + args=args, + context_config=context_config, + teacher_backend=teacher, + teacher_top_k=16, + processing_class=tokenizer, + data_collator=prompt_baking_collator, + ) + + +# ---------- compute_loss with external teacher ---------- + + +def test_gkd_compute_loss_returns_scalar(gkd_trainer): + clear_mask_cache() + loss = gkd_trainer.compute_loss( + gkd_trainer.model, + {"user_messages": ["What is two plus two?"], "responses": ["four."]}, + ) + assert loss.dim() == 0 + assert loss.item() >= 0 + + +def test_gkd_compute_loss_nonzero_for_valid_batch(gkd_trainer): + clear_mask_cache() + loss = gkd_trainer.compute_loss( + gkd_trainer.model, + { + "user_messages": ["What is two plus two?", "Capital of France?"], + "responses": ["four.", "Paris."], + }, + ) + assert loss.item() > 0 + + +def test_gkd_compute_loss_zero_when_no_responses(gkd_trainer): + clear_mask_cache() + loss = gkd_trainer.compute_loss( + gkd_trainer.model, + {"user_messages": ["q"], "responses": [""]}, + ) + assert loss.item() == 0.0 + + +def test_gkd_loss_is_differentiable_only_through_student(gkd_trainer): + """Backprop should hit student LoRA params, not teacher.""" + clear_mask_cache() + loss = gkd_trainer.compute_loss( + gkd_trainer.model, + {"user_messages": ["hi"], "responses": ["hey."]}, + ) + loss.backward() + student_has_grad = any( + p.grad is not None and p.grad.abs().sum() > 0 + for p in gkd_trainer.model.parameters() + if p.requires_grad + ) + assert student_has_grad + # Teacher params should not have grads — they were registered with requires_grad=False. + teacher_has_grad = any( + p.grad is not None for p in gkd_trainer.teacher_backend.model.parameters() + ) + assert not teacher_has_grad + + +def test_gkd_uses_topk_path(gkd_trainer): + """Smoke: the teacher dispatch should pick the sparse path.""" + clear_mask_cache() + payload = gkd_trainer._teacher_forward( + gkd_trainer.model, + gkd_trainer._build_batch( + {"user_messages": ["q"], "responses": ["a."]}, gkd_trainer.model + ), + ) + assert payload[0] == "topk" + assert payload[1].shape[-1] <= gkd_trainer.teacher_top_k + + +def test_gkd_prediction_step_returns_triple(gkd_trainer): + """Eval path also takes the topk route and returns a detached scalar.""" + clear_mask_cache() + out = gkd_trainer.prediction_step( + gkd_trainer.model, + {"user_messages": ["q"], "responses": ["a."]}, + prediction_loss_only=True, + ) + assert isinstance(out, tuple) and len(out) == 3 + loss = out[0] + assert loss.dim() == 0 + assert not loss.requires_grad + + +# ---------- generate via external teacher ---------- + + +def test_gkd_generate_trajectory_uses_external_teacher(gkd_trainer): + """When teacher_backend is set, _generate_trajectory delegates to it.""" + # The teacher's tokenizer needs a chat template for generate() — real + # chat-tuned models (Gemma3-it, Qwen) ship one, gpt2 doesn't. + gkd_trainer.teacher_backend.tokenizer.chat_template = CHAT_TEMPLATE + out = gkd_trainer._generate_trajectory("Hello?") + assert isinstance(out, str) + + +# ---------- local-toggle path still works (regression) ---------- + + +# ---------- JSD wiring ---------- + + +def test_gkd_jsd_beta_routes_to_jsd(gkd_trainer): + """gkd_jsd_beta > 0 should route _kl_from_topk through topk_jsd.""" + clear_mask_cache() + gkd_trainer.gkd_jsd_beta = 0.5 + try: + loss = gkd_trainer.compute_loss( + gkd_trainer.model, + {"user_messages": ["hello"], "responses": ["world"]}, + ) + assert loss.item() >= 0 + finally: + gkd_trainer.gkd_jsd_beta = 0.0 + + +def test_gkd_jsd_beta_default_matches_forward_kl(gkd_trainer): + """β=0 default and explicit β=0 should produce the same loss.""" + clear_mask_cache() + inputs = {"user_messages": ["hello"], "responses": ["world"]} + loss_default = gkd_trainer.compute_loss(gkd_trainer.model, inputs) + gkd_trainer.gkd_jsd_beta = 0.0 + clear_mask_cache() + loss_explicit = gkd_trainer.compute_loss(gkd_trainer.model, inputs) + assert torch.allclose(loss_default, loss_explicit, atol=1e-5) + + +def test_trainer_rejects_invalid_jsd_beta(): + tokenizer = _mk_tokenizer() + model = AutoModelForCausalLM.from_pretrained("gpt2") + model = get_peft_model( + model, + PeftLoraConfig( + r=4, lora_alpha=8, target_modules=["c_attn"], task_type="CAUSAL_LM" + ), + ) + with pytest.raises(ValueError, match="gkd_jsd_beta"): + ContextBakingTrainer( + model=model, + args=BakeryConfig( + output_dir="/tmp/bakery_bad", + per_device_train_batch_size=1, + num_train_epochs=1, + logging_steps=1, + report_to="none", + use_cpu=True, + ), + context_config=ContextConfig( + prefix_messages=[{"role": "system", "content": "s"}] + ), + gkd_jsd_beta=2.0, + processing_class=tokenizer, + data_collator=prompt_baking_collator, + ) + + +def test_trainer_rejects_invalid_on_policy_fraction(): + tokenizer = _mk_tokenizer() + model = AutoModelForCausalLM.from_pretrained("gpt2") + model = get_peft_model( + model, + PeftLoraConfig( + r=4, lora_alpha=8, target_modules=["c_attn"], task_type="CAUSAL_LM" + ), + ) + with pytest.raises(ValueError, match="gkd_on_policy_fraction"): + ContextBakingTrainer( + model=model, + args=BakeryConfig( + output_dir="/tmp/bakery_bad2", + per_device_train_batch_size=1, + num_train_epochs=1, + logging_steps=1, + report_to="none", + use_cpu=True, + ), + context_config=ContextConfig( + prefix_messages=[{"role": "system", "content": "s"}] + ), + gkd_on_policy_fraction=-0.1, + processing_class=tokenizer, + data_collator=prompt_baking_collator, + ) + + +# ---------- on-policy routing ---------- + + +def test_on_policy_fraction_one_always_samples_from_student(gkd_trainer, monkeypatch): + """gkd_on_policy_fraction=1.0 → _generate_trajectory uses student sampler.""" + called = {"student": 0, "teacher": 0} + monkeypatch.setattr( + gkd_trainer, + "_sample_from_student", + lambda u: ( + called.__setitem__("student", called["student"] + 1), + "student-sample", + )[1], + ) + monkeypatch.setattr( + gkd_trainer.teacher_backend, + "generate", + lambda msgs, max_new_tokens=256: ( + called.__setitem__("teacher", called["teacher"] + 1), + "teacher-sample", + )[1], + ) + gkd_trainer.gkd_on_policy_fraction = 1.0 + try: + for _ in range(5): + out = gkd_trainer._generate_trajectory("q?") + assert out == "student-sample" + assert called["student"] == 5 + assert called["teacher"] == 0 + finally: + gkd_trainer.gkd_on_policy_fraction = 0.0 + + +def test_on_policy_fraction_zero_never_samples_from_student(gkd_trainer, monkeypatch): + """Default (0.0) → never uses the student sampler.""" + called = {"student": 0, "teacher": 0} + monkeypatch.setattr( + gkd_trainer, + "_sample_from_student", + lambda u: ( + called.__setitem__("student", called["student"] + 1), + "student-sample", + )[1], + ) + monkeypatch.setattr( + gkd_trainer.teacher_backend, + "generate", + lambda msgs, max_new_tokens=256: ( + called.__setitem__("teacher", called["teacher"] + 1), + "teacher-sample", + )[1], + ) + for _ in range(5): + out = gkd_trainer._generate_trajectory("q?") + assert out == "teacher-sample" + assert called["student"] == 0 + assert called["teacher"] == 5 + + +def test_local_toggle_path_unchanged(): + """No teacher_backend → falls back to adapter-toggle KL, dense path.""" + tokenizer = _mk_tokenizer() + model = AutoModelForCausalLM.from_pretrained("gpt2") + model = get_peft_model( + model, + PeftLoraConfig( + r=4, lora_alpha=8, target_modules=["c_attn"], task_type="CAUSAL_LM" + ), + ) + trainer = ContextBakingTrainer( + model=model, + args=BakeryConfig( + output_dir="/tmp/bakery_local_toggle", + num_trajectories=1, + trajectory_length=8, + per_device_train_batch_size=1, + num_train_epochs=1, + logging_steps=1, + report_to="none", + use_cpu=True, + ), + context_config=ContextConfig( + prefix_messages=[{"role": "system", "content": "s"}] + ), + processing_class=tokenizer, + data_collator=prompt_baking_collator, + ) + clear_mask_cache() + payload = trainer._teacher_forward( + trainer.model, + trainer._build_batch( + {"user_messages": ["q"], "responses": ["a"]}, trainer.model + ), + ) + assert payload[0] == "dense"