diff --git a/tests/experimental/test_distillation_trainer.py b/tests/experimental/test_distillation_trainer.py index 9bec167e98e..0d0fcf1bd09 100644 --- a/tests/experimental/test_distillation_trainer.py +++ b/tests/experimental/test_distillation_trainer.py @@ -290,9 +290,14 @@ def test_num_items_in_batch_counts_the_tokens_trained_on(self, monkeypatch): recorded = [] # (denominator applied, completion tokens in this microbatch) original = DistillationTrainer._reduce_divergence_loss - def _recording(jsd, labels=None, reduction="batchmean", num_items_in_batch=None): - recorded.append((num_items_in_batch, int((labels != -100).sum()))) - return original(jsd, labels=labels, reduction=reduction, num_items_in_batch=num_items_in_batch) + def _recording(jsd, completion_mask=None, reduction="batchmean", num_items_in_batch=None): + # `generalized_jsd_loss(reduction="none")` also routes through here with `completion_mask=None`; only the + # loss-reducing call (with a mask) carries the denominator under test. + if completion_mask is not None: + recorded.append((num_items_in_batch, int(completion_mask.sum()))) + return original( + jsd, completion_mask=completion_mask, reduction=reduction, num_items_in_batch=num_items_in_batch + ) monkeypatch.setattr(DistillationTrainer, "_reduce_divergence_loss", staticmethod(_recording)) @@ -385,6 +390,7 @@ def test_loss_normalizes_by_num_items_in_batch(self): "input_ids": input_ids, "attention_mask": torch.ones_like(input_ids), "labels": labels, + "completion_mask": (labels != -100).int(), "prompts": input_ids[:, :prompt_length], "prompt_attention_mask": torch.ones(2, prompt_length, dtype=torch.long, device=device), } diff --git a/trl/experimental/distillation/distillation_trainer.py b/trl/experimental/distillation/distillation_trainer.py index 780ae31bb41..e7ed831db42 100644 --- a/trl/experimental/distillation/distillation_trainer.py +++ b/trl/experimental/distillation/distillation_trainer.py @@ -211,6 +211,7 @@ def __call__(self, examples: list[dict[str, Any]]) -> dict[str, torch.Tensor]: "input_ids": input_ids_t, "attention_mask": attention_mask_t, "labels": labels_t, + "completion_mask": (labels_t != self.ignore_index).int(), "prompts": prompts_t, "prompt_attention_mask": prompt_mask_t, } @@ -882,16 +883,16 @@ def _build_sequence_batch( # ────────────────────────────────────────────────────────────────────── @staticmethod - def _reduce_divergence_loss(jsd, labels=None, reduction="batchmean", num_items_in_batch=None): - """Reduce a per-token divergence tensor using the trainer's label mask semantics. + def _reduce_divergence_loss(jsd, completion_mask=None, reduction="batchmean", num_items_in_batch=None): + """Reduce a per-token divergence tensor over the valid completion tokens. When `num_items_in_batch` is provided (as under gradient accumulation), the divergence is reduced as `sum / num_items_in_batch`, matching the gradient-accumulation-correct behavior of HF's default cross-entropy. Otherwise it falls back to the local `reduction` (default `batchmean`). See issue #4719. """ mask = None - if labels is not None: - mask = labels != -100 + if completion_mask is not None: + mask = completion_mask.bool() jsd = jsd[mask] if num_items_in_batch is not None: @@ -904,7 +905,7 @@ def _reduce_divergence_loss(jsd, labels=None, reduction="batchmean", num_items_i # clamp_min(1) avoids 0/0 -> nan when a sample has no unmasked positions # (e.g. completion fully truncated). jsd[mask] is empty -> jsd.sum() == 0, # so 0/1 == 0 with a valid grad path. - denom = mask.sum().clamp_min(1) if labels is not None else max(jsd.size(0), 1) + denom = mask.sum().clamp_min(1) if completion_mask is not None else max(jsd.size(0), 1) return jsd.sum() / denom elif reduction == "sum": return jsd.sum() @@ -945,7 +946,10 @@ def generalized_jsd_loss( jsd = _jsd_divergence(student_log_probs, teacher_log_probs, beta) return DistillationTrainer._reduce_divergence_loss( - jsd, labels=labels, reduction=reduction, num_items_in_batch=num_items_in_batch + jsd, + completion_mask=(labels != -100) if labels is not None else None, + reduction=reduction, + num_items_in_batch=num_items_in_batch, ) def _get_teacher_logits(self, inputs: dict[str, torch.Tensor | Any]) -> torch.Tensor: @@ -976,19 +980,20 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N attention_mask=inputs["attention_mask"], ) prompt_length = self._compute_prompt_length(inputs) - labels = inputs["labels"][:, prompt_length:] + completion_mask = inputs["completion_mask"][:, prompt_length:] - # Local teacher: exact full-vocabulary (optionally top-k) generalized JSD/KL loss. teacher_logits = self._get_teacher_logits(inputs) student_logits = student_outputs.logits[:, prompt_length - 1 : -1, :] teacher_logits = teacher_logits[:, prompt_length - 1 : -1, :] - loss = self.generalized_jsd_loss( + jsd = self.generalized_jsd_loss( student_logits=student_logits, teacher_logits=teacher_logits, - labels=labels, beta=self.beta, temperature=self.temperature, - num_items_in_batch=num_items_in_batch, + reduction="none", + ) + loss = self._reduce_divergence_loss( + jsd, completion_mask=completion_mask, num_items_in_batch=num_items_in_batch ) return (loss, student_outputs) if return_outputs else loss @@ -1036,8 +1041,10 @@ def _compute_liger_loss(self, model, inputs, num_items_in_batch=None): student_hidden = student_hidden.reshape(-1, student_hidden.shape[-1]) teacher_hidden = teacher_hidden.reshape(-1, teacher_hidden.shape[-1]) - labels_mask = inputs["labels"] != -100 - masked_input_ids = torch.where(labels_mask, inputs["input_ids"], torch.full_like(inputs["input_ids"], -100)) + completion_mask = inputs["completion_mask"].bool() + masked_input_ids = torch.where( + completion_mask, inputs["input_ids"], torch.full_like(inputs["input_ids"], -100) + ) true_labels = masked_input_ids[:, 1:].reshape(-1) student_head = unwrapped_student.get_output_embeddings() diff --git a/trl/experimental/server_distillation/server_distillation_trainer.py b/trl/experimental/server_distillation/server_distillation_trainer.py index a760840ec10..bff7e61abd3 100644 --- a/trl/experimental/server_distillation/server_distillation_trainer.py +++ b/trl/experimental/server_distillation/server_distillation_trainer.py @@ -191,7 +191,7 @@ def _compute_sparse_top_1_divergence_loss( jsd = _jsd_divergence(student_sparse_log_probs, teacher_sparse_log_probs, self.beta, support_mask) return self._reduce_divergence_loss( - jsd, labels=labels, reduction="batchmean", num_items_in_batch=num_items_in_batch + jsd, completion_mask=labels != -100, reduction="batchmean", num_items_in_batch=num_items_in_batch ) def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None): @@ -435,4 +435,4 @@ def _compute_server_forward_kl_loss( ) # See `_compute_server_sparse_top_1_divergence_loss`: the server path normalizes locally, not by # num_items_in_batch, because the teacher window may not cover every student completion token. - return self._reduce_divergence_loss(jsd, labels=labels, reduction="batchmean") + return self._reduce_divergence_loss(jsd, completion_mask=labels != -100, reduction="batchmean")