[DistillationTrainer refactor] Accept a prompt column alongside messages

docs/source/paper_index.md

@@ -1634,21 +1634,21 @@ Papers relating to training a student model with the help of a teacher model.
 
 **📜 Paper**: https://huggingface.co/papers/2306.13649
 
-Introduces Generalized Knowledge Distillation (GKD), which addresses distribution mismatch in KD for auto-regressive models by training the student on its own generated outputs with teacher feedback, instead of a fixed set of sequences. GKD supports flexible loss functions (e.g. beyond KL when the student cannot match the teacher) and integrates with RL fine-tuning (RLHF). The paper reports results on summarization, translation, arithmetic reasoning, and instruction-tuning. Used in TRL via [`experimental.distillation.DistillationTrainer`] and [`experimental.gkd.GKDTrainer`]. To reproduce the paper's setting, use this configuration:
+Introduces Generalized Knowledge Distillation (GKD), which addresses distribution mismatch in KD for auto-regressive models by training the student on its own generated outputs with teacher feedback, instead of a fixed set of sequences. GKD supports flexible loss functions (e.g. beyond KL when the student cannot match the teacher) and integrates with RL fine-tuning (RLHF). The paper reports results on summarization, translation, arithmetic reasoning, and instruction-tuning. Used in TRL via [`experimental.gkd.GKDTrainer`], which exposes the paper's on/off-policy mixing (`lmbda`). [`experimental.distillation.DistillationTrainer`] implements the same generalized-JSD objective for the always-on-policy case. To reproduce the paper's setting, use this configuration:
 
 ```python
-from trl.experimental.distillation import DistillationConfig
+from trl.experimental.gkd import GKDConfig
 
 # XSum summarization task (Table A.1 of the paper)
-training_args = DistillationConfig(
+training_args = GKDConfig(
     lmbda=0.5,  # λ student data fraction (Section 3 of the paper)
     beta=0.5,  # β Generalized JSD interpolation, 0=KL, 1=reverse KL (Section 3 of the paper)
     temperature=1.0,  # student training temperature (Appendix A of the paper)
     max_steps=40000,  # training steps (Table A.1 of the paper)
     learning_rate=3e-4,  # learning rate (Table A.1 of the paper)
     per_device_train_batch_size=32,  # batch size (Table A.1 of the paper)
     warmup_steps=2000,  # warm-up steps (Table A.1 of the paper)
-    max_completion_length=64,  # max output tokens (Table A.1 of the paper)
+    max_new_tokens=64,  # max output tokens (Table A.1 of the paper)
 )
 ```
 

tests/experimental/test_distillation_trainer.py

@@ -198,7 +198,6 @@ def _make_args(self, **kwargs):
             "disable_tqdm": True,
             "use_cpu": True,
             "bf16": False,
-            "lmbda": 0.0,
             "max_length": 128,
             "max_completion_length": 32,
             "model_init_kwargs": {"dtype": "float32", "device_map": None},
@@ -254,17 +253,11 @@ def test_distillation_trainer_train_runs_with_local_teacher(self):
         assert train_result.metrics["train_loss"] >= -1e-4
         assert "model.safetensors" in os.listdir(self.tmp_dir + "/checkpoint-2")
 
-    @pytest.mark.parametrize("lmbda", [0.0, 1.0])
-    def test_train_updates_params_on_and_off_policy(self, lmbda):
-        """Pin both policy modes end to end before `lmbda` is removed.
-
-        `lmbda=0.0` trains on the dataset's own completions, `lmbda=1.0` on completions the student generates. The
-        trainer is scheduled to become always-on-policy, so the off-policy case is pinned here to make its removal a
-        deliberate deletion rather than a silent one.
-        """
+    def test_train_updates_params(self):
+        """Training is always on-policy: the student generates completions, the teacher scores them, params move."""
         # Higher lr than the default: gradients are tiny on this model and the default lr can stall the update, which
         # would make the assertion below vacuous.
-        trainer = self._make_local_trainer(lmbda=lmbda, max_steps=2, learning_rate=0.1)
+        trainer = self._make_local_trainer(max_steps=2, learning_rate=0.1)
 
         # Diverge the teacher from the student so the divergence (and thus the gradient) is well above fp noise; with
         # matched weights it would be ~0 and the update below could pass on noise alone.
@@ -281,6 +274,27 @@ def test_train_updates_params_on_and_off_policy(self, lmbda):
         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 "

tests/experimental/test_server_distillation_trainer.py

@@ -278,7 +278,6 @@ def _run_one_step(self, bs, ga, monkeypatch):
             teacher_model_server_url="http://fake-teacher.invalid:8000",
             loss_top_k=1,
             beta=1.0,
-            lmbda=0.0,
             loss_add_tail=True,
             save_strategy="no",
             report_to="none",

trl/experimental/distillation/distillation_config.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import warnings
 from dataclasses import dataclass, field
 from typing import Any
 
@@ -50,9 +49,6 @@ class DistillationConfig(_BaseConfig):
         temperature (`float`, *optional*, defaults to `1.0`):
             Temperature for sampling during generation and for computing the distillation loss. Higher values produce
             softer probability distributions.
-        lmbda (`float`, *optional*, defaults to `1.0`):
-            Probability of using on-policy (student-generated) data for each gradient accumulation slice. A value of
-            `0.0` means fully off-policy (dataset completions only), `1.0` means fully on-policy.
         beta (`float`, *optional*, defaults to `1.0`):
             Interpolation coefficient for the Generalized Jensen-Shannon Divergence loss. When `0.0`, the loss is the
             forward KL divergence. When `1.0`, the loss is the reverse KL divergence. When `0.5`, it is the standard
@@ -165,13 +161,6 @@ class DistillationConfig(_BaseConfig):
             "help": "Temperature for sampling and loss computation. Higher values produce softer distributions."
         },
     )
-    lmbda: float = field(
-        default=1.0,
-        metadata={
-            "help": "Probability of using on-policy (student-generated) data per gradient accumulation slice. "
-            "0.0 = fully off-policy, 1.0 = fully on-policy."
-        },
-    )
     beta: float = field(
         default=1.0,
         metadata={
@@ -314,8 +303,6 @@ class DistillationConfig(_BaseConfig):
     def __post_init__(self):
         super().__post_init__()
 
-        if self.lmbda < 0.0 or self.lmbda > 1.0:
-            raise ValueError(f"lmbda must be in [0.0, 1.0], got {self.lmbda}.")
         if self.beta < 0.0 or self.beta > 1.0:
             raise ValueError(f"beta must be in [0.0, 1.0], got {self.beta}.")
 
@@ -342,11 +329,3 @@ def __post_init__(self):
                 f"gradient_accumulation_steps. Got {self.generation_batch_size} * {self.num_generations} != "
                 f"{self.per_device_train_batch_size} * {self.gradient_accumulation_steps}."
             )
-
-        if self.num_generations > 1 and self.lmbda < 1.0:
-            warnings.warn(
-                f"num_generations={self.num_generations} with lmbda={self.lmbda} means off-policy batches include "
-                f"{self.num_generations} copies of each sample. Consider lmbda=1.0 when num_generations > 1.",
-                UserWarning,
-                stacklevel=2,
-            )