From d2a74acdad767279263152fc21ff065fcbde5d3e Mon Sep 17 00:00:00 2001 From: marksverdhei Date: Sun, 29 Mar 2026 01:21:52 +0000 Subject: [PATCH 1/3] perf: vectorize per-sample logit slicing to reduce KL kernel launches (issue #6 item 3) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The loss computation loop called compute_kl_divergence N times (once per sample per batch), each requiring its own softmax + log_softmax + kl_div kernel launch. For batch size B, this is 3*B separate CUDA kernel calls. Replace with a single batched call: 1. Compute per-sample response start positions (CPU, cheap) 2. Build padded batch tensors [valid, max_resp_len, V] via CPU loop (copies only response slices — no wasted full-sequence allocation) 3. Call compute_kl_divergence once with per_sample=True 4. Average the per-sample losses The same vectorization is applied to the sequential_eval path in prediction_step to keep both code paths consistent. For batch size B=8 this reduces CUDA kernel launches from 24 (8 × 3) to 3 (1 × 3), at the cost of one extra CPU loop to build the aligned batch tensors. Net effect: fewer GPU-CPU round trips, better GPU utilization, especially for larger batch sizes. Co-Authored-By: Claude Sonnet 4.6 --- src/bakery/trainer.py | 137 ++++++++++++++++++++++++------------------ 1 file changed, 78 insertions(+), 59 deletions(-) diff --git a/src/bakery/trainer.py b/src/bakery/trainer.py index 29dbb43..4a842bb 100644 --- a/src/bakery/trainer.py +++ b/src/bakery/trainer.py @@ -239,43 +239,54 @@ def compute_loss( s_seq_len = student_inputs["input_ids"].shape[1] t_real_lengths = teacher_inputs["attention_mask"].sum(dim=1) s_real_lengths = student_inputs["attention_mask"].sum(dim=1) - - losses = [] - - for i in range(len(pairs)): - t_pad = t_seq_len - t_real_lengths[i].item() - s_pad = s_seq_len - s_real_lengths[i].item() - t_start = int(t_pad) + teacher_prompt_lengths[i] - s_start = int(s_pad) + student_prompt_lengths[i] - - # Logits at position t predict token t+1, so shift back by 1 to get - # the logits that correspond to each response token. Slice to -1 - # because the last position predicts a token beyond the sequence. - t_logits = teacher_outputs.logits[i : i + 1, t_start - 1 : -1, :] - s_logits = student_outputs.logits[i : i + 1, s_start - 1 : -1, :] - - t_mask = teacher_inputs["attention_mask"][i : i + 1, t_start:] - s_mask = student_inputs["attention_mask"][i : i + 1, s_start:] - - min_len = min(t_logits.shape[1], s_logits.shape[1]) - if min_len == 0: - continue - - t_logits = t_logits[:, :min_len, :] - s_logits = s_logits[:, :min_len, :] - mask = (t_mask[:, :min_len] * s_mask[:, :min_len]).float() - - loss = compute_kl_divergence( - t_logits.detach(), s_logits, mask, self.kl_temperature - ) - losses.append(loss) - - if not losses: + B = len(pairs) + V = teacher_outputs.logits.shape[-1] + + # Compute per-sample response start positions (in logit space, shifted -1 + # so that logit[t] predicts token[t+1]). + t_starts = [ + int(t_seq_len - t_real_lengths[i].item()) + teacher_prompt_lengths[i] + for i in range(B) + ] + s_starts = [ + int(s_seq_len - s_real_lengths[i].item()) + student_prompt_lengths[i] + for i in range(B) + ] + # Response length for sample i: from start to seq_end (exclusive), capped + # at the other sequence's response length to keep teacher/student aligned. + t_resp_lens = [t_seq_len - 1 - (t_starts[i] - 1) for i in range(B)] + s_resp_lens = [s_seq_len - 1 - (s_starts[i] - 1) for i in range(B)] + min_resp_lens = [min(t_resp_lens[i], s_resp_lens[i]) for i in range(B)] + + # Filter out zero-length samples (degenerate prompts/responses). + valid = [i for i, L in enumerate(min_resp_lens) if L > 0] + if not valid: logger.warning("No aligned logit pairs after slicing — returning zero loss") zero = torch.tensor(0.0, device=self.args.device, requires_grad=True) return (zero, None) if return_outputs else zero - total_loss = torch.stack(losses).mean() + max_resp_len = max(min_resp_lens[i] for i in valid) + + # Build batched logit tensors [|valid|, max_resp_len, V] by copying each + # sample's response slice. This CPU loop is cheap (shapes only differ in + # sequence position); the expensive softmax/KL runs once on the batch. + t_batch = teacher_outputs.logits.new_zeros(len(valid), max_resp_len, V) + s_batch = student_outputs.logits.new_zeros(len(valid), max_resp_len, V) + mask_batch = teacher_outputs.logits.new_zeros(len(valid), max_resp_len) + + for out_idx, i in enumerate(valid): + L = min_resp_lens[i] + ts = t_starts[i] - 1 # logit position for first response token + ss = s_starts[i] - 1 + t_batch[out_idx, :L] = teacher_outputs.logits[i, ts : ts + L] + s_batch[out_idx, :L] = student_outputs.logits[i, ss : ss + L] + mask_batch[out_idx, :L] = 1.0 + + per_sample_losses = compute_kl_divergence( + t_batch.detach(), s_batch, mask_batch, self.kl_temperature, + per_sample=True, + ) + total_loss = per_sample_losses.mean() return (total_loss, None) if return_outputs else total_loss def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None): @@ -359,39 +370,47 @@ def _make_fwd(tok_inputs): s_seq_len = student_inputs["input_ids"].shape[1] t_real_lengths = teacher_inputs["attention_mask"].sum(dim=1) s_real_lengths = student_inputs["attention_mask"].sum(dim=1) + B = len(pairs) + V = teacher_logits.shape[-1] - losses = [] - for i in range(len(pairs)): - t_pad = t_seq_len - t_real_lengths[i].item() - s_pad = s_seq_len - s_real_lengths[i].item() - t_start = int(t_pad) + teacher_prompt_lengths[i] - s_start = int(s_pad) + student_prompt_lengths[i] - - t_logits = teacher_logits[i : i + 1, t_start - 1 : -1, :].to(model.device) - s_logits = student_outputs.logits[i : i + 1, s_start - 1 : -1, :] - t_mask = teacher_inputs["attention_mask"][i : i + 1, t_start:] - s_mask = student_inputs["attention_mask"][i : i + 1, s_start:] - - min_len = min(t_logits.shape[1], s_logits.shape[1]) - if min_len == 0: - continue - mask = (t_mask[:, :min_len] * s_mask[:, :min_len]).float() - losses.append( - compute_kl_divergence( - t_logits[:, :min_len, :], - s_logits[:, :min_len, :], - mask, - self.kl_temperature, - ) - ) + t_starts = [ + int(t_seq_len - t_real_lengths[i].item()) + teacher_prompt_lengths[i] + for i in range(B) + ] + s_starts = [ + int(s_seq_len - s_real_lengths[i].item()) + student_prompt_lengths[i] + for i in range(B) + ] + t_resp_lens = [t_seq_len - 1 - (t_starts[i] - 1) for i in range(B)] + s_resp_lens = [s_seq_len - 1 - (s_starts[i] - 1) for i in range(B)] + min_resp_lens = [min(t_resp_lens[i], s_resp_lens[i]) for i in range(B)] - if not losses: + valid = [i for i, L in enumerate(min_resp_lens) if L > 0] + if not valid: return ( torch.tensor(0.0, device=self.args.device, requires_grad=True), None, None, ) - return (torch.stack(losses).mean().detach(), None, None) + + max_resp_len = max(min_resp_lens[i] for i in valid) + dev = student_outputs.logits.device + t_batch = student_outputs.logits.new_zeros(len(valid), max_resp_len, V) + s_batch = student_outputs.logits.new_zeros(len(valid), max_resp_len, V) + mask_batch = student_outputs.logits.new_zeros(len(valid), max_resp_len) + + for out_idx, i in enumerate(valid): + L = min_resp_lens[i] + ts = t_starts[i] - 1 + ss = s_starts[i] - 1 + t_batch[out_idx, :L] = teacher_logits[i, ts : ts + L].to(dev) + s_batch[out_idx, :L] = student_outputs.logits[i, ss : ss + L] + mask_batch[out_idx, :L] = 1.0 + + per_sample_losses = compute_kl_divergence( + t_batch, s_batch, mask_batch, self.kl_temperature, per_sample=True + ) + return (per_sample_losses.mean().detach(), None, None) def training_step(self, model, inputs, num_items_in_batch=None) -> torch.Tensor: """Generate trajectories on-the-fly if no precomputed responses.""" From 60aa639226cd515425c13906214e75a5e2a2a5b4 Mon Sep 17 00:00:00 2001 From: marksverdhei Date: Sun, 29 Mar 2026 17:51:28 +0000 Subject: [PATCH 2/3] refactor: extract _compute_batched_kl helper and add numerical equivalence test Address PR #20 review feedback: 1. Add test_batched_kl_matches_per_sample_loop that verifies the batched KL computation matches a per-sample reference loop on fixed-seed data. 2. Extract duplicated logic between compute_loss and prediction_step into shared helpers: _prepare_pairs, _build_texts_and_lengths, _make_fwd_kwargs, and _compute_batched_kl. 3. Simplify t_seq_len - 1 - (t_starts[i] - 1) to t_seq_len - t_starts[i]. Co-Authored-By: Claude Opus 4.6 --- src/bakery/trainer.py | 254 ++++++++++++++++++++---------------------- tests/test_trainer.py | 98 ++++++++++++++++ 2 files changed, 216 insertions(+), 136 deletions(-) diff --git a/src/bakery/trainer.py b/src/bakery/trainer.py index 4a842bb..e5a0bf4 100644 --- a/src/bakery/trainer.py +++ b/src/bakery/trainer.py @@ -147,30 +147,27 @@ def _generate_trajectory(self, user_message: str) -> str: # -- Loss computation -- - def compute_loss( - self, model, inputs, return_outputs=False, num_items_in_batch=None - ): - """Compute KL divergence loss with batched forward passes.""" + def _prepare_pairs(self, inputs): + """Extract and validate (user_message, response) pairs from inputs. + + Returns a list of (user_msg, response) tuples with empty responses + filtered out, or None if the batch is empty/invalid. + """ user_messages = inputs.get("user_messages", []) responses = inputs.get("responses", []) - if not user_messages or not responses: - logger.warning( - "Batch has no user_messages or responses — returning zero loss" - ) - loss = torch.tensor(0.0, device=self.args.device, requires_grad=True) - return (loss, None) if return_outputs else loss - + return None pairs = [ (msg, resp) for msg, resp in zip(user_messages, responses) if resp.strip() ] - if not pairs: - logger.warning( - "All responses in batch are empty/whitespace — returning zero loss" - ) - loss = torch.tensor(0.0, device=self.args.device, requires_grad=True) - return (loss, None) if return_outputs else loss + return pairs if pairs else None + + def _build_texts_and_lengths(self, pairs): + """Build teacher/student chat texts and prompt lengths for each pair. + Returns (teacher_texts, student_texts, teacher_prompt_lengths, + student_prompt_lengths). + """ teacher_texts, student_texts = [], [] teacher_prompt_lengths, student_prompt_lengths = [], [] @@ -196,51 +193,46 @@ def compute_loss( teacher_prompt_lengths.append(t_len) student_prompt_lengths.append(s_len) - with padding_side(self.processing_class, "left"): - teacher_inputs = self._tokenize( - teacher_texts, return_tensors="pt", padding=True - ).to(model.device) - student_inputs = self._tokenize( - student_texts, return_tensors="pt", padding=True - ).to(model.device) + return teacher_texts, student_texts, teacher_prompt_lengths, student_prompt_lengths - teacher_fwd = dict( - input_ids=teacher_inputs["input_ids"], - attention_mask=teacher_inputs["attention_mask"], + def _make_fwd_kwargs(self, model, tok_inputs): + """Build forward-pass keyword arguments, handling token_type_ids.""" + fwd = dict( + input_ids=tok_inputs["input_ids"], + attention_mask=tok_inputs["attention_mask"], ) - student_fwd = dict( - input_ids=student_inputs["input_ids"], - attention_mask=student_inputs["attention_mask"], - ) - # Some architectures (e.g. Gemma 3) require token_type_ids during training. - # The tokenizer may not return them, so create zeros if the model expects them. if hasattr(model.config, "model_type") and model.config.model_type in ( "gemma3", ): - teacher_fwd["token_type_ids"] = torch.zeros_like( - teacher_inputs["input_ids"] - ) - student_fwd["token_type_ids"] = torch.zeros_like( - student_inputs["input_ids"] - ) - elif "token_type_ids" in teacher_inputs: - teacher_fwd["token_type_ids"] = teacher_inputs["token_type_ids"] - student_fwd["token_type_ids"] = student_inputs["token_type_ids"] + fwd["token_type_ids"] = torch.zeros_like(tok_inputs["input_ids"]) + elif "token_type_ids" in tok_inputs: + fwd["token_type_ids"] = tok_inputs["token_type_ids"] + return fwd - with torch.no_grad(): - with disable_adapters(model): - teacher_outputs = model(**teacher_fwd) + def _compute_batched_kl( + self, + teacher_logits, + student_logits, + teacher_inputs, + student_inputs, + teacher_prompt_lengths, + student_prompt_lengths, + B, + ): + """Compute batched KL divergence from aligned teacher/student logits. - student_outputs = model(**student_fwd) + Slices response-only logits from each sequence (accounting for + left-padding offsets), assembles them into aligned batch tensors, + and returns per-sample KL losses. - # With left-padding, each sequence has leading pad tokens that shift - # the real content rightward. Compute per-sequence padding offsets. + Returns per-sample loss tensor of shape [|valid|], or None if no + valid aligned logit pairs exist. + """ t_seq_len = teacher_inputs["input_ids"].shape[1] s_seq_len = student_inputs["input_ids"].shape[1] t_real_lengths = teacher_inputs["attention_mask"].sum(dim=1) s_real_lengths = student_inputs["attention_mask"].sum(dim=1) - B = len(pairs) - V = teacher_outputs.logits.shape[-1] + V = teacher_logits.shape[-1] # Compute per-sample response start positions (in logit space, shifted -1 # so that logit[t] predicts token[t+1]). @@ -254,38 +246,84 @@ def compute_loss( ] # Response length for sample i: from start to seq_end (exclusive), capped # at the other sequence's response length to keep teacher/student aligned. - t_resp_lens = [t_seq_len - 1 - (t_starts[i] - 1) for i in range(B)] - s_resp_lens = [s_seq_len - 1 - (s_starts[i] - 1) for i in range(B)] + t_resp_lens = [t_seq_len - t_starts[i] for i in range(B)] + s_resp_lens = [s_seq_len - s_starts[i] for i in range(B)] min_resp_lens = [min(t_resp_lens[i], s_resp_lens[i]) for i in range(B)] # Filter out zero-length samples (degenerate prompts/responses). valid = [i for i, L in enumerate(min_resp_lens) if L > 0] if not valid: - logger.warning("No aligned logit pairs after slicing — returning zero loss") - zero = torch.tensor(0.0, device=self.args.device, requires_grad=True) - return (zero, None) if return_outputs else zero + return None max_resp_len = max(min_resp_lens[i] for i in valid) # Build batched logit tensors [|valid|, max_resp_len, V] by copying each # sample's response slice. This CPU loop is cheap (shapes only differ in # sequence position); the expensive softmax/KL runs once on the batch. - t_batch = teacher_outputs.logits.new_zeros(len(valid), max_resp_len, V) - s_batch = student_outputs.logits.new_zeros(len(valid), max_resp_len, V) - mask_batch = teacher_outputs.logits.new_zeros(len(valid), max_resp_len) + dev = student_logits.device + t_batch = student_logits.new_zeros(len(valid), max_resp_len, V) + s_batch = student_logits.new_zeros(len(valid), max_resp_len, V) + mask_batch = student_logits.new_zeros(len(valid), max_resp_len) for out_idx, i in enumerate(valid): L = min_resp_lens[i] ts = t_starts[i] - 1 # logit position for first response token ss = s_starts[i] - 1 - t_batch[out_idx, :L] = teacher_outputs.logits[i, ts : ts + L] - s_batch[out_idx, :L] = student_outputs.logits[i, ss : ss + L] + t_batch[out_idx, :L] = teacher_logits[i, ts : ts + L].to(dev) + s_batch[out_idx, :L] = student_logits[i, ss : ss + L] mask_batch[out_idx, :L] = 1.0 per_sample_losses = compute_kl_divergence( t_batch.detach(), s_batch, mask_batch, self.kl_temperature, per_sample=True, ) + return per_sample_losses + + def compute_loss( + self, model, inputs, return_outputs=False, num_items_in_batch=None + ): + """Compute KL divergence loss with batched forward passes.""" + pairs = self._prepare_pairs(inputs) + if pairs is None: + logger.warning( + "Batch has no valid user_messages/responses — returning zero loss" + ) + loss = torch.tensor(0.0, device=self.args.device, requires_grad=True) + return (loss, None) if return_outputs else loss + + teacher_texts, student_texts, teacher_prompt_lengths, student_prompt_lengths = ( + self._build_texts_and_lengths(pairs) + ) + + with padding_side(self.processing_class, "left"): + teacher_inputs = self._tokenize( + teacher_texts, return_tensors="pt", padding=True + ).to(model.device) + student_inputs = self._tokenize( + student_texts, return_tensors="pt", padding=True + ).to(model.device) + + with torch.no_grad(): + with disable_adapters(model): + teacher_outputs = model(**self._make_fwd_kwargs(model, teacher_inputs)) + + student_outputs = model(**self._make_fwd_kwargs(model, student_inputs)) + + per_sample_losses = self._compute_batched_kl( + teacher_outputs.logits, + student_outputs.logits, + teacher_inputs, + student_inputs, + teacher_prompt_lengths, + student_prompt_lengths, + len(pairs), + ) + + if per_sample_losses is None: + logger.warning("No aligned logit pairs after slicing — returning zero loss") + zero = torch.tensor(0.0, device=self.args.device, requires_grad=True) + return (zero, None) if return_outputs else zero + total_loss = per_sample_losses.mean() return (total_loss, None) if return_outputs else total_loss @@ -303,9 +341,7 @@ def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None) return (loss.detach(), None, None) # Sequential eval: teacher → offload logits to CPU → student - user_messages = inputs.get("user_messages", []) - responses = inputs.get("responses", []) - pairs = [(m, r) for m, r in zip(user_messages, responses) if r.strip()] + pairs = self._prepare_pairs(inputs) if not pairs: return ( torch.tensor(0.0, device=self.args.device, requires_grad=True), @@ -313,27 +349,9 @@ def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None) None, ) - teacher_texts, student_texts = [], [] - teacher_prompt_lengths, student_prompt_lengths = [], [] - for user_msg, response in pairs: - t_msgs = [ - {"role": "system", "content": self.system_prompt}, - {"role": "user", "content": user_msg}, - {"role": "assistant", "content": response}, - ] - s_msgs = [ - {"role": "user", "content": user_msg}, - {"role": "assistant", "content": response}, - ] - teacher_texts.append( - self.processing_class.apply_chat_template(t_msgs, tokenize=False) - ) - student_texts.append( - self.processing_class.apply_chat_template(s_msgs, tokenize=False) - ) - t_len, s_len = self._get_prompt_lengths(user_msg) - teacher_prompt_lengths.append(t_len) - student_prompt_lengths.append(s_len) + teacher_texts, student_texts, teacher_prompt_lengths, student_prompt_lengths = ( + self._build_texts_and_lengths(pairs) + ) with padding_side(self.processing_class, "left"): teacher_inputs = self._tokenize( @@ -343,73 +361,37 @@ def prediction_step(self, model, inputs, prediction_loss_only, ignore_keys=None) student_texts, return_tensors="pt", padding=True ).to(model.device) - def _make_fwd(tok_inputs): - fwd = dict( - input_ids=tok_inputs["input_ids"], - attention_mask=tok_inputs["attention_mask"], - ) - if hasattr(model.config, "model_type") and model.config.model_type in ( - "gemma3", - ): - fwd["token_type_ids"] = torch.zeros_like(tok_inputs["input_ids"]) - elif "token_type_ids" in tok_inputs: - fwd["token_type_ids"] = tok_inputs["token_type_ids"] - return fwd - # Accelerate replaces model.forward with a wrapper that upcasts to fp32. # Bypass by calling the CLASS forward method directly. 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, **_make_fwd(teacher_inputs)).logits.cpu() + teacher_logits = fwd_fn( + base, **self._make_fwd_kwargs(base, teacher_inputs) + ).logits.cpu() torch.cuda.empty_cache() - student_outputs = fwd_fn(base, **_make_fwd(student_inputs)) - - t_seq_len = teacher_inputs["input_ids"].shape[1] - s_seq_len = student_inputs["input_ids"].shape[1] - t_real_lengths = teacher_inputs["attention_mask"].sum(dim=1) - s_real_lengths = student_inputs["attention_mask"].sum(dim=1) - B = len(pairs) - V = teacher_logits.shape[-1] + student_outputs = fwd_fn( + base, **self._make_fwd_kwargs(base, student_inputs) + ) - t_starts = [ - int(t_seq_len - t_real_lengths[i].item()) + teacher_prompt_lengths[i] - for i in range(B) - ] - s_starts = [ - int(s_seq_len - s_real_lengths[i].item()) + student_prompt_lengths[i] - for i in range(B) - ] - t_resp_lens = [t_seq_len - 1 - (t_starts[i] - 1) for i in range(B)] - s_resp_lens = [s_seq_len - 1 - (s_starts[i] - 1) for i in range(B)] - min_resp_lens = [min(t_resp_lens[i], s_resp_lens[i]) for i in range(B)] + per_sample_losses = self._compute_batched_kl( + teacher_logits, + student_outputs.logits, + teacher_inputs, + student_inputs, + teacher_prompt_lengths, + student_prompt_lengths, + len(pairs), + ) - valid = [i for i, L in enumerate(min_resp_lens) if L > 0] - if not valid: + if per_sample_losses is None: return ( torch.tensor(0.0, device=self.args.device, requires_grad=True), None, None, ) - max_resp_len = max(min_resp_lens[i] for i in valid) - dev = student_outputs.logits.device - t_batch = student_outputs.logits.new_zeros(len(valid), max_resp_len, V) - s_batch = student_outputs.logits.new_zeros(len(valid), max_resp_len, V) - mask_batch = student_outputs.logits.new_zeros(len(valid), max_resp_len) - - for out_idx, i in enumerate(valid): - L = min_resp_lens[i] - ts = t_starts[i] - 1 - ss = s_starts[i] - 1 - t_batch[out_idx, :L] = teacher_logits[i, ts : ts + L].to(dev) - s_batch[out_idx, :L] = student_outputs.logits[i, ss : ss + L] - mask_batch[out_idx, :L] = 1.0 - - per_sample_losses = compute_kl_divergence( - t_batch, s_batch, mask_batch, self.kl_temperature, per_sample=True - ) return (per_sample_losses.mean().detach(), None, None) def training_step(self, model, inputs, num_items_in_batch=None) -> torch.Tensor: diff --git a/tests/test_trainer.py b/tests/test_trainer.py index f1dc148..326fab8 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -1,10 +1,12 @@ """Tests for PromptBakingTrainer using a tiny GPT-2 model with LoRA.""" +import torch from transformers import AutoModelForCausalLM, AutoTokenizer from peft import LoraConfig as PeftLoraConfig, get_peft_model from bakery.config import BakeryConfig from bakery.data import create_dataset, prompt_baking_collator +from bakery.kl import compute_kl_divergence, disable_adapters, padding_side from bakery.trainer import PromptBakingTrainer @@ -283,3 +285,99 @@ def test_prediction_step_sequential_eval_returns_triple(): assert result[1] is None and result[2] is None loss = result[0] assert loss.dim() == 0 + + +# --------------------------------------------------------------------------- +# Numerical equivalence: batched vs per-sample loop +# --------------------------------------------------------------------------- + + +def test_batched_kl_matches_per_sample_loop(): + """Verify that the batched _compute_batched_kl produces the same result + as computing KL divergence one sample at a time in a loop. + + This guards against regressions when refactoring the vectorized path. + """ + torch.manual_seed(42) + + trainer = _make_trainer( + prompts=["What is 2+2?", "Explain gravity"], + responses=["The answer is 4.", "Gravity is a fundamental force of nature."], + batch_size=2, + ) + model = trainer.model + + user_messages = ["What is 2+2?", "Explain gravity"] + responses = ["The answer is 4.", "Gravity is a fundamental force of nature."] + pairs = list(zip(user_messages, responses)) + + teacher_texts, student_texts, t_prompt_lens, s_prompt_lens = ( + trainer._build_texts_and_lengths(pairs) + ) + + with padding_side(trainer.processing_class, "left"): + teacher_inputs = trainer._tokenize( + teacher_texts, return_tensors="pt", padding=True + ).to(model.device) + student_inputs = trainer._tokenize( + student_texts, return_tensors="pt", padding=True + ).to(model.device) + + with torch.no_grad(): + with disable_adapters(model): + teacher_logits = model( + **trainer._make_fwd_kwargs(model, teacher_inputs) + ).logits + student_logits = model( + **trainer._make_fwd_kwargs(model, student_inputs) + ).logits + + # --- Batched path (the code under test) --- + batched_losses = trainer._compute_batched_kl( + teacher_logits, + student_logits, + teacher_inputs, + student_inputs, + t_prompt_lens, + s_prompt_lens, + len(pairs), + ) + assert batched_losses is not None + + # --- Reference: per-sample loop (the old approach) --- + per_sample_losses = [] + t_seq_len = teacher_inputs["input_ids"].shape[1] + s_seq_len = student_inputs["input_ids"].shape[1] + t_real_lengths = teacher_inputs["attention_mask"].sum(dim=1) + s_real_lengths = student_inputs["attention_mask"].sum(dim=1) + + for i in range(len(pairs)): + t_start = int(t_seq_len - t_real_lengths[i].item()) + t_prompt_lens[i] + s_start = int(s_seq_len - s_real_lengths[i].item()) + s_prompt_lens[i] + t_resp_len = t_seq_len - t_start + s_resp_len = s_seq_len - s_start + L = min(t_resp_len, s_resp_len) + if L <= 0: + continue + + ts = t_start - 1 # logit position for first response token + ss = s_start - 1 + t_logits_i = teacher_logits[i, ts : ts + L].unsqueeze(0) + s_logits_i = student_logits[i, ss : ss + L].unsqueeze(0) + mask_i = torch.ones(1, L, device=model.device) + + loss_i = compute_kl_divergence( + t_logits_i.detach(), s_logits_i, mask_i, + trainer.kl_temperature, per_sample=True, + ) + per_sample_losses.append(loss_i.squeeze(0)) + + assert len(per_sample_losses) == batched_losses.shape[0] + loop_losses = torch.stack(per_sample_losses) + + assert torch.allclose(batched_losses, loop_losses, atol=1e-5), ( + f"Batched and per-sample loop KL losses differ:\n" + f" batched: {batched_losses}\n" + f" loop: {loop_losses}\n" + f" max diff: {(batched_losses - loop_losses).abs().max().item()}" + ) From 6770512983df5979ad7d8123b689effe700d6073 Mon Sep 17 00:00:00 2001 From: marksverdhei Date: Sun, 29 Mar 2026 17:51:58 +0000 Subject: [PATCH 3/3] style: apply ruff formatting --- src/bakery/trainer.py | 12 ++++++++++-- tests/test_trainer.py | 11 ++++++----- 2 files changed, 16 insertions(+), 7 deletions(-) diff --git a/src/bakery/trainer.py b/src/bakery/trainer.py index e5a0bf4..d42e222 100644 --- a/src/bakery/trainer.py +++ b/src/bakery/trainer.py @@ -193,7 +193,12 @@ def _build_texts_and_lengths(self, pairs): teacher_prompt_lengths.append(t_len) student_prompt_lengths.append(s_len) - return teacher_texts, student_texts, teacher_prompt_lengths, student_prompt_lengths + return ( + teacher_texts, + student_texts, + teacher_prompt_lengths, + student_prompt_lengths, + ) def _make_fwd_kwargs(self, model, tok_inputs): """Build forward-pass keyword arguments, handling token_type_ids.""" @@ -274,7 +279,10 @@ def _compute_batched_kl( mask_batch[out_idx, :L] = 1.0 per_sample_losses = compute_kl_divergence( - t_batch.detach(), s_batch, mask_batch, self.kl_temperature, + t_batch.detach(), + s_batch, + mask_batch, + self.kl_temperature, per_sample=True, ) return per_sample_losses diff --git a/tests/test_trainer.py b/tests/test_trainer.py index 326fab8..5d9d517 100644 --- a/tests/test_trainer.py +++ b/tests/test_trainer.py @@ -328,9 +328,7 @@ def test_batched_kl_matches_per_sample_loop(): teacher_logits = model( **trainer._make_fwd_kwargs(model, teacher_inputs) ).logits - student_logits = model( - **trainer._make_fwd_kwargs(model, student_inputs) - ).logits + student_logits = model(**trainer._make_fwd_kwargs(model, student_inputs)).logits # --- Batched path (the code under test) --- batched_losses = trainer._compute_batched_kl( @@ -367,8 +365,11 @@ def test_batched_kl_matches_per_sample_loop(): mask_i = torch.ones(1, L, device=model.device) loss_i = compute_kl_divergence( - t_logits_i.detach(), s_logits_i, mask_i, - trainer.kl_temperature, per_sample=True, + t_logits_i.detach(), + s_logits_i, + mask_i, + trainer.kl_temperature, + per_sample=True, ) per_sample_losses.append(loss_i.squeeze(0))