[DistillationTrainer refactor] Accept a prompt column alongside messages

tests/experimental/test_distillation_trainer.py

@@ -274,6 +274,27 @@ def test_train_updates_params(self):
         for name, param in previous_params.items():
             assert not torch.equal(param, trainer.model.get_parameter(name)), f"Parameter {name} has not changed."
 
+    @pytest.mark.xfail(
+        strict=True,
+        reason="Prompt-only datasets carry no dataset-completion tokens, so num_items_in_batch (counted from the raw "
+        "dataloader labels before on-policy generation replaces them) is 0 and the loss is sum / 0 = NaN. Un-xfail "
+        "when the num_items_in_batch denominator is fixed (plan 5.6).",
+    )
+    def test_train_runs_with_prompt_only_dataset(self):
+        """The forward-looking prompt-only format should train end to end once the loss denominator is fixed."""
+        dataset = load_dataset("trl-internal-testing/zen", "conversational_prompt_only", split="train")
+        trainer = DistillationTrainer(
+            model=self.model_id,
+            teacher_model=self.model_id,
+            args=self._make_args(max_steps=1, learning_rate=0.1),
+            train_dataset=dataset,
+            processing_class=self.tokenizer,
+        )
+
+        trainer.train()
+
+        assert all(torch.isfinite(param).all() for param in trainer.model.parameters())
+
     @pytest.mark.xfail(
         reason="On-policy, num_items_in_batch is computed by transformers from the raw dataloader labels before "
         "generation replaces the completions, and _RepeatBatchDataLoader repeats one generation batch across the "

trl/experimental/distillation/distillation_trainer.py

@@ -148,9 +148,15 @@ def _jsd_divergence(student_log_probs, teacher_log_probs, beta, support_mask=Non
 class _DistillationCollator:
     """Data collator for the distillation trainer with independent prompt/completion budgets.
 
+    Accepts either of two dataset formats:
+
+    - a ``prompt`` column (prompt-only, the forward-looking format shared with GRPO): the student generates the
+      completion on-policy, so there is nothing to train on in the dataset;
+    - a ``messages`` column (language-modeling format): the prompt is everything before the last assistant turn and the
+      completion is that turn.
+
     Unlike ``DataCollatorForChatML``, this collator tokenizes prompts and completions separately so that long
-    completions can never truncate the prompt to empty. It also handles prompt-only data (no assistant completions),
-    which is what on-policy distillation uses.
+    completions can never truncate the prompt to empty.
     """
 
     def __init__(
@@ -176,11 +182,16 @@ def __call__(self, examples: list[dict[str, Any]]) -> dict[str, torch.Tensor]:
         all_prompt_ids: list[list[int]] = []
 
         for example in examples:
-            messages = example[self.messages_key]
-
-            # Split: prompt = everything before the last assistant turn, completion = last assistant turn
-            has_completion = len(messages) > 1 and messages[-1].get("role") == "assistant"
-            prompt_messages = messages[:-1] if has_completion else messages
+            if "prompt" in example:
+                # Prompt-only dataset (the forward-looking format, shared with GRPO): no completion to train on, the
+                # student generates one on-policy.
+                prompt_messages = example["prompt"]
+                has_completion = False
+            else:
+                messages = example[self.messages_key]
+                # Split: prompt = everything before the last assistant turn, completion = last assistant turn
+                has_completion = len(messages) > 1 and messages[-1].get("role") == "assistant"
+                prompt_messages = messages[:-1] if has_completion else messages
 
             # Tokenize prompt with its own budget using the tokenizer's truncation side
             formatted_prompt = self.tokenizer.apply_chat_template(
@@ -612,7 +623,7 @@ def _get_completion_lengths(self, generated_tokens: torch.Tensor, prompt_width:
 
     def _set_signature_columns_if_needed(self):
         super()._set_signature_columns_if_needed()
-        extra_columns = ["prompts", "prompt_attention_mask", "messages", "chat_template_kwargs", "tools"]
+        extra_columns = ["prompt", "prompts", "prompt_attention_mask", "messages", "chat_template_kwargs", "tools"]
         if self._signature_columns is None:
             self._signature_columns = extra_columns
         else: