[DistillationTrainer refactor] Remove messages-format support and prompt-length config

tests/experimental/test_distillation_trainer.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 import os
-import warnings
 
 import pytest
 import torch
@@ -22,7 +21,7 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 from trl.experimental.distillation import DistillationConfig, DistillationTrainer
-from trl.experimental.distillation.distillation_trainer import _DistillationCollator, _RepeatBatchDataLoader
+from trl.experimental.distillation.distillation_trainer import _RepeatBatchDataLoader
 from trl.experimental.gkd.gkd_trainer import GKDTrainer
 
 from ..testing_utils import TrlTestCase, require_liger_kernel, require_torch_accelerator
@@ -188,19 +187,6 @@ def setup_method(self):
         self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)
         self.tokenizer.pad_token = self.tokenizer.eos_token
 
-    def test_messages_format_dataset_is_deprecated(self):
-        """The `messages` (language-modeling) format is deprecated in favour of a prompt-only `prompt` column."""
-        messages_example = {"messages": [{"role": "user", "content": "Hi"}, {"role": "assistant", "content": "Hi!"}]}
-        with pytest.warns(FutureWarning, match="`messages`-format"):
-            _DistillationCollator(self.tokenizer, max_length=128, max_prompt_length=64)([messages_example])
-
-        # The forward-looking `prompt` column must not trigger the deprecation.
-        with warnings.catch_warnings():
-            warnings.simplefilter("error", FutureWarning)
-            _DistillationCollator(self.tokenizer, max_length=128, max_prompt_length=64)(
-                [{"prompt": [{"role": "user", "content": "Hi"}]}]
-            )
-
     def _make_args(self, **kwargs):
         args = {
             "output_dir": self.tmp_dir,
@@ -212,7 +198,6 @@ def _make_args(self, **kwargs):
             "disable_tqdm": True,
             "use_cpu": True,
             "bf16": False,
-            "max_length": 128,
             "max_completion_length": 32,
             "model_init_kwargs": {"dtype": "float32", "device_map": None},
             "teacher_model_init_kwargs": {"dtype": "float32", "device_map": None},
@@ -221,9 +206,7 @@ def _make_args(self, **kwargs):
         return DistillationConfig(**args)
 
     def _make_local_trainer(self, **kwargs):
-        # Messages-format: `_make_batch` reads completion tokens straight from the collator (no generation), which the
-        # `prompt`-only format cannot provide until generation replaces the collator. Switched to prompt-only then.
-        dataset = load_dataset("trl-internal-testing/zen", "conversational_language_modeling", split="train")
+        dataset = load_dataset("trl-internal-testing/zen", "conversational_prompt_only", split="train")
         return DistillationTrainer(
             model=self.model_id,
             teacher_model=self.model_id,
@@ -232,14 +215,6 @@ def _make_local_trainer(self, **kwargs):
             processing_class=self.tokenizer,
         )
 
-    def _make_batch(self, trainer):
-        examples = [trainer.train_dataset[i] for i in range(2)]
-        return trainer.data_collator(examples)
-
-    @staticmethod
-    def _move_batch_to_device(batch, device):
-        return {key: value.to(device) for key, value in batch.items()}
-
     def test_distillation_trainer_train_runs_with_local_teacher(self):
         training_args = self._make_args(
             dataloader_drop_last=True,
@@ -398,11 +373,24 @@ def test_loss_normalizes_by_num_items_in_batch(self):
             for p in trainer.teacher_model.parameters():
                 p.add_(0.5 * torch.randn_like(p))
 
-        batch = self._move_batch_to_device(self._make_batch(trainer), trainer.accelerator.device)
+        # The collator is prompt-only (completions come from on-policy generation); build a batch with completion
+        # tokens directly to exercise the loss reduction.
+        device = trainer.accelerator.device
+        prompt_length, completion_length = 4, 3
+        seq_length = prompt_length + completion_length
+        input_ids = torch.randint(0, trainer.model.config.vocab_size, (2, seq_length), device=device)
+        labels = input_ids.clone()
+        labels[:, :prompt_length] = -100
+        batch = {
+            "input_ids": input_ids,
+            "attention_mask": torch.ones_like(input_ids),
+            "labels": labels,
+            "prompts": input_ids[:, :prompt_length],
+            "prompt_attention_mask": torch.ones(2, prompt_length, dtype=torch.long, device=device),
+        }
 
         # Number of valid (non-ignored) tokens in the local batch, sliced the same way `compute_loss` does.
-        prompt_length = trainer._compute_prompt_length(batch)
-        num_valid = (batch["labels"][:, prompt_length:] != -100).sum()
+        num_valid = (labels[:, prompt_length:] != -100).sum()
 
         trainer.model.eval()
         with torch.no_grad():

tests/experimental/test_server_distillation_trainer.py

@@ -69,16 +69,8 @@ def _canned_teacher_logprobs(**kwargs):
 def _variable_length_dataset():
     return Dataset.from_list(
         [
-            {"messages": [{"role": "user", "content": "What's 2+2?"}, {"role": "assistant", "content": "4."}]},
-            {
-                "messages": [
-                    {"role": "user", "content": "Name three primary colors."},
-                    {
-                        "role": "assistant",
-                        "content": "Red, green, and blue are the three primary colors commonly used in additive color mixing.",
-                    },
-                ]
-            },
+            {"prompt": [{"role": "user", "content": "What's 2+2?"}]},
+            {"prompt": [{"role": "user", "content": "Name three primary colors."}]},
         ]
     )
 
@@ -272,8 +264,6 @@ def _run_one_step(self, bs, ga, monkeypatch):
             per_device_train_batch_size=bs,
             gradient_accumulation_steps=ga,
             learning_rate=1e-4,
-            max_length=64,
-            max_prompt_length=32,
             max_completion_length=32,
             teacher_model_server_url="http://fake-teacher.invalid:8000",
             loss_top_k=1,

trl/experimental/distillation/distillation_config.py

@@ -41,8 +41,6 @@ class DistillationConfig(_BaseConfig):
             Whether to allow loading models and tokenizers that ship custom Python code from the Hub. Forwarded to
             [`~transformers.AutoModelForCausalLM.from_pretrained`] and
             [`~transformers.AutoTokenizer.from_pretrained`], for both the student and teacher.
-        max_length (`int` or `None`, *optional*, defaults to `1024`):
-            Maximum total sequence length (prompt + completion) for tokenization and truncation.
 
         > Parameters that control the distillation
 
@@ -55,9 +53,6 @@ class DistillationConfig(_BaseConfig):
             JSD.
         max_completion_length (`int`, *optional*, defaults to `512`):
             Maximum number of tokens to generate per completion during on-policy generation.
-        max_prompt_length (`int` or `None`, *optional*):
-            Maximum number of tokens for the prompt. If `None`, auto-computed as `max_length - max_completion_length`.
-            Prompts are truncated according to the tokenizer's `truncation_side` setting.
         disable_dropout (`bool`, *optional*, defaults to `True`):
             Whether to disable dropout in the student model during training.
 
@@ -143,11 +138,6 @@ class DistillationConfig(_BaseConfig):
             "student and teacher."
         },
     )
-    max_length: int | None = field(
-        default=1024,
-        metadata={"help": "Maximum total sequence length (prompt + completion) for tokenization and truncation."},
-    )
-
     # Overridden defaults
     learning_rate: float = field(
         default=1e-6,
@@ -172,14 +162,6 @@ class DistillationConfig(_BaseConfig):
         default=512,
         metadata={"help": "Maximum number of tokens to generate per completion."},
     )
-    max_prompt_length: int | None = field(
-        default=None,
-        metadata={
-            "help": "Maximum number of tokens for the prompt. If None, auto-computed as "
-            "max_length - max_completion_length. Prompts are truncated according to the "
-            "tokenizer's truncation_side setting."
-        },
-    )
     disable_dropout: bool = field(
         default=True,
         metadata={"help": "Whether to disable dropout in the student model during training."},
@@ -306,15 +288,6 @@ def __post_init__(self):
         if self.beta < 0.0 or self.beta > 1.0:
             raise ValueError(f"beta must be in [0.0, 1.0], got {self.beta}.")
 
-        if self.max_length is not None and self.max_completion_length >= self.max_length:
-            raise ValueError(
-                f"max_completion_length ({self.max_completion_length}) must be smaller than "
-                f"max_length ({self.max_length}) to leave room for the prompt."
-            )
-
-        if self.max_prompt_length is None and self.max_length is not None:
-            self.max_prompt_length = self.max_length - self.max_completion_length
-
         if self.num_generations < 1:
             raise ValueError(f"num_generations must be at least 1, got {self.num_generations}.")
 

trl/experimental/distillation/distillation_trainer.py

@@ -14,7 +14,6 @@
 
 import random
 import textwrap
-import warnings
 from collections import defaultdict
 from collections.abc import Callable
 from contextlib import nullcontext
@@ -147,33 +146,19 @@ 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.
+    """Data collator for the distillation trainer.
 
-    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.
+    Accepts a prompt-only dataset (a ``prompt`` column, the format shared with GRPO): the student generates the
+    completion on-policy, so there is nothing to train on in the dataset.
     """
 
     def __init__(
         self,
         tokenizer: "PreTrainedTokenizerBase",
-        max_length: int,
-        max_prompt_length: int,
-        messages_key: str = "messages",
         ignore_index: int = -100,
     ):
         self.tokenizer = tokenizer
-        self.max_length = max_length
-        self.max_prompt_length = max_prompt_length
-        self.messages_key = messages_key
         self.ignore_index = ignore_index
-        self._warned_messages_deprecated = False
 
         if tokenizer.pad_token_id is None:
             raise ValueError("The tokenizer does not have a pad token. Please set `pad_token_id` in the tokenizer.")
@@ -184,69 +169,14 @@ def __call__(self, examples: list[dict[str, Any]]) -> dict[str, torch.Tensor]:
         all_prompt_ids: list[list[int]] = []
 
         for example in examples:
-            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:
-                if not self._warned_messages_deprecated:
-                    warnings.warn(
-                        "Passing a `messages`-format (conversational language-modeling) dataset to "
-                        "`DistillationTrainer` is deprecated and will be removed. Use a prompt-only dataset with a "
-                        "`prompt` column instead.",
-                        FutureWarning,
-                    )
-                    self._warned_messages_deprecated = True
-                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(
-                prompt_messages, tokenize=False, add_generation_prompt=True
+                example["prompt"], tokenize=False, add_generation_prompt=True
             )
-            prompt_ids = self.tokenizer(
-                formatted_prompt,
-                truncation=True,
-                max_length=self.max_prompt_length,
-                padding=False,
-                add_special_tokens=False,
-            )["input_ids"]
-
-            if has_completion:
-                # Tokenize the full message (prompt + completion) without truncation first
-                formatted_full = self.tokenizer.apply_chat_template(
-                    messages, tokenize=False, add_generation_prompt=False
-                )
-                full_ids = self.tokenizer(formatted_full, truncation=False, padding=False, add_special_tokens=False)[
-                    "input_ids"
-                ]
-
-                # Identify completion tokens: everything after the prompt in the full sequence.
-                # Use the un-truncated prompt length as the split point.
-                formatted_prompt_ids = self.tokenizer(
-                    formatted_prompt, truncation=False, padding=False, add_special_tokens=False
-                )["input_ids"]
-                completion_ids = full_ids[len(formatted_prompt_ids) :]
-
-                # Trim completion so prompt + completion <= max_length
-                max_comp = self.max_length - len(prompt_ids)
-                if max_comp > 0 and len(completion_ids) > max_comp:
-                    completion_ids = completion_ids[:max_comp]
-                elif max_comp <= 0:
-                    completion_ids = []
-
-                input_ids = prompt_ids + completion_ids
-                labels = [self.ignore_index] * len(prompt_ids) + list(completion_ids)
-            else:
-                # Prompt-only: no completion to train on (on-policy will generate one)
-                input_ids = list(prompt_ids)
-                labels = [self.ignore_index] * len(prompt_ids)
+            prompt_ids = self.tokenizer(formatted_prompt, padding=False, add_special_tokens=False)["input_ids"]
 
-            all_input_ids.append(input_ids)
-            all_labels.append(labels)
+            # Prompt-only: no completion to train on (on-policy will generate one)
+            all_input_ids.append(list(prompt_ids))
+            all_labels.append([self.ignore_index] * len(prompt_ids))
             all_prompt_ids.append(list(prompt_ids))
 
         # Convert to tensors and left-pad
@@ -433,11 +363,7 @@ def __init__(
 
         # ── Data collator ──
         if data_collator is None:
-            data_collator = _DistillationCollator(
-                tokenizer=processing_class,
-                max_length=args.max_length,
-                max_prompt_length=args.max_prompt_length,
-            )
+            data_collator = _DistillationCollator(tokenizer=processing_class)
 
         # ── Liger fused JSD loss ──
         self.use_liger_loss = False