[DistillationTrainer refactor] Emit prompt_ids / prompt_mask / completion_ids

tests/experimental/test_distillation_trainer.py

@@ -426,6 +426,27 @@ def _capturing(self, model, inputs, *args, **kwargs):
         assert "completion_mask" in inputs
         assert torch.equal(inputs["completion_mask"].bool(), inputs["labels"] != -100)
 
+    def test_generated_batch_emits_prompt_and_completion_ids(self, monkeypatch):
+        """The generated batch emits GRPO-style `prompt_ids`/`prompt_mask`/`completion_ids` alongside the old keys."""
+        trainer = self._make_local_trainer()
+        captured = {}
+        original = DistillationTrainer.compute_loss
+
+        def _capturing(self, model, inputs, *args, **kwargs):
+            captured.setdefault("inputs", {k: v.clone() if torch.is_tensor(v) else v for k, v in inputs.items()})
+            return original(self, model, inputs, *args, **kwargs)
+
+        monkeypatch.setattr(DistillationTrainer, "compute_loss", _capturing)
+        trainer.train()
+
+        inputs = captured["inputs"]
+        for key in ("prompt_ids", "prompt_mask", "completion_ids"):
+            assert key in inputs
+        # cat(prompt_ids, completion_ids) reconstructs input_ids; the new keys mirror the existing prompt tensors.
+        assert torch.equal(torch.cat([inputs["prompt_ids"], inputs["completion_ids"]], dim=1), inputs["input_ids"])
+        assert torch.equal(inputs["prompt_ids"], inputs["prompts"])
+        assert torch.equal(inputs["prompt_mask"], inputs["prompt_attention_mask"])
+
     @require_liger_kernel
     @require_torch_accelerator
     def test_distillation_trainer_with_liger(self):

trl/experimental/distillation/distillation_trainer.py

@@ -214,6 +214,9 @@ def __call__(self, examples: list[dict[str, Any]]) -> dict[str, torch.Tensor]:
             "completion_mask": (labels_t != self.ignore_index).int(),
             "prompts": prompts_t,
             "prompt_attention_mask": prompt_mask_t,
+            "prompt_ids": prompts_t,
+            "prompt_mask": prompt_mask_t,
+            "completion_ids": input_ids_t[:, prompts_t.shape[1] :],
         }
 
 
@@ -773,6 +776,9 @@ def _generate_with_model(self, slices: list[dict[str, torch.Tensor | Any]], on_p
                 updated["attention_mask"] = new_attention_mask
                 updated["labels"] = new_labels
                 updated["completion_mask"] = new_completion_mask
+                updated["prompt_ids"] = slice_inputs["prompts"]
+                updated["prompt_mask"] = prompt_mask
+                updated["completion_ids"] = generated_tokens[:, prompt_width:]
 
                 self._buffered_inputs[slice_idx] = updated
                 self._buffered_text_logs[slice_idx] = (prompt_texts, completion_texts)
@@ -851,6 +857,9 @@ def _store_completions_in_buffer(
             updated["attention_mask"] = new_attention_mask
             updated["labels"] = new_labels
             updated["completion_mask"] = new_completion_mask
+            updated["prompt_ids"] = prompt_ids
+            updated["prompt_mask"] = prompt_attention_mask
+            updated["completion_ids"] = completion_ids_padded
             # Update prompts to match the new padding width so prompt_length is consistent
             updated["prompts"] = prompt_ids
             updated["prompt_attention_mask"] = prompt_attention_mask