From 6983b6fc92c124f51ac99485f0792bf6d5ceee5a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Sun, 19 Jul 2026 19:45:10 +0000
Subject: [PATCH 01/44] Pin the distillation objective before refactoring

The trainer is about to go through a long refactor (top-k support removal,
then a switch to a chunked loss). The implementation is expected to change;
the value it computes is not.

Add three pinning tests:

- `generalized_jsd_loss` against a naive reference written straight from the
  definition. It builds the mixture in probability space rather than reusing
  `F.kl_div`'s inverted argument order and the `logsumexp` trick, so it can
  actually catch an argument-order or mixture-weight regression.
- parity with `GKDTrainer.generalized_jsd_loss`, which implements the same
  objective. This is the cross-trainer contract.
- the temperature path, which is scheduled to become sampling-only. Pinning it
  makes that change a visible diff rather than a silent drift.

All are seeded and cover beta in {0, 0.5, 1}, with and without a label mask.
The existing fixture uses an unseeded `torch.randn`, which cannot pin anything.
---
 .../experimental/test_distillation_trainer.py | 76 +++++++++++++++++++
 1 file changed, 76 insertions(+)

diff --git a/tests/experimental/test_distillation_trainer.py b/tests/experimental/test_distillation_trainer.py
index c3cd2d9443a..7642b112773 100644
--- a/tests/experimental/test_distillation_trainer.py
+++ b/tests/experimental/test_distillation_trainer.py
@@ -31,6 +31,7 @@
     _RepeatBatchDataLoader,
     build_teacher_request_inputs,
 )
+from trl.experimental.gkd.gkd_trainer import GKDTrainer
 
 from ..testing_utils import TrlTestCase, require_liger_kernel, require_torch_accelerator
 
@@ -367,6 +368,81 @@ def test_mask_keeps_forward_and_backward_finite(self):
         assert torch.isfinite(raw_student.grad).all()
 
 
+def _reference_generalized_jsd(student_logits, teacher_logits, labels=None, beta=0.5, temperature=1.0):
+    """Naive reference for the generalized JSD, written straight from the definition.
+
+    Deliberately independent of the implementation: probabilities are formed explicitly and the mixture is built in
+    probability space, so this does not share `F.kl_div`'s inverted argument order nor the `logsumexp` trick. That is
+    what makes it able to catch an argument-order or mixture-weight regression.
+    """
+    student_log_probs = torch.log_softmax(student_logits / temperature, dim=-1)
+    teacher_log_probs = torch.log_softmax(teacher_logits / temperature, dim=-1)
+    student_probs, teacher_probs = student_log_probs.exp(), teacher_log_probs.exp()
+
+    if beta == 0.0:  # forward KL: KL(teacher || student)
+        per_element = teacher_probs * (teacher_log_probs - student_log_probs)
+    elif beta == 1.0:  # reverse KL: KL(student || teacher)
+        per_element = student_probs * (student_log_probs - teacher_log_probs)
+    else:  # generalized JSD against the mixture M = (1 - beta) * student + beta * teacher
+        mixture_log_probs = ((1 - beta) * student_probs + beta * teacher_probs).log()
+        per_element = beta * (teacher_probs * (teacher_log_probs - mixture_log_probs)) + (1 - beta) * (
+            student_probs * (student_log_probs - mixture_log_probs)
+        )
+
+    if labels is None:  # "batchmean" without labels divides by the batch size
+        return per_element.sum() / max(per_element.size(0), 1)
+    mask = labels != -100
+    return per_element[mask].sum() / mask.sum().clamp(min=1)
+
+
+class TestGeneralizedJSDLossIsPinned(TrlTestCase):
+    """Pins the distillation objective while the trainer is refactored.
+
+    The implementation is expected to change (top-k support removal, then the switch to a chunked loss); the value it
+    computes is not. Any diff that moves these numbers is changing the objective and must say so.
+    """
+
+    def setup_method(self):
+        generator = torch.Generator().manual_seed(42)  # seeded: an unseeded fixture cannot pin anything
+        self.student_logits = torch.randn(2, 3, 5, generator=generator)
+        self.teacher_logits = torch.randn(2, 3, 5, generator=generator)
+        self.labels = torch.tensor([[-100, 1, 2], [-100, -100, 3]])
+
+    @pytest.mark.parametrize("beta", [0.0, 0.5, 1.0])
+    @pytest.mark.parametrize("use_labels", [False, True])
+    def test_matches_reference_implementation(self, beta, use_labels):
+        labels = self.labels if use_labels else None
+        loss = DistillationTrainer.generalized_jsd_loss(
+            self.student_logits, self.teacher_logits, labels=labels, beta=beta
+        )
+        expected = _reference_generalized_jsd(self.student_logits, self.teacher_logits, labels=labels, beta=beta)
+        torch.testing.assert_close(loss, expected)
+
+    @pytest.mark.parametrize("beta", [0.0, 0.5, 1.0])
+    @pytest.mark.parametrize("use_labels", [False, True])
+    def test_matches_gkd(self, beta, use_labels):
+        # GKD implements the same objective. Keeping the two in lockstep is the cross-trainer contract: if this breaks,
+        # either the promotion changed the objective or GKD drifted.
+        labels = self.labels if use_labels else None
+        loss = DistillationTrainer.generalized_jsd_loss(
+            self.student_logits, self.teacher_logits, labels=labels, beta=beta
+        )
+        gkd_loss = GKDTrainer.generalized_jsd_loss(self.student_logits, self.teacher_logits, labels=labels, beta=beta)
+        torch.testing.assert_close(loss, gkd_loss)
+
+    @pytest.mark.parametrize("beta", [0.0, 0.5, 1.0])
+    def test_temperature_matches_reference(self, beta):
+        # `temperature` is applied to the loss today. It is scheduled to become sampling-only, so pin it explicitly:
+        # that change must be a visible diff here, not a silent drift.
+        loss = DistillationTrainer.generalized_jsd_loss(
+            self.student_logits, self.teacher_logits, labels=self.labels, beta=beta, temperature=2.0
+        )
+        expected = _reference_generalized_jsd(
+            self.student_logits, self.teacher_logits, labels=self.labels, beta=beta, temperature=2.0
+        )
+        torch.testing.assert_close(loss, expected)
+
+
 class TestGeneralizedJSDLoss(TrlTestCase):
     def setup_method(self):
         self.batch_size = 2

From d443d3b841c9fe01f26ac9b6f699e54aa8a2c30e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Sun, 19 Jul 2026 19:46:19 +0000
Subject: [PATCH 02/44] Pin on-policy and off-policy training end to end

`lmbda` selects between training on the dataset's own completions
(`lmbda=0.0`) and on completions the student generates (`lmbda=1.0`). The
trainer is scheduled to become always-on-policy, so pin both modes now: the
later removal of the off-policy path should be a deliberate deletion with a
visibly failing test, not a silent behaviour change.

Asserts the loss is recorded and that every parameter actually moved, using
`torch.equal` rather than a tolerance so a stalled update cannot pass. Uses a
higher learning rate than the default because gradients are tiny on the test
model and the default rate can stall the update, which would make the
assertion vacuous.
---
 .../experimental/test_distillation_trainer.py | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/tests/experimental/test_distillation_trainer.py b/tests/experimental/test_distillation_trainer.py
index 7642b112773..40e0ec1d9b4 100644
--- a/tests/experimental/test_distillation_trainer.py
+++ b/tests/experimental/test_distillation_trainer.py
@@ -603,6 +603,25 @@ def test_distillation_trainer_train_runs_with_local_teacher(self):
         assert train_result.metrics["train_loss"] >= 0.0
         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.
+        """
+        # 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)
+        previous_params = {name: param.clone() for name, param in trainer.model.named_parameters()}
+
+        trainer.train()
+
+        assert trainer.state.log_history[-1]["train_loss"] is not None
+        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.parametrize(
         "eval_dataset_type",
         [

From a8bb14caa6fd11b1653d4890ccccf6398b161665 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Sun, 19 Jul 2026 19:47:57 +0000
Subject: [PATCH 03/44] Document the num_items_in_batch denominator bug

The divergence loss is reduced as `sum / num_items_in_batch`, so that count is
the loss denominator. It is currently wrong.

`transformers.Trainer` computes it from the raw dataloader batches, before
`_prepare_inputs` runs. Two things follow. `_RepeatBatchDataLoader` yields the
same generation batch once per accumulation step, so the count is
`gradient_accumulation_steps` times too large. And it is derived from the
dataset labels, even on steps whose completions are replaced by generated ones.

The effect is a silently down-scaled loss and gradient. Add an xfail test that
records what the trainer is actually handed during `train()` and compares it
with the tokens the loss is summed over. The existing coverage passes
`num_items_in_batch` explicitly, so it exercises the reduction rather than the
value, and cannot catch this.

Un-xfail when the generation buffer moves to the GRPO-style `_prepare_inputs`
and the denominator is derived from the completions actually trained on.
---
 .../experimental/test_distillation_trainer.py | 37 +++++++++++++++++++
 1 file changed, 37 insertions(+)

diff --git a/tests/experimental/test_distillation_trainer.py b/tests/experimental/test_distillation_trainer.py
index 40e0ec1d9b4..439a9cc2b0e 100644
--- a/tests/experimental/test_distillation_trainer.py
+++ b/tests/experimental/test_distillation_trainer.py
@@ -622,6 +622,43 @@ 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(
+        reason="num_items_in_batch is computed from the raw dataloader batches, before _prepare_inputs runs. "
+        "_RepeatBatchDataLoader yields the same generation batch once per accumulation step, so the count is "
+        "gradient_accumulation_steps times too large, and it is derived from the dataset labels even on steps whose "
+        "completions are replaced by generation. Fixed when the buffer moves to the GRPO-style _prepare_inputs."
+    )
+    def test_num_items_in_batch_counts_the_tokens_trained_on(self):
+        """`num_items_in_batch` is the loss denominator, so it must count the tokens actually trained on.
+
+        The loss is reduced as `sum / num_items_in_batch`, so an inflated count silently scales the loss — and the
+        gradient — down. Documented here as a known failure; un-xfail it in the PR that fixes the denominator.
+        """
+        recorded = []
+
+        class _RecordingTrainer(DistillationTrainer):
+            def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
+                # Prompt positions are already -100, so this is the completion-token count either way.
+                recorded.append((num_items_in_batch, int((inputs["labels"] != -100).sum())))
+                return super().compute_loss(model, inputs, return_outputs, num_items_in_batch)
+
+        dataset = load_dataset("trl-internal-testing/zen", "conversational_language_modeling", split="train")
+        trainer = _RecordingTrainer(
+            model=self.model_id,
+            teacher_model=self.model_id,
+            args=self._make_args(gradient_accumulation_steps=2, max_steps=1),
+            train_dataset=dataset,
+            processing_class=self.tokenizer,
+        )
+
+        trainer.train()
+
+        assert len(recorded) == 2, "expected one compute_loss call per accumulation step"
+        reported = recorded[0][0]
+        assert reported is not None, "transformers did not pass num_items_in_batch"
+        # The denominator should be the number of tokens summed over the accumulation window.
+        assert int(reported) == sum(trained_on for _, trained_on in recorded)
+
     @pytest.mark.parametrize(
         "eval_dataset_type",
         [

From 98eef79e8939d5527d7fc2871df58f6c2884ad9e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Sun, 19 Jul 2026 19:55:33 +0000
Subject: [PATCH 04/44] Move IW-OPD to a dedicated experimental trainer
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

IW-OPD (Importance-Weighted On-Policy Distillation, 2606.22600) is a
teacher-guided policy-gradient method: sampled-token advantage
`A_t = sg[log π_T − log π_0]`, a prefix-decayed importance weight, and a
clipped-surrogate PPO inner loop. It fits neither the stable full-vocabulary
`DistillationTrainer` (supervised KD, no advantages) nor `GRPOTrainer` (which
has no teacher). So rather than delete it, move it to its own home.

- Fork the pre-refactor trainer/config into `trl/experimental/iw_opd/` as
  `IWOPDTrainer` / `IWOPDConfig` — a self-contained snapshot that keeps IW-OPD
  (and, incidentally, the other pre-refactor features) working. Its docstring
  states plainly that it is a frozen, unmaintained snapshot.
- Remove IW-OPD from the base `DistillationTrainer` (loss, config fields,
  rollout-logprob plumbing, the `logprobs=0` vLLM request, and the tests).
- Repoint the IW-OPD paper-index entry to `experimental.iw_opd` (the GKD and
  general on-policy-distillation entries stay on `DistillationTrainer`).
- Register `IWOPDTrainer` in the telemetry allowlist; the moved tests live in
  `tests/experimental/test_iw_opd_trainer.py`.

This merges what were two PRs (remove IW-OPD + drop its paper entry) into one
"move" so nothing is lost. `compute_loss` in the base still drops from 7
dispatch paths to 5.
---
 docs/source/paper_index.md                    |   10 +-
 .../experimental/test_distillation_trainer.py |  204 --
 tests/experimental/test_iw_opd_trainer.py     |  899 ++++++++
 .../distillation/distillation_config.py       |   47 -
 .../distillation/distillation_trainer.py      |  136 +-
 trl/experimental/iw_opd/__init__.py           |   19 +
 trl/experimental/iw_opd/iw_opd_config.py      |  501 +++++
 trl/experimental/iw_opd/iw_opd_trainer.py     | 1911 +++++++++++++++++
 trl/trainer/base_trainer.py                   |    1 +
 9 files changed, 3342 insertions(+), 386 deletions(-)
 create mode 100644 tests/experimental/test_iw_opd_trainer.py
 create mode 100644 trl/experimental/iw_opd/__init__.py
 create mode 100644 trl/experimental/iw_opd/iw_opd_config.py
 create mode 100644 trl/experimental/iw_opd/iw_opd_trainer.py

diff --git a/docs/source/paper_index.md b/docs/source/paper_index.md
index a8fb3ae782f..5e815a8a86f 100644
--- a/docs/source/paper_index.md
+++ b/docs/source/paper_index.md
@@ -1687,15 +1687,15 @@ training_args = DistillationConfig(
 
 **📜 Paper**: https://huggingface.co/papers/2606.22600
 
-Introduces Importance-Weighted On-Policy Distillation (IW-OPD), which addresses the position bias in OPD by reweighting sampled-token distillation updates according to accumulated teacher-student prefix discrepancy. Early tokens keep larger weights, while later tokens after high drift are downweighted. Used in TRL via [`experimental.distillation.DistillationTrainer`] with `distillation_objective="iw_opd"`.
+Introduces Importance-Weighted On-Policy Distillation (IW-OPD), which addresses the position bias in OPD by reweighting sampled-token distillation updates according to accumulated teacher-student prefix discrepancy. Early tokens keep larger weights, while later tokens after high drift are downweighted. Used in TRL via [`experimental.iw_opd.IWOPDTrainer`] with `distillation_objective="iw_opd"`.
 
-The paper reports its main experiments with a verl PPO trainer and vLLM rollouts. `DistillationTrainer` exposes the matching distillation and rollout settings below; PPO-specific settings from the paper such as clipping range `0.2`, dual-clip constant `3.0`, PPO epochs, entropy coefficient, KL reward penalty, auxiliary KL, and rollout importance correction are not `DistillationConfig` parameters.
+The paper optimizes IW-OPD with a clipped policy-gradient setup (verl) and vLLM rollouts. `IWOPDTrainer` exposes the matching distillation and rollout settings below; policy-optimization settings from the paper such as clipping range `0.2`, dual-clip constant `3.0`, inner PPO epochs, entropy coefficient, KL reward penalty, auxiliary KL, and rollout importance correction are not `IWOPDConfig` parameters.
 
 ```python
-from trl.experimental.distillation import DistillationConfig
+from trl.experimental.iw_opd import IWOPDConfig
 
-# Table 6 and Algorithm 1 of the paper, mapped to DistillationConfig where available.
-training_args = DistillationConfig(
+# Table 6 and Algorithm 1 of the paper, mapped to IWOPDConfig where available.
+training_args = IWOPDConfig(
     distillation_objective="iw_opd",
     iw_opd_gamma=0.5,  # γ amplification, Algorithm 1 and Appendix C.3
     lmbda=1.0,  # fully on-policy rollouts
diff --git a/tests/experimental/test_distillation_trainer.py b/tests/experimental/test_distillation_trainer.py
index 439a9cc2b0e..d9f5b1bf89a 100644
--- a/tests/experimental/test_distillation_trainer.py
+++ b/tests/experimental/test_distillation_trainer.py
@@ -14,8 +14,6 @@
 
 import math
 import os
-from collections import defaultdict
-from types import SimpleNamespace
 from unittest.mock import MagicMock
 
 import pytest
@@ -154,50 +152,6 @@ def test_distillation_config_rejects_invalid_reverse_kl_top_1_mode(tmp_path):
         DistillationConfig(**_make_distillation_config_kwargs(tmp_path), reverse_kl_top_1_mode="invalid")
 
 
-def test_distillation_config_rejects_invalid_distillation_objective(tmp_path):
-    with pytest.raises(ValueError, match="distillation_objective must be one of"):
-        DistillationConfig(**_make_distillation_config_kwargs(tmp_path), distillation_objective="invalid")
-
-
-def test_distillation_config_rejects_invalid_iw_opd_gamma(tmp_path):
-    with pytest.raises(ValueError, match="iw_opd_gamma must be non-negative"):
-        DistillationConfig(**_make_distillation_config_kwargs(tmp_path), iw_opd_gamma=-0.1)
-
-
-def test_distillation_config_rejects_invalid_iw_opd_epsilon(tmp_path):
-    with pytest.raises(ValueError, match="iw_opd_epsilon must be positive"):
-        DistillationConfig(**_make_distillation_config_kwargs(tmp_path), iw_opd_epsilon=0.0)
-
-
-def test_distillation_config_rejects_iw_opd_with_liger(tmp_path):
-    with pytest.raises(
-        ValueError, match="use_liger_kernel=True is not supported with distillation_objective='iw_opd'"
-    ):
-        DistillationConfig(
-            **_make_distillation_config_kwargs(tmp_path),
-            distillation_objective="iw_opd",
-            use_liger_kernel=True,
-        )
-
-
-def test_distillation_config_rejects_iw_opd_when_not_fully_on_policy(tmp_path):
-    with pytest.raises(ValueError, match="distillation_objective='iw_opd' requires lmbda=1.0"):
-        DistillationConfig(
-            **_make_distillation_config_kwargs(tmp_path),
-            distillation_objective="iw_opd",
-            lmbda=0.5,
-        )
-
-
-def test_distillation_config_rejects_iw_opd_with_argmax_top_1_mode(tmp_path):
-    with pytest.raises(ValueError, match="distillation_objective='iw_opd' requires reverse_kl_top_1_mode='sampled'"):
-        DistillationConfig(
-            **_make_distillation_config_kwargs(tmp_path),
-            distillation_objective="iw_opd",
-            reverse_kl_top_1_mode="argmax",
-        )
-
-
 def test_distillation_config_rejects_teacher_server_with_reverse_kl_argmax(tmp_path):
     with pytest.raises(ValueError, match="reverse_kl_top_1_mode='argmax' is not supported"):
         DistillationConfig(
@@ -800,164 +754,6 @@ def test_sampled_mode_matches_between_local_and_external_teachers(self, monkeypa
         assert teacher_client.calls[0]["top_logprobs"] == 1
         torch.testing.assert_close(local_loss, server_loss)
 
-    def test_iw_opd_prefix_weights_and_loss(self):
-        trainer = DistillationTrainer.__new__(DistillationTrainer)
-        trainer.temperature = 1.0
-        trainer.iw_opd_gamma = 0.5
-        trainer.iw_opd_epsilon = 1e-8
-        trainer._metrics = {"train": defaultdict(list), "eval": defaultdict(list)}
-        trainer.model = SimpleNamespace(training=True)
-
-        student_logits = torch.zeros(2, 3, 4, requires_grad=True)
-        completion_tokens = torch.tensor([[0, 1, 2], [1, 2, 3]])
-        labels = torch.tensor([[0, 1, 2], [1, -100, -100]])
-        student_logprob = -math.log(4)
-        teacher_actual_logprobs = torch.tensor(
-            [
-                [student_logprob + 0.2, student_logprob - 0.3, student_logprob + 0.5],
-                [student_logprob + 0.4, float("-inf"), float("-inf")],
-            ]
-        )
-
-        loss = trainer._compute_iw_opd_loss(
-            student_logits=student_logits,
-            completion_tokens=completion_tokens,
-            labels=labels,
-            teacher_actual_logprobs=teacher_actual_logprobs,
-        )
-
-        advantages = torch.tensor([[0.2, -0.3, 0.5], [0.4, 0.0, 0.0]])
-        weights = torch.tensor([[1.5, 1.4, 1.25], [1.5, 1.0, 1.0]])
-        expected = -student_logprob * (weights * advantages).sum() / 4
-
-        torch.testing.assert_close(loss, expected)
-        assert trainer._metrics["train"]["iw_opd/max_weight"] == [1.5]
-        assert trainer._metrics["train"]["iw_opd/min_weight"] == [1.25]
-
-        # Under gradient accumulation the sum is divided by the global token count instead of the local one
-        loss_ga = trainer._compute_iw_opd_loss(
-            student_logits=student_logits,
-            completion_tokens=completion_tokens,
-            labels=labels,
-            teacher_actual_logprobs=teacher_actual_logprobs,
-            num_items_in_batch=8,
-        )
-        torch.testing.assert_close(loss_ga, expected * 4 / 8)
-
-        loss.backward()
-        assert student_logits.grad is not None
-        assert torch.isfinite(student_logits.grad).all()
-
-    def test_iw_opd_uses_cached_rollout_logprobs(self):
-        trainer = DistillationTrainer.__new__(DistillationTrainer)
-        trainer.temperature = 1.0
-        trainer.iw_opd_gamma = 0.0
-        trainer.iw_opd_epsilon = 1e-8
-        trainer._metrics = {"train": defaultdict(list), "eval": defaultdict(list)}
-        trainer.model = SimpleNamespace(training=True)
-
-        student_logits = torch.zeros(1, 2, 4, requires_grad=True)
-        completion_tokens = torch.tensor([[0, 1]])
-        labels = torch.tensor([[0, 1]])
-        current_student_logprob = -math.log(4)
-        rollout_logprobs = torch.tensor([[-0.5, -0.25]])
-        teacher_actual_logprobs = torch.tensor([[-0.2, -0.05]])
-
-        loss = trainer._compute_iw_opd_loss(
-            student_logits=student_logits,
-            completion_tokens=completion_tokens,
-            labels=labels,
-            teacher_actual_logprobs=teacher_actual_logprobs,
-            rollout_logprobs=rollout_logprobs,
-        )
-
-        expected_advantage_sum = (teacher_actual_logprobs - rollout_logprobs).sum()
-        expected = -current_student_logprob * expected_advantage_sum / 2
-        torch.testing.assert_close(loss, expected)
-
-        # Buffer layout: with mixed prompt lengths, prompt_length (shortest real prompt) is smaller than the
-        # padded prompt width, so rollout logprobs must be stored at full sequence width to survive the
-        # labels-style slice in compute_loss.
-        class _Tok:
-            pad_token_id = 0
-
-            @staticmethod
-            def batch_decode(sequences, skip_special_tokens=False, clean_up_tokenization_spaces=False):
-                return [" ".join(str(int(t)) for t in seq) for seq in sequences]
-
-        trainer = DistillationTrainer.__new__(DistillationTrainer)
-        trainer.accelerator = SimpleNamespace(device=torch.device("cpu"))
-        trainer.processing_class = _Tok()
-        trainer.generation_config = SimpleNamespace(max_new_tokens=3)
-        trainer._buffered_inputs = [None]
-        trainer._buffered_text_logs = [None]
-
-        trainer._store_completions_in_buffer(
-            slices=[{}],
-            on_policy_indices=[0],
-            local_slice_indices=[0, 0],
-            local_prompts=[torch.tensor([5]), torch.tensor([6, 7])],
-            completion_ids=[[11, 12, 13], [14]],
-            logprobs=[[[-0.1], [-0.2], [-0.3]], [[-0.4]]],
-            logprob_token_ids=[[[11], [12], [13]], [[14]]],
-        )
-
-        buffered = trainer._buffered_inputs[0]
-        labels = buffered["labels"]
-        rollout_logprobs = buffered["rollout_logprobs"]
-        assert rollout_logprobs.shape == labels.shape
-
-        prompt_length = trainer._compute_prompt_length(buffered)
-        assert prompt_length == 1
-        sliced_labels = labels[:, prompt_length:]
-        sliced_rollout = rollout_logprobs[:, prompt_length:]
-        expected = {(0, 11): -0.1, (0, 12): -0.2, (0, 13): -0.3, (1, 14): -0.4}
-        valid = sliced_labels != -100
-        assert valid.sum() == 4
-        for row, col in valid.nonzero().tolist():
-            token = int(sliced_labels[row, col])
-            torch.testing.assert_close(sliced_rollout[row, col].item(), expected[(row, token)])
-
-    def test_iw_opd_train_runs_with_local_teacher(self):
-        trainer = self._make_local_trainer(
-            distillation_objective="iw_opd",
-            lmbda=1.0,
-            max_completion_length=8,
-        )
-
-        train_result = trainer.train()
-
-        assert train_result.metrics["train_loss"] is not None
-        assert math.isfinite(train_result.metrics["train_loss"])
-        assert any("iw_opd/mean_weight" in record for record in trainer.state.log_history)
-
-    def test_iw_opd_server_path_uses_actual_logprobs(self, monkeypatch):
-        import trl.generation.vllm_client as vllm_client_module
-
-        teacher_client = RecordingTeacherClient()
-        monkeypatch.setattr(vllm_client_module, "VLLMClient", lambda *args, **kwargs: teacher_client)
-
-        trainer = self._make_server_trainer(distillation_objective="iw_opd", lmbda=1.0)
-        inputs = self._move_batch_to_device(self._make_batch(trainer), trainer.accelerator.device)
-        sequences, prompt_lengths, _ = build_teacher_request_inputs(
-            inputs["input_ids"],
-            inputs["attention_mask"],
-            prompt_attention_mask=inputs["prompt_attention_mask"],
-            labels=inputs["labels"],
-        )
-        teacher_client.result = _canned_teacher_logprobs(
-            sequences=sequences,
-            prompt_lengths=prompt_lengths,
-            top_logprobs=1,
-        )
-
-        loss = trainer.compute_loss(trainer.model, inputs)
-
-        assert torch.isfinite(loss)
-        assert teacher_client.calls[0]["top_logprobs"] == 1
-
-
-class TestDistillationTrainerServerPath(TrlTestCase):
     @classmethod
     def setup_class(cls):
         model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
diff --git a/tests/experimental/test_iw_opd_trainer.py b/tests/experimental/test_iw_opd_trainer.py
new file mode 100644
index 00000000000..0a5d67c0ecc
--- /dev/null
+++ b/tests/experimental/test_iw_opd_trainer.py
@@ -0,0 +1,899 @@
+# Copyright 2020-2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+import os
+from collections import defaultdict
+from types import SimpleNamespace
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+import torch.nn.functional as F
+from datasets import Dataset, DatasetDict, IterableDatasetDict, load_dataset
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from trl.experimental.iw_opd import IWOPDConfig, IWOPDTrainer
+from trl.experimental.iw_opd.iw_opd_trainer import (
+    _add_tail_bucket,
+    _jsd_divergence,
+    _RepeatBatchDataLoader,
+    build_teacher_request_inputs,
+)
+
+from ..testing_utils import TrlTestCase, require_liger_kernel, require_torch_accelerator
+
+
+def _make_distillation_config_kwargs(tmp_path):
+    return {"output_dir": str(tmp_path), "report_to": "none", "use_cpu": True, "bf16": False}
+
+
+def _build_server_result(teacher_logits, inputs, temperature=1.0):
+    """Simulate a vLLM server response with variable-length per-sample completions."""
+    _, _, completion_lengths = build_teacher_request_inputs(
+        inputs["input_ids"],
+        inputs["attention_mask"],
+        prompt_attention_mask=inputs.get("prompt_attention_mask"),
+        labels=inputs.get("labels"),
+    )
+
+    label_mask = inputs["labels"] != -100
+    actual_logprobs = []
+    logprobs = []
+    logprob_token_ids = []
+
+    for i, comp_len in enumerate(completion_lengths):
+        if comp_len == 0:
+            actual_logprobs.append([])
+            logprobs.append([])
+            logprob_token_ids.append([])
+            continue
+
+        comp_start = int(torch.nonzero(label_mask[i], as_tuple=False)[0].item())
+        sample_logits = teacher_logits[i, comp_start - 1 : comp_start - 1 + comp_len, :]
+        sample_log_probs = F.log_softmax(sample_logits / temperature, dim=-1)
+        comp_tokens = inputs["input_ids"][i, comp_start : comp_start + comp_len]
+
+        top1_ids = sample_logits.argmax(dim=-1, keepdim=True)
+        top1_lps = sample_log_probs.gather(dim=-1, index=top1_ids)
+        actual_lps = sample_log_probs.gather(dim=-1, index=comp_tokens.unsqueeze(-1))
+
+        actual_logprobs.append(actual_lps.tolist())
+        logprobs.append(top1_lps.tolist())
+        logprob_token_ids.append(top1_ids.tolist())
+
+    return {
+        "actual_logprobs": actual_logprobs,
+        "logprobs": logprobs,
+        "logprob_token_ids": logprob_token_ids,
+    }
+
+
+class RecordingTeacherClient:
+    def __init__(self):
+        self.calls = []
+        self.result = None
+
+    def get_sequence_logprobs(self, sequences, prompt_lengths, top_logprobs, temperature):
+        self.calls.append(
+            {
+                "sequences": sequences,
+                "prompt_lengths": prompt_lengths,
+                "top_logprobs": top_logprobs,
+                "temperature": temperature,
+            }
+        )
+        return self.result
+
+
+def _ragged_server_response():
+    # Two samples with completion lengths 1 and 3 respectively; matches the wire format
+    # of VLLMClient.get_sequence_logprobs (per-sample shape (comp_len, top_k=1)).
+    return {
+        "logprobs": [[[-2.3]], [[-1.1], [-0.4], [-3.0]]],
+        "logprob_token_ids": [[[90]], [[90], [9217], [100]]],
+        "actual_logprobs": [[[-2.3]], [[-1.1], [-0.4], [-3.0]]],
+    }
+
+
+def _canned_teacher_logprobs(**kwargs):
+    # Fabricate ragged per-sample logprobs matching the requested sequence shapes.
+    sequences = kwargs["sequences"]
+    prompt_lengths = kwargs["prompt_lengths"]
+    top_k = kwargs.get("top_logprobs", 1)
+    logprobs, token_ids, actual = [], [], []
+    for seq, plen in zip(sequences, prompt_lengths, strict=True):
+        comp_len = len(seq) - plen
+        logprobs.append([[-1.0 - 0.05 * i] * top_k for i in range(comp_len)])
+        token_ids.append([[int(seq[plen + i])] * top_k for i in range(comp_len)])
+        actual.append([[-1.0 - 0.05 * i] for i in range(comp_len)])
+    return {"logprobs": logprobs, "logprob_token_ids": token_ids, "actual_logprobs": actual}
+
+
+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.",
+                    },
+                ]
+            },
+        ]
+    )
+
+
+def test_distillation_config_rejects_liger_with_teacher_server(tmp_path):
+    with pytest.raises(ValueError, match="use_liger_kernel=True is not supported with use_teacher_server=True"):
+        IWOPDConfig(
+            **_make_distillation_config_kwargs(tmp_path),
+            use_teacher_server=True,
+            teacher_model_server_url="http://localhost:8000",
+            use_liger_kernel=True,
+        )
+
+
+def test_distillation_config_rejects_invalid_reverse_kl_top_1_mode(tmp_path):
+    with pytest.raises(ValueError, match="reverse_kl_top_1_mode must be one of"):
+        IWOPDConfig(**_make_distillation_config_kwargs(tmp_path), reverse_kl_top_1_mode="invalid")
+
+
+def test_distillation_config_rejects_invalid_distillation_objective(tmp_path):
+    with pytest.raises(ValueError, match="distillation_objective must be one of"):
+        IWOPDConfig(**_make_distillation_config_kwargs(tmp_path), distillation_objective="invalid")
+
+
+def test_distillation_config_rejects_invalid_iw_opd_gamma(tmp_path):
+    with pytest.raises(ValueError, match="iw_opd_gamma must be non-negative"):
+        IWOPDConfig(**_make_distillation_config_kwargs(tmp_path), iw_opd_gamma=-0.1)
+
+
+def test_distillation_config_rejects_invalid_iw_opd_epsilon(tmp_path):
+    with pytest.raises(ValueError, match="iw_opd_epsilon must be positive"):
+        IWOPDConfig(**_make_distillation_config_kwargs(tmp_path), iw_opd_epsilon=0.0)
+
+
+def test_distillation_config_rejects_iw_opd_with_liger(tmp_path):
+    with pytest.raises(
+        ValueError, match="use_liger_kernel=True is not supported with distillation_objective='iw_opd'"
+    ):
+        IWOPDConfig(
+            **_make_distillation_config_kwargs(tmp_path),
+            distillation_objective="iw_opd",
+            use_liger_kernel=True,
+        )
+
+
+def test_distillation_config_rejects_iw_opd_when_not_fully_on_policy(tmp_path):
+    with pytest.raises(ValueError, match="distillation_objective='iw_opd' requires lmbda=1.0"):
+        IWOPDConfig(
+            **_make_distillation_config_kwargs(tmp_path),
+            distillation_objective="iw_opd",
+            lmbda=0.5,
+        )
+
+
+def test_distillation_config_rejects_iw_opd_with_argmax_top_1_mode(tmp_path):
+    with pytest.raises(ValueError, match="distillation_objective='iw_opd' requires reverse_kl_top_1_mode='sampled'"):
+        IWOPDConfig(
+            **_make_distillation_config_kwargs(tmp_path),
+            distillation_objective="iw_opd",
+            reverse_kl_top_1_mode="argmax",
+        )
+
+
+def test_distillation_config_rejects_teacher_server_with_reverse_kl_argmax(tmp_path):
+    with pytest.raises(ValueError, match="reverse_kl_top_1_mode='argmax' is not supported"):
+        IWOPDConfig(
+            **_make_distillation_config_kwargs(tmp_path),
+            use_teacher_server=True,
+            teacher_model_server_url="http://localhost:8000",
+            reverse_kl_top_1_mode="argmax",
+        )
+
+
+def test_distillation_config_rejects_teacher_server_mixed_loss_without_top_1(tmp_path):
+    with pytest.raises(ValueError, match="loss_top_k must be 1 when using use_teacher_server=True with beta>0"):
+        IWOPDConfig(
+            **_make_distillation_config_kwargs(tmp_path),
+            use_teacher_server=True,
+            teacher_model_server_url="http://localhost:8000",
+            beta=0.5,
+            loss_top_k=2,
+        )
+
+
+def test_distillation_config_requires_teacher_server_url(tmp_path):
+    with pytest.raises(ValueError, match="teacher_model_server_url must be set when use_teacher_server=True"):
+        IWOPDConfig(
+            **_make_distillation_config_kwargs(tmp_path),
+            use_teacher_server=True,
+            beta=0.5,
+            loss_top_k=1,
+        )
+
+
+@pytest.mark.parametrize(
+    (
+        "input_ids",
+        "attention_mask",
+        "prompt_attention_mask",
+        "labels",
+        "expected_sequences",
+        "expected_prompt_lengths",
+        "expected_completion_lengths",
+    ),
+    [
+        pytest.param(
+            torch.tensor([[0, 0, 10, 11, 12, 20, 21], [30, 31, 32, 33, 34, 40, 41]], dtype=torch.long),
+            torch.tensor([[0, 0, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1]], dtype=torch.long),
+            None,
+            torch.tensor(
+                [[-100, -100, -100, -100, -100, 20, 21], [-100, -100, -100, -100, -100, 40, 41]],
+                dtype=torch.long,
+            ),
+            [[10, 11, 12, 20, 21], [30, 31, 32, 33, 34, 40, 41]],
+            [3, 5],
+            [2, 2],
+            id="variable_prompt_lengths",
+        ),
+        pytest.param(
+            torch.tensor(
+                [[0, 0, 10, 11, 12, 20, 21, 0, 0], [30, 31, 32, 33, 34, 40, 41, 0, 0]],
+                dtype=torch.long,
+            ),
+            torch.tensor(
+                [[0, 0, 1, 1, 1, 1, 1, 0, 0], [1, 1, 1, 1, 1, 1, 1, 0, 0]],
+                dtype=torch.long,
+            ),
+            torch.tensor([[0, 0, 1, 1, 1], [1, 1, 1, 1, 1]], dtype=torch.long),
+            torch.tensor(
+                [
+                    [-100, -100, -100, -100, -100, 20, 21, -100, -100],
+                    [-100, -100, -100, -100, -100, 40, 41, -100, -100],
+                ],
+                dtype=torch.long,
+            ),
+            [[10, 11, 12, 20, 21], [30, 31, 32, 33, 34, 40, 41]],
+            [3, 5],
+            [2, 2],
+            id="padded_on_policy_completions",
+        ),
+        pytest.param(
+            torch.tensor([[10, 11, 12]], dtype=torch.long),
+            torch.tensor([[1, 1, 1]], dtype=torch.long),
+            None,
+            torch.tensor([[-100, -100, -100]], dtype=torch.long),
+            [[10, 11, 12]],
+            [3],
+            [0],
+            id="empty_completion",
+        ),
+    ],
+)
+def test_build_teacher_request_inputs(
+    input_ids,
+    attention_mask,
+    prompt_attention_mask,
+    labels,
+    expected_sequences,
+    expected_prompt_lengths,
+    expected_completion_lengths,
+):
+    sequences, prompt_lengths, completion_lengths = build_teacher_request_inputs(
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        prompt_attention_mask=prompt_attention_mask,
+        labels=labels,
+    )
+
+    assert sequences == expected_sequences
+    assert prompt_lengths == expected_prompt_lengths
+    assert completion_lengths == expected_completion_lengths
+
+
+class TestGetTeacherTokenLogprobsFromServer(TrlTestCase):
+    def test_variable_lengths_use_neg_inf_sentinel_at_padding(self):
+        mock_self = MagicMock()
+        mock_self.teacher_client.get_sequence_logprobs = MagicMock(return_value=_ragged_server_response())
+        mock_self.loss_top_k = 1
+        mock_self.temperature = 1.0
+
+        inputs = {
+            "input_ids": torch.tensor([[10, 11, 90, 0, 0], [10, 11, 90, 9217, 100]]),
+            "attention_mask": torch.tensor([[1, 1, 1, 0, 0], [1, 1, 1, 1, 1]]),
+            "labels": torch.tensor([[-100, -100, 90, -100, -100], [-100, -100, 90, 9217, 100]]),
+        }
+
+        out = IWOPDTrainer._get_teacher_token_logprobs_from_server(mock_self, inputs, aligned_prompt_length=2)
+
+        assert out["actual_logprobs"].shape == (2, 3)
+        assert out["topk_logprobs"].shape == (2, 3, 1)
+
+        # Real completion positions preserved.
+        assert out["actual_logprobs"][0, 0].item() == pytest.approx(-2.3, rel=1e-5)
+        assert out["actual_logprobs"][1, 0].item() == pytest.approx(-1.1, rel=1e-5)
+        assert out["actual_logprobs"][1, 2].item() == pytest.approx(-3.0, rel=1e-5)
+
+        # Sample 0 is 1 token long; positions 1 and 2 are padded with the -inf sentinel.
+        assert out["actual_logprobs"][0, 1].item() == float("-inf")
+        assert out["actual_logprobs"][0, 2].item() == float("-inf")
+        assert out["topk_logprobs"][0, 1, 0].item() == float("-inf")
+
+        # Sample 1 is full-length and fully finite.
+        assert torch.isfinite(out["actual_logprobs"][1, :]).all()
+
+
+class TestServerReverseKLPaddingMask(TrlTestCase):
+    def test_mask_keeps_forward_and_backward_finite(self):
+        # Simulates the getter's output: sample 0 has completion length 1 (positions 1-2
+        # padded with -inf), sample 1 is full-length.
+        teacher_topk = torch.tensor(
+            [[[-2.3], [float("-inf")], [float("-inf")]], [[-1.1], [-0.4], [-3.0]]],
+            dtype=torch.float32,
+        )
+        labels = torch.tensor([[90, -100, -100], [90, 9217, 100]])
+
+        # Strategy B: neutralise -inf at labels == -100 before the divergence math.
+        pad_mask = (labels == -100).unsqueeze(-1)
+        zero = torch.zeros((), dtype=teacher_topk.dtype)
+        teacher_topk = torch.where(pad_mask, zero, teacher_topk)
+
+        valid_mask = torch.ones_like(teacher_topk, dtype=torch.bool)
+        teacher_with_tail, support_mask = _add_tail_bucket(teacher_topk, valid_mask)
+        assert torch.isfinite(teacher_with_tail).all()
+
+        raw_student = torch.randn(2, 3, 2, requires_grad=True)
+        student_log_probs = F.log_softmax(raw_student, dim=-1)
+        loss = _jsd_divergence(student_log_probs, teacher_with_tail, beta=1.0, support_mask=support_mask)
+        assert torch.isfinite(loss).all()
+
+        loss.sum().backward()
+        assert torch.isfinite(raw_student.grad).all()
+
+
+class TestGeneralizedJSDLoss(TrlTestCase):
+    def setup_method(self):
+        self.batch_size = 2
+        self.seq_length = 3
+        self.vocab_size = 5
+        self.student_logits = torch.randn(self.batch_size, self.seq_length, self.vocab_size)
+        self.teacher_logits = torch.randn(self.batch_size, self.seq_length, self.vocab_size)
+
+    def test_uniform_distribution(self):
+        logits = torch.ones(1, 1, self.vocab_size)
+        loss = IWOPDTrainer.generalized_jsd_loss(logits, logits)
+        assert round(abs(loss.item() - 0), 5) == 0
+
+    def test_generalized_jsd_loss_edge_cases(self):
+        # Setup
+        student_logits = torch.log(torch.tensor([[0.1, 0.9]])).unsqueeze(0)
+        teacher_logits = torch.log(torch.tensor([[0.9, 0.1]])).unsqueeze(0)
+
+        # Case 1: beta = 1 (should be equivalent to KL(student || teacher))
+        loss_beta_1 = IWOPDTrainer.generalized_jsd_loss(student_logits, teacher_logits, beta=1)
+        expected_loss_beta_1 = F.kl_div(
+            F.log_softmax(teacher_logits, dim=-1), F.softmax(student_logits, dim=-1), reduction="batchmean"
+        )
+        assert round(abs(loss_beta_1.item() - expected_loss_beta_1.item()), 5) == 0
+
+        # Case 2: beta = 0 (should be equivalent to KL(teacher || student))
+        loss_beta_0 = IWOPDTrainer.generalized_jsd_loss(student_logits, teacher_logits, beta=0)
+        expected_loss_beta_0 = F.kl_div(
+            F.log_softmax(student_logits, dim=-1), F.softmax(teacher_logits, dim=-1), reduction="batchmean"
+        )
+        assert round(abs(loss_beta_0.item() - expected_loss_beta_0.item()), 5) == 0
+
+    def test_output_shape(self):
+        loss = IWOPDTrainer.generalized_jsd_loss(self.student_logits, self.teacher_logits)
+        assert torch.is_tensor(loss)
+        assert loss.shape == torch.Size([])
+
+    def test_beta_values(self):
+        loss_beta_0 = IWOPDTrainer.generalized_jsd_loss(self.student_logits, self.teacher_logits, beta=0)
+        loss_beta_1 = IWOPDTrainer.generalized_jsd_loss(self.student_logits, self.teacher_logits, beta=1)
+        assert loss_beta_0 != loss_beta_1
+
+    def test_temperature_scaling(self):
+        loss_temp_1 = IWOPDTrainer.generalized_jsd_loss(self.student_logits, self.teacher_logits, temperature=1)
+        loss_temp_2 = IWOPDTrainer.generalized_jsd_loss(self.student_logits, self.teacher_logits, temperature=2)
+        assert loss_temp_1 != loss_temp_2
+
+    def test_reduction_methods(self):
+        loss_batchmean = IWOPDTrainer.generalized_jsd_loss(
+            self.student_logits, self.teacher_logits, reduction="batchmean"
+        )
+        loss_sum = IWOPDTrainer.generalized_jsd_loss(self.student_logits, self.teacher_logits, reduction="sum")
+        loss_mean = IWOPDTrainer.generalized_jsd_loss(self.student_logits, self.teacher_logits, reduction="mean")
+        loss_none = IWOPDTrainer.generalized_jsd_loss(self.student_logits, self.teacher_logits, reduction="none")
+
+        assert loss_batchmean.shape == torch.Size([])
+        assert loss_sum.shape == torch.Size([])
+        assert loss_mean.shape == torch.Size([])
+        assert loss_none.shape == self.student_logits.shape
+
+    def test_symmetry(self):
+        student_teacher = IWOPDTrainer.generalized_jsd_loss(self.student_logits, self.teacher_logits, beta=0.1)
+        teacher_student = IWOPDTrainer.generalized_jsd_loss(self.teacher_logits, self.student_logits, beta=0.1)
+        assert student_teacher != teacher_student
+
+        student_teacher = IWOPDTrainer.generalized_jsd_loss(self.student_logits, self.teacher_logits, beta=0.5)
+        teacher_student = IWOPDTrainer.generalized_jsd_loss(self.teacher_logits, self.student_logits, beta=0.5)
+        assert student_teacher == teacher_student
+
+    def test_zero_loss_for_identical_inputs(self):
+        identical_logits = torch.randn(self.batch_size, self.seq_length, self.vocab_size)
+        loss = IWOPDTrainer.generalized_jsd_loss(identical_logits, identical_logits)
+        assert round(abs(loss.item() - 0), 6) == 0
+
+
+class TestDistillationTrainer(TrlTestCase):
+    def setup_method(self):
+        self.model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+
+    def _make_args(self, **kwargs):
+        args = {
+            "output_dir": self.tmp_dir,
+            "per_device_train_batch_size": 2,
+            "gradient_accumulation_steps": 1,
+            "max_steps": 1,
+            "save_strategy": "no",
+            "report_to": "none",
+            "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},
+            "teacher_model_init_kwargs": {"dtype": "float32", "device_map": None},
+        }
+        args.update(kwargs)
+        return IWOPDConfig(**args)
+
+    def _make_local_trainer(self, **kwargs):
+        dataset = load_dataset("trl-internal-testing/zen", "conversational_language_modeling", split="train")
+        return IWOPDTrainer(
+            model=self.model_id,
+            teacher_model=self.model_id,
+            args=self._make_args(**kwargs),
+            train_dataset=dataset,
+            processing_class=self.tokenizer,
+        )
+
+    def _make_server_trainer(self, **kwargs):
+        dataset = load_dataset("trl-internal-testing/zen", "conversational_language_modeling", split="train")
+        return IWOPDTrainer(
+            model=self.model_id,
+            teacher_model=None,
+            args=self._make_args(use_teacher_server=True, teacher_model_server_url="http://localhost:8000", **kwargs),
+            train_dataset=dataset,
+            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,
+            eval_strategy="steps",
+            max_steps=4,
+            eval_steps=2,
+            save_strategy="steps",
+            save_steps=2,
+            per_device_eval_batch_size=2,
+        )
+        dataset = load_dataset("trl-internal-testing/zen", "conversational_language_modeling")
+        trainer = IWOPDTrainer(
+            model=self.model_id,
+            teacher_model=self.model_id,
+            args=training_args,
+            train_dataset=dataset["train"],
+            eval_dataset=dataset["test"],
+            processing_class=self.tokenizer,
+        )
+
+        train_result = trainer.train()
+
+        assert trainer.state.log_history[-1]["train_loss"] is not None
+        assert trainer.state.log_history[0]["eval_loss"] is not None
+        assert train_result.metrics["train_loss"] >= 0.0
+        assert "model.safetensors" in os.listdir(self.tmp_dir + "/checkpoint-2")
+
+    @pytest.mark.parametrize(
+        "eval_dataset_type",
+        [
+            "dataset",
+            "iterable_dataset",
+            "dataset_dict",
+            "iterable_dataset_dict",
+            "dict_of_dataset",
+            "dict_of_iterable_dataset",
+            "none",
+        ],
+    )
+    def test_init_with_eval_dataset(self, eval_dataset_type):
+        train_dataset = load_dataset("trl-internal-testing/zen", "conversational_language_modeling", split="train")
+
+        if eval_dataset_type == "none":
+            eval_dataset = None
+        else:
+            streaming = "iterable" in eval_dataset_type
+            eval_split = load_dataset(
+                "trl-internal-testing/zen", "conversational_language_modeling", split="test", streaming=streaming
+            )
+            if eval_dataset_type in ("dataset", "iterable_dataset"):
+                eval_dataset = eval_split
+            elif eval_dataset_type in ("dataset_dict", "iterable_dataset_dict"):
+                dataset_dict_cls = IterableDatasetDict if streaming else DatasetDict
+                eval_dataset = dataset_dict_cls({"data1": eval_split, "data2": eval_split})
+            else:  # "dict_of_dataset" or "dict_of_iterable_dataset"
+                eval_dataset = {"data1": eval_split, "data2": eval_split}
+
+        training_args = IWOPDConfig(output_dir=self.tmp_dir, report_to="none")
+        trainer = IWOPDTrainer(
+            model=self.model_id,
+            teacher_model=self.model_id,
+            args=training_args,
+            train_dataset=train_dataset,
+            eval_dataset=eval_dataset,
+            processing_class=self.tokenizer,
+        )
+
+        # The distillation collator consumes raw examples, so eval datasets are stored as-is (not tokenized).
+        if eval_dataset_type == "none":
+            assert trainer.eval_dataset is None
+        elif isinstance(trainer.eval_dataset, dict):
+            assert set(trainer.eval_dataset.keys()) == {"data1", "data2"}
+        else:
+            assert trainer.eval_dataset is eval_dataset
+
+    @pytest.mark.parametrize("loss_top_k", [0, 1])
+    def test_loss_normalizes_by_num_items_in_batch(self, loss_top_k):
+        # When `num_items_in_batch` is passed (as under gradient accumulation), the divergence loss must be reduced as
+        # sum / num_items_in_batch rather than the local per-microbatch mean. See issue #4719. Both the full-vocabulary
+        # JSD path (loss_top_k=0) and the default mixed top-1 path (loss_top_k=1) route through
+        # `_reduce_divergence_loss`, so both must honor `num_items_in_batch`.
+        trainer = self._make_local_trainer(beta=0.5, loss_top_k=loss_top_k)
+
+        # Diverge the teacher from the student so the divergence is well above fp noise (else the loss is ~0).
+        torch.manual_seed(0)
+        with torch.no_grad():
+            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)
+
+        # 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()
+
+        trainer.model.eval()
+        with torch.no_grad():
+            loss_mean = trainer.compute_loss(trainer.model, batch)  # num_items_in_batch=None -> local mean
+            loss_global = trainer.compute_loss(trainer.model, batch, num_items_in_batch=num_valid)
+            loss_double = trainer.compute_loss(trainer.model, batch, num_items_in_batch=num_valid * 2)
+
+        # With num_items_in_batch equal to the local valid-token count, sum/N equals the local mean.
+        torch.testing.assert_close(loss_global, loss_mean, rtol=1e-4, atol=1e-6)
+        # Doubling the global count exactly halves the loss (sum / num_items is linear in 1/num_items).
+        torch.testing.assert_close(loss_double, loss_mean / 2, rtol=1e-4, atol=1e-6)
+
+    @require_liger_kernel
+    @require_torch_accelerator
+    def test_distillation_trainer_with_liger(self):
+        import importlib
+
+        training_args = self._make_args(use_liger_kernel=True, use_cpu=False)
+        dataset = load_dataset("trl-internal-testing/zen", "conversational_language_modeling", split="train")
+
+        trainer = IWOPDTrainer(
+            model=self.model_id,
+            teacher_model=self.model_id,
+            args=training_args,
+            train_dataset=dataset,
+            processing_class=self.tokenizer,
+        )
+
+        try:
+            assert trainer.use_liger_loss is True
+            trainer.train()
+            assert trainer.state.log_history[-1]["train_loss"] is not None
+        finally:
+            importlib.reload(importlib.import_module(trainer.model.__module__))
+
+    def test_sampled_mode_matches_between_local_and_external_teachers(self, monkeypatch):
+        import trl.generation.vllm_client as vllm_client_module
+
+        teacher_client = RecordingTeacherClient()
+        monkeypatch.setattr(vllm_client_module, "VLLMClient", lambda *args, **kwargs: teacher_client)
+
+        local_trainer = self._make_local_trainer(beta=0.5, loss_top_k=1, reverse_kl_top_1_mode="sampled")
+        server_trainer = self._make_server_trainer(beta=0.5, loss_top_k=1)
+
+        cpu_inputs = self._make_batch(local_trainer)
+        expected_sequences, expected_prompt_lengths, _ = build_teacher_request_inputs(
+            cpu_inputs["input_ids"],
+            cpu_inputs["attention_mask"],
+            prompt_attention_mask=cpu_inputs["prompt_attention_mask"],
+            labels=cpu_inputs["labels"],
+        )
+
+        local_inputs = self._move_batch_to_device(cpu_inputs, local_trainer.accelerator.device)
+        server_inputs = self._move_batch_to_device(cpu_inputs, server_trainer.accelerator.device)
+
+        local_trainer.teacher_model.eval()
+        with torch.no_grad():
+            teacher_logits = local_trainer.teacher_model(
+                input_ids=local_inputs["input_ids"],
+                attention_mask=local_inputs["attention_mask"],
+            ).logits
+            teacher_client.result = _build_server_result(
+                teacher_logits,
+                local_inputs,
+                temperature=local_trainer.temperature,
+            )
+            local_loss = local_trainer.compute_loss(local_trainer.model, local_inputs)
+            server_loss = server_trainer.compute_loss(server_trainer.model, server_inputs)
+
+        assert teacher_client.calls[0]["sequences"] == expected_sequences
+        assert teacher_client.calls[0]["prompt_lengths"] == expected_prompt_lengths
+        assert teacher_client.calls[0]["top_logprobs"] == 1
+        torch.testing.assert_close(local_loss, server_loss)
+
+    def test_iw_opd_prefix_weights_and_loss(self):
+        trainer = IWOPDTrainer.__new__(IWOPDTrainer)
+        trainer.temperature = 1.0
+        trainer.iw_opd_gamma = 0.5
+        trainer.iw_opd_epsilon = 1e-8
+        trainer._metrics = {"train": defaultdict(list), "eval": defaultdict(list)}
+        trainer.model = SimpleNamespace(training=True)
+
+        student_logits = torch.zeros(2, 3, 4, requires_grad=True)
+        completion_tokens = torch.tensor([[0, 1, 2], [1, 2, 3]])
+        labels = torch.tensor([[0, 1, 2], [1, -100, -100]])
+        student_logprob = -math.log(4)
+        teacher_actual_logprobs = torch.tensor(
+            [
+                [student_logprob + 0.2, student_logprob - 0.3, student_logprob + 0.5],
+                [student_logprob + 0.4, float("-inf"), float("-inf")],
+            ]
+        )
+
+        loss = trainer._compute_iw_opd_loss(
+            student_logits=student_logits,
+            completion_tokens=completion_tokens,
+            labels=labels,
+            teacher_actual_logprobs=teacher_actual_logprobs,
+        )
+
+        advantages = torch.tensor([[0.2, -0.3, 0.5], [0.4, 0.0, 0.0]])
+        weights = torch.tensor([[1.5, 1.4, 1.25], [1.5, 1.0, 1.0]])
+        expected = -student_logprob * (weights * advantages).sum() / 4
+
+        torch.testing.assert_close(loss, expected)
+        assert trainer._metrics["train"]["iw_opd/max_weight"] == [1.5]
+        assert trainer._metrics["train"]["iw_opd/min_weight"] == [1.25]
+
+        # Under gradient accumulation the sum is divided by the global token count instead of the local one
+        loss_ga = trainer._compute_iw_opd_loss(
+            student_logits=student_logits,
+            completion_tokens=completion_tokens,
+            labels=labels,
+            teacher_actual_logprobs=teacher_actual_logprobs,
+            num_items_in_batch=8,
+        )
+        torch.testing.assert_close(loss_ga, expected * 4 / 8)
+
+        loss.backward()
+        assert student_logits.grad is not None
+        assert torch.isfinite(student_logits.grad).all()
+
+    def test_iw_opd_uses_cached_rollout_logprobs(self):
+        trainer = IWOPDTrainer.__new__(IWOPDTrainer)
+        trainer.temperature = 1.0
+        trainer.iw_opd_gamma = 0.0
+        trainer.iw_opd_epsilon = 1e-8
+        trainer._metrics = {"train": defaultdict(list), "eval": defaultdict(list)}
+        trainer.model = SimpleNamespace(training=True)
+
+        student_logits = torch.zeros(1, 2, 4, requires_grad=True)
+        completion_tokens = torch.tensor([[0, 1]])
+        labels = torch.tensor([[0, 1]])
+        current_student_logprob = -math.log(4)
+        rollout_logprobs = torch.tensor([[-0.5, -0.25]])
+        teacher_actual_logprobs = torch.tensor([[-0.2, -0.05]])
+
+        loss = trainer._compute_iw_opd_loss(
+            student_logits=student_logits,
+            completion_tokens=completion_tokens,
+            labels=labels,
+            teacher_actual_logprobs=teacher_actual_logprobs,
+            rollout_logprobs=rollout_logprobs,
+        )
+
+        expected_advantage_sum = (teacher_actual_logprobs - rollout_logprobs).sum()
+        expected = -current_student_logprob * expected_advantage_sum / 2
+        torch.testing.assert_close(loss, expected)
+
+        # Buffer layout: with mixed prompt lengths, prompt_length (shortest real prompt) is smaller than the
+        # padded prompt width, so rollout logprobs must be stored at full sequence width to survive the
+        # labels-style slice in compute_loss.
+        class _Tok:
+            pad_token_id = 0
+
+            @staticmethod
+            def batch_decode(sequences, skip_special_tokens=False, clean_up_tokenization_spaces=False):
+                return [" ".join(str(int(t)) for t in seq) for seq in sequences]
+
+        trainer = IWOPDTrainer.__new__(IWOPDTrainer)
+        trainer.accelerator = SimpleNamespace(device=torch.device("cpu"))
+        trainer.processing_class = _Tok()
+        trainer.generation_config = SimpleNamespace(max_new_tokens=3)
+        trainer._buffered_inputs = [None]
+        trainer._buffered_text_logs = [None]
+
+        trainer._store_completions_in_buffer(
+            slices=[{}],
+            on_policy_indices=[0],
+            local_slice_indices=[0, 0],
+            local_prompts=[torch.tensor([5]), torch.tensor([6, 7])],
+            completion_ids=[[11, 12, 13], [14]],
+            logprobs=[[[-0.1], [-0.2], [-0.3]], [[-0.4]]],
+            logprob_token_ids=[[[11], [12], [13]], [[14]]],
+        )
+
+        buffered = trainer._buffered_inputs[0]
+        labels = buffered["labels"]
+        rollout_logprobs = buffered["rollout_logprobs"]
+        assert rollout_logprobs.shape == labels.shape
+
+        prompt_length = trainer._compute_prompt_length(buffered)
+        assert prompt_length == 1
+        sliced_labels = labels[:, prompt_length:]
+        sliced_rollout = rollout_logprobs[:, prompt_length:]
+        expected = {(0, 11): -0.1, (0, 12): -0.2, (0, 13): -0.3, (1, 14): -0.4}
+        valid = sliced_labels != -100
+        assert valid.sum() == 4
+        for row, col in valid.nonzero().tolist():
+            token = int(sliced_labels[row, col])
+            torch.testing.assert_close(sliced_rollout[row, col].item(), expected[(row, token)])
+
+    def test_iw_opd_train_runs_with_local_teacher(self):
+        trainer = self._make_local_trainer(
+            distillation_objective="iw_opd",
+            lmbda=1.0,
+            max_completion_length=8,
+        )
+
+        train_result = trainer.train()
+
+        assert train_result.metrics["train_loss"] is not None
+        assert math.isfinite(train_result.metrics["train_loss"])
+        assert any("iw_opd/mean_weight" in record for record in trainer.state.log_history)
+
+    def test_iw_opd_server_path_uses_actual_logprobs(self, monkeypatch):
+        import trl.generation.vllm_client as vllm_client_module
+
+        teacher_client = RecordingTeacherClient()
+        monkeypatch.setattr(vllm_client_module, "VLLMClient", lambda *args, **kwargs: teacher_client)
+
+        trainer = self._make_server_trainer(distillation_objective="iw_opd", lmbda=1.0)
+        inputs = self._move_batch_to_device(self._make_batch(trainer), trainer.accelerator.device)
+        sequences, prompt_lengths, _ = build_teacher_request_inputs(
+            inputs["input_ids"],
+            inputs["attention_mask"],
+            prompt_attention_mask=inputs["prompt_attention_mask"],
+            labels=inputs["labels"],
+        )
+        teacher_client.result = _canned_teacher_logprobs(
+            sequences=sequences,
+            prompt_lengths=prompt_lengths,
+            top_logprobs=1,
+        )
+
+        loss = trainer.compute_loss(trainer.model, inputs)
+
+        assert torch.isfinite(loss)
+        assert teacher_client.calls[0]["top_logprobs"] == 1
+
+
+class TestDistillationTrainerServerPath(TrlTestCase):
+    @classmethod
+    def setup_class(cls):
+        model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
+        cls.device = "cuda" if torch.cuda.is_available() else "cpu"
+        cls.tokenizer = AutoTokenizer.from_pretrained(model_id)
+        cls.tokenizer.pad_token = cls.tokenizer.eos_token
+        cls.model_id = model_id
+
+    def _run_one_step(self, bs, ga, monkeypatch):
+        from trl.generation import vllm_client as vllm_client_module
+
+        fake_client = MagicMock()
+        fake_client.get_sequence_logprobs.side_effect = _canned_teacher_logprobs
+        monkeypatch.setattr(vllm_client_module, "VLLMClient", lambda *a, **kw: fake_client)
+
+        config = IWOPDConfig(
+            output_dir=self.tmp_dir,
+            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,
+            use_teacher_server=True,
+            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",
+            logging_steps=1,
+            use_cpu=not torch.cuda.is_available(),
+            bf16=False,
+        )
+        model = AutoModelForCausalLM.from_pretrained(self.model_id, dtype=torch.float32).to(self.device)
+        trainer = IWOPDTrainer(
+            model=model,
+            args=config,
+            train_dataset=_variable_length_dataset(),
+            processing_class=self.tokenizer,
+        )
+        trainer.teacher_client = fake_client
+        trainer.train()
+        return [rec for rec in trainer.state.log_history if "grad_norm" in rec]
+
+    @pytest.mark.parametrize(("bs", "ga"), [(1, 2), (2, 1)])
+    def test_reverse_kl_finite_grad_with_ragged_batch(self, bs, ga, monkeypatch):
+        records = self._run_one_step(bs=bs, ga=ga, monkeypatch=monkeypatch)
+        assert records, "Expected at least one grad_norm log entry during training"
+        for record in records:
+            assert math.isfinite(record["grad_norm"]), f"grad_norm={record['grad_norm']} leaked -inf into backward"
+            assert math.isfinite(record["loss"])
+
+
+def test_repeat_batch_dataloader_delegates_set_epoch_via_getattr():
+    class DummyDataLoader:
+        def __init__(self):
+            self.epoch = None
+
+        def __iter__(self):
+            yield {"x": 1}
+
+        def __len__(self):
+            return 1
+
+        def set_epoch(self, epoch):
+            self.epoch = epoch
+
+    dataloader = DummyDataLoader()
+    wrapper = _RepeatBatchDataLoader(dataloader, repeat_count=2)
+
+    wrapper.set_epoch(7)
+
+    assert dataloader.epoch == 7
diff --git a/trl/experimental/distillation/distillation_config.py b/trl/experimental/distillation/distillation_config.py
index eccdefaf58d..c70dd4b2d6d 100644
--- a/trl/experimental/distillation/distillation_config.py
+++ b/trl/experimental/distillation/distillation_config.py
@@ -57,13 +57,6 @@ class DistillationConfig(_BaseConfig):
             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
             JSD.
-        distillation_objective (`str`, *optional*, defaults to `"jsd"`):
-            Objective to optimize. `"jsd"` keeps the existing generalized JSD/KL objective. `"iw_opd"` uses the
-            sampled-token Importance-Weighted On-Policy Distillation objective.
-        iw_opd_gamma (`float`, *optional*, defaults to `0.5`):
-            Importance-weight amplification for `distillation_objective="iw_opd"`.
-        iw_opd_epsilon (`float`, *optional*, defaults to `1e-8`):
-            Stabilizer used when normalizing IW-OPD prefix weights.
         reverse_kl_top_1_mode (`str`, *optional*, defaults to `"sampled"`):
             Selection rule for the reverse-KL top-1 token when `beta > 0` and `loss_top_k == 1`. `"sampled"` uses the
             actual completion token in the batch. `"argmax"` uses the student's highest-probability token. This
@@ -209,21 +202,6 @@ class DistillationConfig(_BaseConfig):
             "0.0 = forward KL, 0.5 = JSD, 1.0 = reverse KL."
         },
     )
-    distillation_objective: str = field(
-        default="jsd",
-        metadata={
-            "help": "Objective to optimize. Use 'jsd' for the existing generalized JSD/KL objective or 'iw_opd' "
-            "for sampled-token Importance-Weighted On-Policy Distillation."
-        },
-    )
-    iw_opd_gamma: float = field(
-        default=0.5,
-        metadata={"help": "Importance-weight amplification for distillation_objective='iw_opd'."},
-    )
-    iw_opd_epsilon: float = field(
-        default=1e-8,
-        metadata={"help": "Stabilizer used when normalizing IW-OPD prefix weights."},
-    )
     reverse_kl_top_1_mode: str = field(
         default="sampled",
         metadata={
@@ -405,12 +383,6 @@ def __post_init__(self):
             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}.")
-        if self.distillation_objective not in {"jsd", "iw_opd"}:
-            raise ValueError("distillation_objective must be one of: 'jsd', 'iw_opd'")
-        if self.iw_opd_gamma < 0.0:
-            raise ValueError(f"iw_opd_gamma must be non-negative, got {self.iw_opd_gamma}.")
-        if self.iw_opd_epsilon <= 0.0:
-            raise ValueError(f"iw_opd_epsilon must be positive, got {self.iw_opd_epsilon}.")
         if self.reverse_kl_top_1_mode not in {"sampled", "argmax"}:
             raise ValueError("reverse_kl_top_1_mode must be one of: 'sampled', 'argmax'")
 
@@ -443,25 +415,6 @@ def __post_init__(self):
                 "use_liger_kernel=True is not supported with use_teacher_server=True because the Liger loss path "
                 "requires a local teacher model."
             )
-        if self.distillation_objective == "iw_opd" and self.use_liger_kernel:
-            raise ValueError(
-                "use_liger_kernel=True is not supported with distillation_objective='iw_opd' because IW-OPD needs "
-                "sampled-token student and teacher logprobs."
-            )
-        if self.distillation_objective == "iw_opd" and self.lmbda < 1.0:
-            raise ValueError(
-                "distillation_objective='iw_opd' requires lmbda=1.0 because IW-OPD is an on-policy objective."
-            )
-        if self.distillation_objective == "iw_opd" and self.reverse_kl_top_1_mode != "sampled":
-            raise ValueError(
-                "distillation_objective='iw_opd' requires reverse_kl_top_1_mode='sampled' because IW-OPD is defined "
-                "on sampled completion tokens."
-            )
-        if self.distillation_objective == "iw_opd" and self.use_vllm and self.vllm_sync_frequency != 1:
-            raise ValueError(
-                "distillation_objective='iw_opd' with use_vllm=True requires vllm_sync_frequency=1 so generated "
-                "rollouts match the training policy as closely as possible."
-            )
         if self.use_teacher_server and (
             self.teacher_model_server_url is None or not self.teacher_model_server_url.strip()
         ):
diff --git a/trl/experimental/distillation/distillation_trainer.py b/trl/experimental/distillation/distillation_trainer.py
index 2dc4707a9e9..be6a1632e95 100644
--- a/trl/experimental/distillation/distillation_trainer.py
+++ b/trl/experimental/distillation/distillation_trainer.py
@@ -590,9 +590,6 @@ def __init__(
         # ── Store config values ──
         self.lmbda = args.lmbda
         self.beta = args.beta
-        self.distillation_objective = args.distillation_objective
-        self.iw_opd_gamma = args.iw_opd_gamma
-        self.iw_opd_epsilon = args.iw_opd_epsilon
         self.temperature = args.temperature
         self.top_p = args.top_p
         self.num_generations = args.num_generations
@@ -671,7 +668,7 @@ def __init__(
                 top_p=args.top_p,
                 top_k=args.top_k,
                 max_completion_length=args.max_completion_length,
-                logprobs=0 if args.distillation_objective == "iw_opd" else None,
+                logprobs=None,
             )
             self.vllm_sync_frequency = args.vllm_sync_frequency
             self._last_vllm_sync_step = -1
@@ -888,13 +885,13 @@ def _generate_student_completions(self, slices: list[dict[str, torch.Tensor | An
 
         # Generate completions — pass token IDs directly, no text decoding
         prompt_ids_list = [p.tolist() for p in local_prompts]
-        _, completion_ids, logprobs, logprob_token_ids = self.vllm_generation.generate(
+        _, completion_ids, _, _ = self.vllm_generation.generate(
             prompts=prompt_ids_list, images=None, num_generations=self.num_generations
         )
 
         # Process completions into the buffer
         self._store_completions_in_buffer(
-            slices, on_policy_indices, local_slice_indices, local_prompts, completion_ids, logprobs, logprob_token_ids
+            slices, on_policy_indices, local_slice_indices, local_prompts, completion_ids
         )
 
     def _generate_with_model(self, slices: list[dict[str, torch.Tensor | Any]], on_policy_indices: list[int]):
@@ -961,8 +958,6 @@ def _store_completions_in_buffer(
         local_slice_indices: list[int],
         local_prompts: list[torch.Tensor],
         completion_ids: list,
-        logprobs: list | None = None,
-        logprob_token_ids: list | None = None,
     ):
         """Process vLLM completions and store them in the buffer.
 
@@ -975,18 +970,9 @@ def _store_completions_in_buffer(
 
         # Group completions and prompt token IDs by slice
         slice_completions = {idx: [] for idx in on_policy_indices}
-        slice_rollout_logprobs = {idx: [] for idx in on_policy_indices}
         slice_prompt_ids = {idx: [] for idx in on_policy_indices}
         for i, slice_idx in enumerate(local_slice_indices):
             slice_completions[slice_idx].append(completion_ids[i])
-            if logprobs is not None:
-                sampled_logprobs = []
-                for step_lps, step_token_ids, sampled_token_id in zip(
-                    logprobs[i], logprob_token_ids[i], completion_ids[i], strict=True
-                ):
-                    sampled_idx = step_token_ids.index(sampled_token_id)
-                    sampled_logprobs.append(step_lps[sampled_idx])
-                slice_rollout_logprobs[slice_idx].append(sampled_logprobs)
             slice_prompt_ids[slice_idx].append(local_prompts[i])
 
         for slice_idx in on_policy_indices:
@@ -1006,28 +992,14 @@ def _store_completions_in_buffer(
 
             # Pad/truncate completions (right-pad to max_completion_length)
             completion_tensors = []
-            rollout_logprob_tensors = []
             completion_ids_for_text = []
             completion_lengths = []
-            for completion_idx, comp_ids in enumerate(slice_completions[slice_idx]):
+            for comp_ids in slice_completions[slice_idx]:
                 t = torch.tensor(comp_ids, device=device)
                 if len(t) > max_completion_length:
                     t = t[:max_completion_length]
                 completion_ids_for_text.append(t.tolist())
                 completion_lengths.append(len(t))
-                if logprobs is not None:
-                    rollout_logprobs = torch.tensor(
-                        slice_rollout_logprobs[slice_idx][completion_idx], device=device, dtype=torch.float32
-                    )
-                    if len(rollout_logprobs) > max_completion_length:
-                        rollout_logprobs = rollout_logprobs[:max_completion_length]
-                    if len(rollout_logprobs) < max_completion_length:
-                        rollout_logprobs = F.pad(
-                            rollout_logprobs,
-                            (0, max_completion_length - len(rollout_logprobs)),
-                            value=0.0,
-                        )
-                    rollout_logprob_tensors.append(rollout_logprobs)
                 if len(t) < max_completion_length:
                     padding = torch.full((max_completion_length - len(t),), pad_token_id, device=device, dtype=t.dtype)
                     t = torch.cat([t, padding])
@@ -1055,9 +1027,6 @@ def _store_completions_in_buffer(
             # Update prompts to match the new padding width so prompt_length is consistent
             updated["prompts"] = prompt_ids
             updated["prompt_attention_mask"] = prompt_attention_mask
-            if logprobs is not None:
-                # Full sequence width (zeros over the prompt) so compute_loss can slice it like labels
-                updated["rollout_logprobs"] = F.pad(torch.stack(rollout_logprob_tensors), (prompt_width, 0), value=0.0)
 
             self._buffered_inputs[slice_idx] = updated
             self._buffered_text_logs[slice_idx] = (prompt_texts, completion_texts)
@@ -1220,63 +1189,6 @@ def _get_reverse_kl_top_1_tokens(
             return student_scores.argmax(dim=-1)
         return completion_tokens
 
-    def _compute_iw_opd_loss(
-        self,
-        student_logits: torch.Tensor,
-        completion_tokens: torch.Tensor,
-        labels: torch.Tensor,
-        teacher_actual_logprobs: torch.Tensor,
-        rollout_logprobs: torch.Tensor | None = None,
-        num_items_in_batch=None,
-    ) -> torch.Tensor:
-        """Compute the sampled-token Importance-Weighted On-Policy Distillation loss."""
-        student_log_probs = F.log_softmax(student_logits / self.temperature, dim=-1)
-        student_actual_logprobs = student_log_probs.gather(dim=-1, index=completion_tokens.unsqueeze(-1)).squeeze(-1)
-
-        valid_mask = labels != -100
-        missing_teacher = valid_mask & ~torch.isfinite(teacher_actual_logprobs)
-        if missing_teacher.any():
-            missing_count = int(missing_teacher.sum().item())
-            total_required = int(valid_mask.sum().item())
-            raise ValueError(
-                f"Teacher logprobs are missing for required IW-OPD positions: {missing_count}/{total_required}."
-            )
-
-        safe_teacher_logprobs = torch.where(valid_mask, teacher_actual_logprobs, 0.0)
-        if rollout_logprobs is None:
-            rollout_logprobs = student_actual_logprobs
-        safe_rollout_logprobs = torch.where(valid_mask, rollout_logprobs, 0.0)
-        a_opd = (safe_teacher_logprobs - safe_rollout_logprobs).detach()
-        a_opd = torch.where(valid_mask, a_opd, torch.zeros_like(a_opd))
-
-        abs_advantages = a_opd.abs()
-        prefix_abs_advantages = abs_advantages.cumsum(dim=1) - abs_advantages
-        total_abs_advantages = abs_advantages.sum(dim=1, keepdim=True).clamp_min(self.iw_opd_epsilon)
-        weights = (1.0 + self.iw_opd_gamma * (1.0 - prefix_abs_advantages / total_abs_advantages)).detach()
-        advantages = weights * a_opd
-
-        loss = -student_actual_logprobs * advantages
-        loss = torch.where(valid_mask, loss, torch.zeros_like(loss))
-
-        with torch.no_grad():
-            mode = "train" if self.model.training else "eval"
-            if valid_mask.any():
-                valid_advantages = advantages[valid_mask]
-                valid_abs_advantages = abs_advantages[valid_mask]
-                valid_weights = weights[valid_mask]
-                self._metrics[mode]["iw_opd/mean_advantage"].append(valid_advantages.mean().item())
-                self._metrics[mode]["iw_opd/mean_abs_advantage"].append(valid_abs_advantages.mean().item())
-                self._metrics[mode]["iw_opd/mean_weight"].append(valid_weights.mean().item())
-                self._metrics[mode]["iw_opd/min_weight"].append(valid_weights.min().item())
-                self._metrics[mode]["iw_opd/max_weight"].append(valid_weights.max().item())
-
-        if num_items_in_batch is not None:
-            loss_sum = loss.sum()
-            if isinstance(num_items_in_batch, torch.Tensor):
-                num_items_in_batch = num_items_in_batch.to(loss_sum.device)
-            return loss_sum / num_items_in_batch
-        return loss.sum() / valid_mask.sum().clamp_min(1)
-
     def _compute_sparse_top_1_divergence_loss(
         self,
         student_log_probs: torch.Tensor,
@@ -1583,10 +1495,6 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
         prompt_length = self._compute_prompt_length(inputs)
         labels = inputs["labels"][:, prompt_length:]
         completion_tokens = inputs["input_ids"][:, prompt_length:]
-        # Cached rollout logprobs are stored at full sequence width, so this slice keeps them aligned with labels
-        rollout_logprobs = inputs.get("rollout_logprobs")
-        if rollout_logprobs is not None:
-            rollout_logprobs = rollout_logprobs[:, prompt_length:]
 
         if self.use_teacher_server:
             # Server path: token-level divergence using teacher logprobs.
@@ -1603,26 +1511,7 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
             completion_tokens = completion_tokens[:, :comp_len]
             trimmed_labels = labels[:, :comp_len]
 
-            if self.distillation_objective == "iw_opd":
-                if rollout_logprobs is not None:
-                    rollout_logprobs = rollout_logprobs[:, :comp_len]
-                # Unlike the JSD server losses, IW-OPD raises on missing teacher coverage instead of skipping
-                # positions, so once the window is verified to cover every valid token, num_items_in_batch is
-                # the correct gradient-accumulation denominator here too.
-                if (labels[:, comp_len:] != -100).any():
-                    raise ValueError(
-                        "Teacher server returned fewer completion logprobs than the student completion length; "
-                        "IW-OPD requires teacher logprobs for every completion token."
-                    )
-                loss = self._compute_iw_opd_loss(
-                    student_logits=student_logits[:, :comp_len, :],
-                    completion_tokens=completion_tokens,
-                    labels=trimmed_labels,
-                    teacher_actual_logprobs=teacher_result["actual_logprobs"],
-                    rollout_logprobs=rollout_logprobs,
-                    num_items_in_batch=num_items_in_batch,
-                )
-            elif self.beta > 0:
+            if self.beta > 0:
                 loss = self._compute_server_sparse_top_1_divergence_loss(
                     teacher_result=teacher_result,
                     student_log_probs=student_log_probs[:, :comp_len, :],
@@ -1640,20 +1529,7 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
             teacher_logits = self._get_teacher_logits(inputs)
             student_logits = student_outputs.logits[:, prompt_length - 1 : -1, :]
             teacher_logits = teacher_logits[:, prompt_length - 1 : -1, :]
-            if self.distillation_objective == "iw_opd":
-                teacher_log_probs = F.log_softmax(teacher_logits / self.temperature, dim=-1)
-                teacher_actual_logprobs = teacher_log_probs.gather(
-                    dim=-1, index=completion_tokens.unsqueeze(-1)
-                ).squeeze(-1)
-                loss = self._compute_iw_opd_loss(
-                    student_logits=student_logits,
-                    completion_tokens=completion_tokens,
-                    labels=labels,
-                    teacher_actual_logprobs=teacher_actual_logprobs,
-                    rollout_logprobs=rollout_logprobs,
-                    num_items_in_batch=num_items_in_batch,
-                )
-            elif self.beta > 0 and self.loss_top_k == 1:
+            if self.beta > 0 and self.loss_top_k == 1:
                 loss = self._compute_local_sparse_top_1_divergence_loss(
                     student_logits=student_logits,
                     teacher_logits=teacher_logits,
diff --git a/trl/experimental/iw_opd/__init__.py b/trl/experimental/iw_opd/__init__.py
new file mode 100644
index 00000000000..3ce82f46796
--- /dev/null
+++ b/trl/experimental/iw_opd/__init__.py
@@ -0,0 +1,19 @@
+# Copyright 2020-2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .iw_opd_config import IWOPDConfig
+from .iw_opd_trainer import IWOPDTrainer
+
+
+__all__ = ["IWOPDConfig", "IWOPDTrainer"]
diff --git a/trl/experimental/iw_opd/iw_opd_config.py b/trl/experimental/iw_opd/iw_opd_config.py
new file mode 100644
index 00000000000..2ba33707e4c
--- /dev/null
+++ b/trl/experimental/iw_opd/iw_opd_config.py
@@ -0,0 +1,501 @@
+# Copyright 2020-2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# 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
+
+from ...trainer.base_config import _BaseConfig
+
+
+@dataclass
+class IWOPDConfig(_BaseConfig):
+    # docstyle-ignore
+    r"""
+    Configuration class for the [`IWOPDTrainer`].
+
+    Extends [`~transformers.TrainingArguments`] with parameters specific to knowledge distillation. This config is
+    independent of [`SFTConfig`] — all necessary fields are declared here.
+
+    Using [`~transformers.HfArgumentParser`] we can turn this class into
+    [argparse](https://docs.python.org/3/library/argparse#module-argparse) arguments that can be specified on the
+    command line.
+
+    Parameters:
+        > Parameters that control the model
+
+        model_init_kwargs (`dict[str, Any]`, *optional*):
+            Keyword arguments for `AutoModelForCausalLM.from_pretrained`, used when the `model` argument of the
+            trainer is provided as a string.
+        trust_remote_code (`bool`, *optional*, defaults to `False`):
+            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
+
+        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
+            JSD.
+        distillation_objective (`str`, *optional*, defaults to `"jsd"`):
+            Objective to optimize. `"jsd"` keeps the existing generalized JSD/KL objective. `"iw_opd"` uses the
+            sampled-token Importance-Weighted On-Policy Distillation objective.
+        iw_opd_gamma (`float`, *optional*, defaults to `0.5`):
+            Importance-weight amplification for `distillation_objective="iw_opd"`.
+        iw_opd_epsilon (`float`, *optional*, defaults to `1e-8`):
+            Stabilizer used when normalizing IW-OPD prefix weights.
+        reverse_kl_top_1_mode (`str`, *optional*, defaults to `"sampled"`):
+            Selection rule for the reverse-KL top-1 token when `beta > 0` and `loss_top_k == 1`. `"sampled"` uses the
+            actual completion token in the batch. `"argmax"` uses the student's highest-probability token. This
+            setting does not affect the forward-KL support, which always uses the teacher's top-1 token. Ignored when
+            `beta == 0` or `loss_top_k != 1`.
+        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.
+
+        > Parameters that control the teacher model
+
+        teacher_model_name_or_path (`str` or `None`, *optional*):
+            Model name or path for the teacher model. Used when the teacher is loaded locally.
+        teacher_model_revision (`str` or `None`, *optional*):
+            Model revision of the teacher model (e.g., branch name, tag, or commit hash).
+        teacher_model_init_kwargs (`dict[str, Any]` or `None`, *optional*):
+            Keyword arguments passed to `AutoModelForCausalLM.from_pretrained` when instantiating the teacher model
+            from a string.
+        use_teacher_server (`bool`, *optional*, defaults to `False`):
+            Whether to use an external vLLM teacher server instead of a local teacher model.
+        teacher_model_server_url (`str` or `None`, *optional*):
+            Base URL of a vLLM server hosting the teacher model (e.g., `"http://localhost:8000"`). When set, teacher
+            logprobs are fetched from the server instead of running a local forward pass when `use_teacher_server=True`.
+        loss_top_k (`int`, *optional*, defaults to `1`):
+            Number of top tokens to use when computing the JSD/KL loss. Both student and teacher distributions are
+            restricted to these K tokens and re-normalized before computing divergence. If 0, the full vocabulary
+            is used. For local teachers, the general support rule is teacher top-k for forward KL, student top-k for
+            reverse KL, and the union for mixed JSD. When `beta > 0` and `loss_top_k == 1`, the forward support still
+            uses the teacher's top-1 token, while the reverse top-1 token is controlled by `reverse_kl_top_1_mode`.
+            When `use_teacher_server=True`, the pure forward path (`beta=0`) requires this to be positive and uses the
+            teacher's top-k logprobs for the forward term. When `beta > 0`, server-backed distillation requires
+            `loss_top_k == 1` and only supports `"sampled"` reverse top-1 tokens.
+        loss_add_tail (`bool`, *optional*, defaults to `True`):
+            Whether to append a tail bucket that represents the remaining probability mass outside the selected top-k
+            support when computing the loss.
+
+        > Parameters that control on-policy generation
+
+        num_generations (`int`, *optional*, defaults to `1`):
+            Number of completions to generate per prompt during on-policy generation.
+        generation_batch_size (`int` or `None`, *optional*):
+            Number of unique prompts per worker per optimizer step. If `None`, computed from
+            `(per_device_train_batch_size * gradient_accumulation_steps) // num_generations`.
+        top_p (`float`, *optional*, defaults to `0.95`):
+            Top-p (nucleus) sampling parameter for on-policy generation.
+        top_k (`int`, *optional*, defaults to `0`):
+            Top-k sampling parameter for on-policy generation. `0` disables top-k filtering.
+
+        > Parameters that control vLLM for student generation
+
+        use_vllm (`bool`, *optional*, defaults to `False`):
+            Whether to use vLLM for generating on-policy completions from the student model.
+        vllm_mode (`str`, *optional*, defaults to `"colocate"`):
+            Mode for student vLLM integration. Either `"server"` or `"colocate"`.
+        vllm_server_base_url (`str` or `None`, *optional*):
+            Base URL for the student vLLM server. If provided, `vllm_server_host` and `vllm_server_port` are ignored.
+        vllm_server_host (`str`, *optional*, defaults to `"0.0.0.0"`):
+            Host of the student vLLM server.
+        vllm_server_port (`int`, *optional*, defaults to `8001`):
+            Port of the student vLLM server.
+        vllm_server_timeout (`float`, *optional*, defaults to `240.0`):
+            Timeout for connecting to the student vLLM server.
+        vllm_group_port (`int`, *optional*, defaults to `51216`):
+            Port for the vLLM weight-update group (NCCL communicator).
+        vllm_gpu_memory_utilization (`float`, *optional*, defaults to `0.3`):
+            GPU memory utilization for the colocated student vLLM engine.
+        vllm_tensor_parallel_size (`int`, *optional*, defaults to `1`):
+            Tensor parallel size for the colocated student vLLM engine.
+        vllm_max_model_length (`int` or `None`, *optional*):
+            Maximum model sequence length for the colocated vLLM engine.
+        vllm_model_impl (`str`, *optional*, defaults to `"vllm"`):
+            Model implementation backend for vLLM. Use `"vllm"` or `"transformers"`.
+        vllm_structured_outputs_regex (`str` or `None`, *optional*):
+            Regex pattern for vLLM structured outputs.
+        vllm_sync_frequency (`int`, *optional*, defaults to `1`):
+            Frequency (in training steps) to synchronize student model weights to the vLLM engine.
+        vllm_enable_sleep_mode (`bool`, *optional*, defaults to `False`):
+            Enable vLLM sleep mode to offload student weights during the optimizer step.
+
+        > Parameters that control logging
+
+        log_completions (`bool`, *optional*, defaults to `False`):
+            Whether to log a sample of (prompt, completion) pairs every `log_completions_steps` steps. If `rich` is
+            installed, it prints the sample. If `wandb` and/or `trackio` logging is enabled, it logs it to `wandb`
+            and/or `trackio`.
+        log_completions_steps (`int`, *optional*, defaults to `100`):
+            Number of steps between logging completions. Only used if `log_completions` is `True`.
+        num_completions_to_print (`int` or `None`, *optional*):
+            Number of completions to print. If `None`, all completions are logged.
+    """
+
+    _VALID_DICT_FIELDS = _BaseConfig._VALID_DICT_FIELDS + ["model_init_kwargs", "teacher_model_init_kwargs"]
+
+    # Model
+    model_init_kwargs: dict[str, Any] | str | None = field(
+        default=None,
+        metadata={
+            "help": "Keyword arguments for `AutoModelForCausalLM.from_pretrained`, used when the `model` argument "
+            "of the trainer is provided as a string."
+        },
+    )
+    trust_remote_code: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether to allow loading models and tokenizers that ship custom Python code from the Hub. "
+            "Forwarded to `AutoModelForCausalLM.from_pretrained` and `AutoTokenizer.from_pretrained`, for both the "
+            "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,
+        metadata={"help": "The initial learning rate for AdamW."},
+    )
+
+    # Distillation core
+    temperature: float = field(
+        default=1.0,
+        metadata={
+            "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={
+            "help": "Interpolation coefficient for the Generalized JSD loss. "
+            "0.0 = forward KL, 0.5 = JSD, 1.0 = reverse KL."
+        },
+    )
+    distillation_objective: str = field(
+        default="jsd",
+        metadata={
+            "help": "Objective to optimize. Use 'jsd' for the existing generalized JSD/KL objective or 'iw_opd' "
+            "for sampled-token Importance-Weighted On-Policy Distillation."
+        },
+    )
+    iw_opd_gamma: float = field(
+        default=0.5,
+        metadata={"help": "Importance-weight amplification for distillation_objective='iw_opd'."},
+    )
+    iw_opd_epsilon: float = field(
+        default=1e-8,
+        metadata={"help": "Stabilizer used when normalizing IW-OPD prefix weights."},
+    )
+    reverse_kl_top_1_mode: str = field(
+        default="sampled",
+        metadata={
+            "help": "Reverse-KL top-1 token selection when beta > 0 and loss_top_k == 1. "
+            "Use 'sampled' for the actual completion token or 'argmax' for the student's top-1 token. "
+            "The forward-KL support always uses the teacher's top-1 token. Ignored when beta == 0 or loss_top_k != 1."
+        },
+    )
+    max_completion_length: int = field(
+        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."},
+    )
+
+    # Teacher model (local)
+    teacher_model_name_or_path: str | None = field(
+        default=None,
+        metadata={"help": "Model name or path for the teacher model."},
+    )
+    teacher_model_revision: str | None = field(
+        default=None,
+        metadata={"help": "Model revision of the teacher model (e.g., branch name, tag, or commit hash)."},
+    )
+    teacher_model_init_kwargs: dict[str, Any] | str | None = field(
+        default=None,
+        metadata={
+            "help": "Keyword arguments for `AutoModelForCausalLM.from_pretrained` when instantiating the teacher."
+        },
+    )
+
+    # Teacher model (external vLLM server)
+    use_teacher_server: bool = field(
+        default=False,
+        metadata={"help": "Whether to use an external vLLM teacher server instead of a local teacher model."},
+    )
+    teacher_model_server_url: str | None = field(
+        default=None,
+        metadata={
+            "help": 'Base URL of a vLLM server hosting the teacher model (e.g., "http://localhost:8000"). '
+            "Required when use_teacher_server=True."
+        },
+    )
+    loss_top_k: int = field(
+        default=1,
+        metadata={
+            "help": "Number of top tokens to use when computing the JSD/KL loss. "
+            "Both student and teacher distributions are restricted to these K tokens "
+            "(selected based on beta: teacher's top-k for forward KL, student's top-k for reverse KL, "
+            "union of both for JSD) and re-normalized before computing divergence. "
+            "If 0, the full vocabulary is used (slower but exact). "
+            "When beta > 0 and loss_top_k == 1, the forward support still uses the teacher's top-1 token, "
+            "while the reverse top-1 token is controlled by reverse_kl_top_1_mode. "
+            "When use_teacher_server=True, beta=0 requires loss_top_k > 0 and uses the teacher's top-k "
+            "logprobs for the forward term. When beta > 0, server-backed distillation requires loss_top_k == 1 "
+            "and only supports 'sampled' reverse top-1 tokens."
+        },
+    )
+    loss_add_tail: bool = field(
+        default=True,
+        metadata={
+            "help": "Whether to append a tail bucket representing the remaining probability mass outside the selected top-k support."
+        },
+    )
+
+    # On-policy generation
+    num_generations: int = field(
+        default=1,
+        metadata={"help": "Number of completions to generate per prompt during on-policy generation."},
+    )
+    generation_batch_size: int | None = field(
+        default=None,
+        metadata={
+            "help": "Number of unique prompts per worker per optimizer step. "
+            "If None, computed from (per_device_train_batch_size * gradient_accumulation_steps) // num_generations."
+        },
+    )
+    top_p: float = field(
+        default=0.95,
+        metadata={"help": "Top-p (nucleus) sampling parameter for on-policy generation."},
+    )
+    top_k: int = field(
+        default=0,
+        metadata={"help": "Top-k sampling parameter for on-policy generation. 0 disables top-k filtering."},
+    )
+
+    # vLLM for student generation
+    use_vllm: bool = field(
+        default=False,
+        metadata={"help": "Whether to use vLLM for generating on-policy completions from the student model."},
+    )
+    vllm_mode: str = field(
+        default="colocate",
+        metadata={"help": 'Mode for student vLLM integration. Either "server" or "colocate".'},
+    )
+    vllm_server_base_url: str | None = field(
+        default=None,
+        metadata={"help": "Base URL for the student vLLM server."},
+    )
+    vllm_server_host: str = field(
+        default="0.0.0.0",
+        metadata={"help": "Host of the student vLLM server."},
+    )
+    vllm_server_port: int = field(
+        default=8001,
+        metadata={"help": "Port of the student vLLM server."},
+    )
+    vllm_server_timeout: float = field(
+        default=240.0,
+        metadata={"help": "Timeout for connecting to the student vLLM server."},
+    )
+    vllm_group_port: int = field(
+        default=51216,
+        metadata={"help": "Port for the vLLM weight-update group (NCCL communicator)."},
+    )
+    vllm_gpu_memory_utilization: float = field(
+        default=0.3,
+        metadata={"help": "GPU memory utilization for the colocated student vLLM engine."},
+    )
+    vllm_tensor_parallel_size: int = field(
+        default=1,
+        metadata={"help": "Tensor parallel size for the colocated student vLLM engine."},
+    )
+    vllm_max_model_length: int | None = field(
+        default=None,
+        metadata={"help": "Maximum model sequence length for the colocated vLLM engine."},
+    )
+    vllm_model_impl: str = field(
+        default="vllm",
+        metadata={"help": 'Model implementation backend for vLLM. Use "vllm" or "transformers".'},
+    )
+    vllm_structured_outputs_regex: str | None = field(
+        default=None,
+        metadata={"help": "Regex pattern for vLLM structured outputs."},
+    )
+    vllm_sync_frequency: int = field(
+        default=1,
+        metadata={"help": "Frequency (in training steps) to synchronize student weights to the vLLM engine."},
+    )
+    vllm_enable_sleep_mode: bool = field(
+        default=False,
+        metadata={"help": "Enable vLLM sleep mode to offload student weights during the optimizer step."},
+    )
+
+    # W&B
+
+    # Logging
+    log_completions: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether to log a sample of (prompt, completion) pairs every `log_completions_steps` steps. If `rich` is "
+            "installed, it prints the sample. If `wandb` and/or `trackio` logging is enabled, it logs it to `wandb` "
+            "and/or `trackio`."
+        },
+    )
+    log_completions_steps: int = field(
+        default=100,
+        metadata={"help": "Number of steps between logging completions."},
+    )
+    num_completions_to_print: int | None = field(
+        default=None,
+        metadata={"help": "Number of completions to print. If None, all completions are logged."},
+    )
+
+    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}.")
+        if self.distillation_objective not in {"jsd", "iw_opd"}:
+            raise ValueError("distillation_objective must be one of: 'jsd', 'iw_opd'")
+        if self.iw_opd_gamma < 0.0:
+            raise ValueError(f"iw_opd_gamma must be non-negative, got {self.iw_opd_gamma}.")
+        if self.iw_opd_epsilon <= 0.0:
+            raise ValueError(f"iw_opd_epsilon must be positive, got {self.iw_opd_epsilon}.")
+        if self.reverse_kl_top_1_mode not in {"sampled", "argmax"}:
+            raise ValueError("reverse_kl_top_1_mode must be one of: 'sampled', 'argmax'")
+
+        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}.")
+
+        local_sequence_batch_size = self.per_device_train_batch_size * self.gradient_accumulation_steps
+        if self.generation_batch_size is None:
+            self.generation_batch_size = local_sequence_batch_size // self.num_generations
+        if self.generation_batch_size < 1:
+            raise ValueError(f"generation_batch_size must be at least 1, got {self.generation_batch_size}.")
+        if self.generation_batch_size * self.num_generations != local_sequence_batch_size:
+            raise ValueError(
+                "generation_batch_size * num_generations must equal per_device_train_batch_size * "
+                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.use_teacher_server and self.use_liger_kernel:
+            raise ValueError(
+                "use_liger_kernel=True is not supported with use_teacher_server=True because the Liger loss path "
+                "requires a local teacher model."
+            )
+        if self.distillation_objective == "iw_opd" and self.use_liger_kernel:
+            raise ValueError(
+                "use_liger_kernel=True is not supported with distillation_objective='iw_opd' because IW-OPD needs "
+                "sampled-token student and teacher logprobs."
+            )
+        if self.distillation_objective == "iw_opd" and self.lmbda < 1.0:
+            raise ValueError(
+                "distillation_objective='iw_opd' requires lmbda=1.0 because IW-OPD is an on-policy objective."
+            )
+        if self.distillation_objective == "iw_opd" and self.reverse_kl_top_1_mode != "sampled":
+            raise ValueError(
+                "distillation_objective='iw_opd' requires reverse_kl_top_1_mode='sampled' because IW-OPD is defined "
+                "on sampled completion tokens."
+            )
+        if self.distillation_objective == "iw_opd" and self.use_vllm and self.vllm_sync_frequency != 1:
+            raise ValueError(
+                "distillation_objective='iw_opd' with use_vllm=True requires vllm_sync_frequency=1 so generated "
+                "rollouts match the training policy as closely as possible."
+            )
+        if self.use_teacher_server and (
+            self.teacher_model_server_url is None or not self.teacher_model_server_url.strip()
+        ):
+            raise ValueError("teacher_model_server_url must be set when use_teacher_server=True.")
+
+        if self.use_teacher_server and self.beta == 0 and self.loss_top_k < 1:
+            raise ValueError(
+                f"loss_top_k must be positive when using use_teacher_server=True with beta=0 "
+                f"(got loss_top_k={self.loss_top_k}). The pure forward server path only has access to the "
+                f"teacher's top-k logprobs, so it cannot compute the exact full-vocabulary loss when loss_top_k=0."
+            )
+        if self.use_teacher_server and self.reverse_kl_top_1_mode == "argmax":
+            raise ValueError(
+                "reverse_kl_top_1_mode='argmax' is not supported with use_teacher_server=True because the server "
+                "cannot provide teacher logprobs for arbitrary student-selected tokens."
+            )
+        if self.use_teacher_server and self.beta > 0 and self.loss_top_k != 1:
+            raise ValueError(
+                f"loss_top_k must be 1 when using use_teacher_server=True with beta>0 "
+                f"(got loss_top_k={self.loss_top_k}). Mixed forward/reverse distillation with an external teacher "
+                "is only implemented for top-1 support."
+            )
+        if self.reverse_kl_top_1_mode != "sampled" and (self.beta == 0 or self.loss_top_k != 1):
+            warnings.warn(
+                f"reverse_kl_top_1_mode='{self.reverse_kl_top_1_mode}' has no effect when beta={self.beta} "
+                f"and loss_top_k={self.loss_top_k}. It only applies when beta > 0 and loss_top_k == 1.",
+                UserWarning,
+                stacklevel=2,
+            )
+
+        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,
+            )
diff --git a/trl/experimental/iw_opd/iw_opd_trainer.py b/trl/experimental/iw_opd/iw_opd_trainer.py
new file mode 100644
index 00000000000..1acaa8f32c4
--- /dev/null
+++ b/trl/experimental/iw_opd/iw_opd_trainer.py
@@ -0,0 +1,1911 @@
+# Copyright 2020-2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import random
+import textwrap
+import warnings
+from collections import defaultdict
+from collections.abc import Callable
+from contextlib import nullcontext
+from functools import partial
+from typing import Any, Optional
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+import torch.nn.functional as F
+from accelerate.utils import DistributedType, broadcast_object_list, gather_object
+from datasets import Dataset
+from packaging.version import Version
+from torch.utils.data import DataLoader
+from transformers import AutoTokenizer, TrainerCallback, is_trackio_available, is_wandb_available
+from transformers.data.data_collator import DataCollator
+from transformers.feature_extraction_utils import FeatureExtractionMixin
+from transformers.generation.configuration_utils import GenerationConfig
+from transformers.image_processing_utils import BaseImageProcessor
+from transformers.modeling_utils import PreTrainedModel
+from transformers.processing_utils import ProcessorMixin
+from transformers.tokenization_utils_base import PreTrainedTokenizerBase
+from transformers.trainer_utils import EvalPrediction, seed_worker
+from transformers.utils import is_liger_kernel_available, is_peft_available, is_rich_available
+
+from ...extras.profiling import profiling_decorator
+from ...generation.vllm_generation import VLLMGeneration
+from ...import_utils import is_vllm_available
+from ...models import prepare_deepspeed
+from ...models.utils import _ForwardRedirection, unwrap_model_for_generation
+from ...trainer.base_trainer import _BaseTrainer
+from ...trainer.utils import RepeatSampler, create_model_from_path, disable_dropout_in_model, pad, split_tensor_dict
+from .iw_opd_config import IWOPDConfig
+
+
+if is_liger_kernel_available():
+    from liger_kernel.chunked_loss import LigerFusedLinearJSDLoss
+
+
+if is_peft_available():
+    import peft
+    from peft import PeftConfig, get_peft_model
+
+
+if is_rich_available():
+    from rich.console import Console
+    from rich.panel import Panel
+    from rich.table import Table
+    from rich.text import Text
+
+
+if is_trackio_available():
+    import trackio
+
+
+if is_wandb_available():
+    import wandb
+
+
+def _print_completions_sample(prompts: list[str], completions: list[str], step: int, num_samples: int = None) -> None:
+    """Print a sample of prompt-completion pairs using rich."""
+    if not is_rich_available():
+        return
+
+    console = Console()
+    table = Table(show_header=True, header_style="bold white", expand=True)
+    table.add_column("Prompt", style="bright_yellow")
+    table.add_column("Completion", style="bright_green")
+
+    if num_samples is not None:
+        if num_samples >= len(prompts):
+            num_samples = None
+        elif num_samples <= 0:
+            return
+
+    if num_samples is not None:
+        indices = random.sample(range(len(prompts)), num_samples)
+        prompts = [prompts[i] for i in indices]
+        completions = [completions[i] for i in indices]
+
+    for prompt, completion in zip(prompts, completions, strict=True):
+        table.add_row(Text(prompt), Text(completion))
+        table.add_section()
+
+    panel = Panel(table, expand=False, title=f"Step {step}", border_style="bold white")
+    console.print(panel)
+
+
+def _add_tail_bucket(log_probs, valid_mask):
+    """Append a (K+1)-th tail element: log(1 - sum(exp(top_k_logps))).
+
+    This creates a proper probability distribution over K+1 elements, preventing trivial zero loss when top_k is small
+    (especially top_k=1).
+    """
+    log_sum = torch.logsumexp(log_probs, dim=-1, keepdim=True)
+    log_sum = torch.clamp(log_sum, max=-1e-7)  # ensure sum < 1
+    tail = torch.log(-torch.expm1(log_sum))  # log(1 - exp(log_sum))
+    tail_mask = torch.ones_like(valid_mask[..., :1], dtype=torch.bool)
+    return torch.cat([log_probs, tail], dim=-1), torch.cat([valid_mask, tail_mask], dim=-1)
+
+
+def _jsd_divergence(student_log_probs, teacher_log_probs, beta, support_mask=None):
+    """Compute JSD (or forward/reverse KL) from log-probability tensors.
+
+    When *support_mask* is not None, uses manual computation with masked positions zeroed. When None, uses
+    ``F.kl_div``.
+    """
+    if support_mask is not None:
+        safe_student = torch.where(support_mask, student_log_probs, torch.zeros_like(student_log_probs))
+        safe_teacher = torch.where(support_mask, teacher_log_probs, torch.zeros_like(teacher_log_probs))
+        student_probs = torch.where(support_mask, student_log_probs.exp(), torch.zeros_like(student_log_probs))
+        teacher_probs = torch.where(support_mask, teacher_log_probs.exp(), torch.zeros_like(teacher_log_probs))
+
+        if beta == 0:
+            return torch.nan_to_num(teacher_probs * (safe_teacher - safe_student), nan=0.0)
+        elif beta == 1:
+            return torch.nan_to_num(student_probs * (safe_student - safe_teacher), nan=0.0)
+        else:
+            beta_t = torch.tensor(beta, dtype=student_log_probs.dtype, device=student_log_probs.device)
+            tiny = torch.finfo(student_probs.dtype).tiny
+            mixture_probs = (1 - beta_t) * student_probs + beta_t * teacher_probs
+            safe_mixture = torch.where(
+                support_mask,
+                torch.log(mixture_probs.clamp_min(tiny)),
+                torch.zeros_like(student_log_probs),
+            )
+            kl_teacher = torch.nan_to_num(teacher_probs * (safe_teacher - safe_mixture), nan=0.0)
+            kl_student = torch.nan_to_num(student_probs * (safe_student - safe_mixture), nan=0.0)
+            return beta_t * kl_teacher + (1 - beta_t) * kl_student
+    else:
+        if beta == 0:
+            return F.kl_div(student_log_probs, teacher_log_probs, reduction="none", log_target=True)
+        elif beta == 1:
+            return F.kl_div(teacher_log_probs, student_log_probs, reduction="none", log_target=True)
+        else:
+            beta_t = torch.tensor(beta, dtype=student_log_probs.dtype, device=student_log_probs.device)
+            mixture_log_probs = torch.logsumexp(
+                torch.stack([student_log_probs + torch.log1p(-beta_t), teacher_log_probs + torch.log(beta_t)]),
+                dim=0,
+            )
+            kl_teacher = F.kl_div(mixture_log_probs, teacher_log_probs, reduction="none", log_target=True)
+            kl_student = F.kl_div(mixture_log_probs, student_log_probs, reduction="none", log_target=True)
+            return beta_t * kl_teacher + (1 - beta_t) * kl_student
+
+
+def build_teacher_request_inputs(
+    input_ids: torch.Tensor,
+    attention_mask: torch.Tensor,
+    prompt_attention_mask: torch.Tensor | None = None,
+    labels: torch.Tensor | None = None,
+) -> tuple[list[list[int]], list[int], list[int]]:
+    """Trim padded batch tensors into per-sample sequences for teacher-server requests."""
+
+    if input_ids.shape != attention_mask.shape:
+        raise ValueError(
+            f"input_ids and attention_mask must have the same shape, got {input_ids.shape} and {attention_mask.shape}."
+        )
+
+    input_ids_cpu = input_ids.detach().cpu()
+    attention_mask_cpu = attention_mask.detach().cpu().bool()
+
+    if prompt_attention_mask is not None:
+        prompt_lengths = prompt_attention_mask.detach().cpu().sum(dim=1).to(torch.long)
+    else:
+        if labels is None:
+            raise ValueError("labels are required when prompt_attention_mask is not provided.")
+        if labels.shape != input_ids.shape:
+            raise ValueError(f"labels must match input_ids shape, got {labels.shape} and {input_ids.shape}.")
+        full_lengths = attention_mask_cpu.sum(dim=1).to(torch.long)
+        completion_lengths = (labels.detach().cpu() != -100).sum(dim=1).to(torch.long)
+        prompt_lengths = full_lengths - completion_lengths
+
+    trimmed_input_ids: list[list[int]] = []
+    prompt_lengths_list: list[int] = []
+    completion_lengths_list: list[int] = []
+
+    for row, mask, prompt_length in zip(input_ids_cpu, attention_mask_cpu, prompt_lengths, strict=True):
+        trimmed_row = row[mask]
+        prompt_len = int(prompt_length.item())
+        if prompt_len < 0 or prompt_len > trimmed_row.numel():
+            raise ValueError(
+                f"Invalid prompt length {prompt_len} for trimmed sequence of length {trimmed_row.numel()}."
+            )
+        trimmed_input_ids.append(trimmed_row.tolist())
+        prompt_lengths_list.append(prompt_len)
+        completion_lengths_list.append(int(trimmed_row.numel()) - prompt_len)
+
+    return trimmed_input_ids, prompt_lengths_list, completion_lengths_list
+
+
+class _DistillationCollator:
+    """Data collator for the distillation trainer with independent prompt/completion budgets.
+
+    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) for
+    pure on-policy distillation (``lmbda=1``).
+    """
+
+    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
+
+        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.")
+
+    def __call__(self, examples: list[dict[str, Any]]) -> dict[str, torch.Tensor]:
+        all_input_ids: list[list[int]] = []
+        all_labels: list[list[int]] = []
+        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
+
+            # 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
+            )
+            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)
+
+            all_input_ids.append(input_ids)
+            all_labels.append(labels)
+            all_prompt_ids.append(list(prompt_ids))
+
+        # Convert to tensors and left-pad
+        pad_id = self.tokenizer.pad_token_id
+        input_ids_t = pad(
+            [torch.tensor(ids, dtype=torch.long) for ids in all_input_ids],
+            padding_side="left",
+            padding_value=pad_id,
+        )
+        attention_mask_t = pad(
+            [torch.ones(len(ids), dtype=torch.long) for ids in all_input_ids],
+            padding_side="left",
+            padding_value=0,
+        )
+        labels_t = pad(
+            [torch.tensor(lab, dtype=torch.long) for lab in all_labels],
+            padding_side="left",
+            padding_value=self.ignore_index,
+        )
+        prompts_t = pad(
+            [torch.tensor(ids, dtype=torch.long) for ids in all_prompt_ids],
+            padding_side="left",
+            padding_value=pad_id,
+        )
+        prompt_mask_t = pad(
+            [torch.ones(len(ids), dtype=torch.long) for ids in all_prompt_ids],
+            padding_side="left",
+            padding_value=0,
+        )
+
+        return {
+            "input_ids": input_ids_t,
+            "attention_mask": attention_mask_t,
+            "labels": labels_t,
+            "prompts": prompts_t,
+            "prompt_attention_mask": prompt_mask_t,
+        }
+
+
+class _RepeatBatchDataLoader:
+    """Repeats each collated batch ``repeat_count`` times without re-collation.
+
+    ``RepeatSampler`` with ``repeat_count > 1`` causes the DataLoader to re-collate (re-tokenize) the same examples on
+    every repeat, which is wasteful. This wrapper instead keeps ``repeat_count=1`` in the sampler and repeats the
+    already-collated tensor dict, avoiding redundant tokenization.
+    """
+
+    def __init__(self, dataloader, repeat_count: int):
+        self.dataloader = dataloader
+        self.repeat_count = repeat_count
+
+    def __iter__(self):
+        for batch in self.dataloader:
+            for _ in range(self.repeat_count):
+                yield batch
+
+    def __len__(self):
+        return len(self.dataloader) * self.repeat_count
+
+    def set_epoch(self, epoch: int):
+        if hasattr(self.dataloader, "set_epoch"):
+            self.dataloader.set_epoch(epoch)
+
+    def __getattr__(self, attr):
+        return getattr(self.dataloader, attr)
+
+
+class IWOPDTrainer(_BaseTrainer):
+    """
+    Trainer for Importance-Weighted On-Policy Distillation (IW-OPD).
+
+    IW-OPD ([On the Position Bias of On-Policy Distillation](https://huggingface.co/papers/2606.22600)) reweights the
+    sampled-token distillation signal by accumulated teacher-student prefix discrepancy, downweighting later tokens
+    whose supervision has drifted out of distribution. Select it with `distillation_objective="iw_opd"`.
+
+    .. note::
+        This trainer is a frozen snapshot of the pre-refactor ``DistillationTrainer``, kept as the home for IW-OPD
+        (a teacher-guided policy-gradient method that does not fit the stable, full-vocabulary ``DistillationTrainer``).
+        It is not actively maintained and will not track that trainer's later improvements.
+
+    Supports:
+    - Generalized JSD loss (forward KL, reverse KL, or interpolated JSD via `beta`)
+    - Sampled-token IW-OPD objective (`distillation_objective="iw_opd"`)
+    - On-policy / off-policy mixing via `lmbda` (buffered across gradient accumulation)
+    - Local teacher model or external teacher via vLLM server
+    - Student on-policy generation via vLLM or model.generate()
+    - Liger kernel for memory-efficient fused JSD loss
+    """
+
+    _tag_names = ["trl", "iw-opd"]
+    _name = "IW-OPD"
+    _paper = {
+        "title": "On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes",
+        "id": "2306.13649",
+        # docstyle-ignore
+        "citation": textwrap.dedent("""\
+            @inproceedings{agarwal2024on-policy,
+                title        = {{On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes}},
+                author       = {Rishabh Agarwal and Nino Vieillard and Yongchao Zhou and Piotr Stanczyk and Sabela Ramos Garea and Matthieu Geist and Olivier Bachem},
+                year         = 2024,
+                booktitle    = {The Twelfth International Conference on Learning Representations, {ICLR} 2024, Vienna, Austria, May 7-11, 2024},
+                publisher    = {OpenReview.net},
+                url          = {https://openreview.net/forum?id=3zKtaqxLhW},
+            }"""),
+    }
+
+    def __init__(
+        self,
+        model: PreTrainedModel | nn.Module | str | None = None,
+        teacher_model: PreTrainedModel | nn.Module | str = None,
+        args: IWOPDConfig | None = None,
+        data_collator: DataCollator | None = None,  # type: ignore
+        train_dataset: Dataset | None = None,
+        eval_dataset: Dataset | dict[str, Dataset] | None = None,
+        processing_class: PreTrainedTokenizerBase
+        | BaseImageProcessor
+        | FeatureExtractionMixin
+        | ProcessorMixin
+        | None = None,
+        compute_metrics: Callable[[EvalPrediction], dict] | None = None,
+        callbacks: list[TrainerCallback] | None = None,
+        optimizers: tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR] = (None, None),
+        preprocess_logits_for_metrics: Callable[[torch.Tensor, torch.Tensor], torch.Tensor] | None = None,
+        peft_config: Optional["PeftConfig"] = None,
+    ):
+        if args is None:
+            args = IWOPDConfig(output_dir="tmp_iw_opd")
+
+        # ── Student model loading ──
+        model_init_kwargs = args.model_init_kwargs or {}
+        if isinstance(model_init_kwargs, str):
+            import json
+
+            model_init_kwargs = json.loads(model_init_kwargs)
+        teacher_model_init_kwargs = args.teacher_model_init_kwargs or {}
+        if isinstance(teacher_model_init_kwargs, str):
+            import json
+
+            teacher_model_init_kwargs = json.loads(teacher_model_init_kwargs)
+        if isinstance(model, str):
+            model_name_or_path = model
+            model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code)
+            # Distributed training requires device_map=None ("auto" fails)
+            if args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]:
+                model_init_kwargs["device_map"] = None
+            model = create_model_from_path(model, **model_init_kwargs)
+        else:
+            model_name_or_path = model.config._name_or_path if model is not None else None
+
+        # ── Processing class (tokenizer) ──
+        if processing_class is None and model_name_or_path is not None:
+            processing_class = AutoTokenizer.from_pretrained(
+                model_name_or_path, trust_remote_code=args.trust_remote_code
+            )
+        if processing_class is not None:
+            if getattr(processing_class, "pad_token", None) is None:
+                processing_class.pad_token = processing_class.eos_token
+
+        # ── PEFT ──
+        if peft_config is not None:
+            if not is_peft_available():
+                raise ImportError(
+                    "You passed `peft_config` but the `peft` library is not installed. "
+                    "Install it with `pip install trl[peft]`."
+                )
+            if not isinstance(peft_config, PeftConfig):
+                raise TypeError(
+                    f"`peft_config` must be a `peft.PeftConfig` instance (e.g. `peft.LoraConfig`), "
+                    f"got {type(peft_config).__name__}."
+                )
+            # ZeRO-3 + PEFT for non-quantized models:
+            # - PEFT's default autocast_adapter_dtype=True upcasts LoRA adapter params to fp32 even when the base model is bf16.
+            # - ZeRO-3's _allgather_params_coalesced allocates output buffers using the dtype of the first persistent parameter,
+            #   so mixed-dtype persistent_parameters (bf16 base + fp32 LoRA) cause a TypeError on the first optimizer step.
+            # - Passing autocast_adapter_dtype=False keeps adapter params in the base model dtype (bf16), fixing the mismatch.
+            # - This is safe: the fp32 upcast is a QLoRA-specific concern (low-bit quantized base models), not needed for
+            #   non-quantized bf16 training.
+            # - See:
+            #   - TRL issue: https://github.com/huggingface/trl/issues/6089
+            #   - Upstream issue: https://github.com/deepspeedai/DeepSpeed/issues/8072
+            # - autocast_adapter_dtype was introduced in PEFT 0.12.0; before, no upcast existed: no need to pass the kwarg
+            _is_quantized_model = getattr(model, "is_loaded_in_4bit", False) or getattr(
+                model, "is_loaded_in_8bit", False
+            )
+            get_peft_model_kwargs = {}
+            if (
+                args.deepspeed_plugin is not None
+                and args.deepspeed_plugin.zero_stage == 3
+                and not _is_quantized_model
+                and Version(peft.__version__) >= Version("0.12.0")
+            ):
+                get_peft_model_kwargs["autocast_adapter_dtype"] = False
+            model = get_peft_model(model, peft_config, **get_peft_model_kwargs)
+
+        # ── 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,
+            )
+
+        # ── Liger fused JSD loss ──
+        self.use_liger_loss = False
+        if args.use_liger_kernel:
+            self.liger_loss = LigerFusedLinearJSDLoss(
+                beta=args.beta,
+                ignore_index=-100,
+                temperature=args.temperature,
+                compiled=False,
+                weight_hard_loss=0.0,
+                weight_soft_loss=1.0,
+            )
+            self.use_liger_loss = True
+            self._forward_redirection = _ForwardRedirection()
+
+        # ── Teacher model setup ──
+        self.teacher_client = None
+        self.use_teacher_server = args.use_teacher_server
+        self.teacher_model_server_url = args.teacher_model_server_url
+        self._local_teacher_tokenizer_matches_student = True
+        if self.use_teacher_server:
+            from ...generation.vllm_client import VLLMClient
+
+            self.teacher_client = VLLMClient(base_url=self.teacher_model_server_url, connection_timeout=60.0)
+            teacher_model = None
+        elif teacher_model is not None:
+            if args.teacher_model_init_kwargs is not None and not isinstance(teacher_model, str):
+                raise ValueError(
+                    "You passed teacher_model_init_kwargs to the config, but your teacher_model is already "
+                    "instantiated."
+                )
+
+            teacher_model_name_or_path = (
+                teacher_model
+                if isinstance(teacher_model, str)
+                else getattr(getattr(teacher_model, "config", None), "_name_or_path", None)
+            )
+            if teacher_model_name_or_path is None:
+                raise ValueError(
+                    "IWOPDTrainer requires a local teacher model with `config._name_or_path` set so its "
+                    "tokenizer can be validated against the student tokenizer."
+                )
+
+            teacher_model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code)
+            teacher_tokenizer_kwargs = {}
+            teacher_revision = teacher_model_init_kwargs.get("revision", args.teacher_model_revision)
+            if teacher_revision is not None:
+                teacher_tokenizer_kwargs["revision"] = teacher_revision
+            teacher_tokenizer_kwargs["trust_remote_code"] = teacher_model_init_kwargs["trust_remote_code"]
+            teacher_processing_class = AutoTokenizer.from_pretrained(
+                teacher_model_name_or_path, **teacher_tokenizer_kwargs
+            )
+            if getattr(teacher_processing_class, "pad_token", None) is None:
+                teacher_processing_class.pad_token = teacher_processing_class.eos_token
+            self._local_teacher_tokenizer_matches_student = self._local_teacher_tokenizers_match(
+                processing_class, teacher_processing_class
+            )
+            if not self._local_teacher_tokenizer_matches_student:
+                warnings.warn(
+                    "IWOPDTrainer's built-in local-teacher loss assumes the student and teacher share the "
+                    "same tokenizer. The loaded local teacher tokenizer does not match the student tokenizer, so "
+                    "the teacher weights will be left unchanged for subclass overrides. Direct use of the base "
+                    "IWOPDTrainer with this local teacher remains unsupported.",
+                    UserWarning,
+                    stacklevel=2,
+                )
+
+            if isinstance(teacher_model, str):
+                dtype = teacher_model_init_kwargs.get("dtype")
+                teacher_model_init_kwargs["dtype"] = dtype if dtype in ["auto", None] else getattr(torch, dtype)
+
+            if isinstance(teacher_model, str):
+                init_kwargs = dict(teacher_model_init_kwargs)
+                if args.teacher_model_revision is not None:
+                    init_kwargs.setdefault("revision", args.teacher_model_revision)
+                # Distributed training requires device_map=None ("auto" fails)
+                if args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]:
+                    init_kwargs["device_map"] = None
+                teacher_model = create_model_from_path(teacher_model, **init_kwargs)
+
+        # Trainer does not need to remove unused columns — the collator handles raw data
+        args.remove_unused_columns = False
+
+        super().__init__(
+            model=model,
+            args=args,
+            data_collator=data_collator,
+            train_dataset=train_dataset,
+            eval_dataset=eval_dataset,
+            processing_class=processing_class,
+            compute_metrics=compute_metrics,
+            callbacks=callbacks,
+            optimizers=optimizers,
+            preprocess_logits_for_metrics=preprocess_logits_for_metrics,
+        )
+
+        # ── Prepare teacher model (after super().__init__ so accelerator is ready) ──
+        if teacher_model is not None:
+            if self._local_teacher_tokenizer_matches_student:
+                teacher_model.resize_token_embeddings(self.model.config.get_text_config().vocab_size)
+            if self.is_deepspeed_enabled:
+                self.teacher_model = prepare_deepspeed(teacher_model, self.accelerator)
+            else:
+                self.teacher_model = self.accelerator.prepare_model(teacher_model, evaluation_mode=True)
+        else:
+            self.teacher_model = None
+
+        if args.disable_dropout:
+            disable_dropout_in_model(self.model)
+
+        # ── Store config values ──
+        self.lmbda = args.lmbda
+        self.beta = args.beta
+        self.distillation_objective = args.distillation_objective
+        self.iw_opd_gamma = args.iw_opd_gamma
+        self.iw_opd_epsilon = args.iw_opd_epsilon
+        self.temperature = args.temperature
+        self.top_p = args.top_p
+        self.num_generations = args.num_generations
+        self.reverse_kl_top_1_mode = args.reverse_kl_top_1_mode
+        self.loss_top_k = args.loss_top_k
+        self.loss_add_tail = args.loss_add_tail
+
+        # ── Buffer state ──
+        self._buffered_inputs = None
+        self._buffered_on_policy_flags = None
+        self._buffered_text_logs = None
+        self._buffer_step = 0
+
+        # ── Loss tracking ──
+        self._on_policy_loss_total = 0.0
+        self._off_policy_loss_total = 0.0
+        self._on_policy_step_equiv = 0.0
+        self._off_policy_step_equiv = 0.0
+
+        # ── Generation config ──
+        generation_kwargs = {
+            "max_new_tokens": args.max_completion_length,
+            "temperature": args.temperature,
+            "top_p": args.top_p,
+            "do_sample": True,
+            "top_k": args.top_k,
+            "pad_token_id": self.processing_class.pad_token_id,
+        }
+        self.generation_config = GenerationConfig(**generation_kwargs)
+        self.generation_kwargs = generation_kwargs
+        if (
+            hasattr(self.model.generation_config, "eos_token_id")
+            and self.model.generation_config.eos_token_id is not None
+        ):
+            self.generation_config.eos_token_id = self.model.generation_config.eos_token_id
+
+        # ── Metrics & Logging ──
+        self._metrics = {"train": defaultdict(list), "eval": defaultdict(list)}
+        self._total_train_tokens = 0
+        self.log_completions = args.log_completions
+        self.log_completions_steps = args.log_completions_steps
+        self.num_completions_to_print = args.num_completions_to_print
+
+        self._textual_logs = {
+            "prompt": [],
+            "completion": [],
+        }
+
+        # ── vLLM for student generation ──
+        self.use_vllm = args.use_vllm
+        if self.use_vllm:
+            if not is_vllm_available():
+                raise ImportError(
+                    "vLLM is not available and use_vllm is set to True. Please install vLLM with "
+                    "`pip install vllm` to use it."
+                )
+            self.vllm_generation = VLLMGeneration(
+                model=self.model,
+                accelerator=self.accelerator,
+                processing_class=self.processing_class,
+                mode=args.vllm_mode,
+                structured_outputs_regex=args.vllm_structured_outputs_regex,
+                server_base_url=args.vllm_server_base_url,
+                server_host=args.vllm_server_host,
+                server_port=args.vllm_server_port,
+                group_port=args.vllm_group_port,
+                server_timeout=args.vllm_server_timeout,
+                tensor_parallel_size=args.vllm_tensor_parallel_size,
+                gpu_memory_utilization=args.vllm_gpu_memory_utilization,
+                max_model_length=args.vllm_max_model_length,
+                max_num_seqs=args.per_device_train_batch_size * args.gradient_accumulation_steps,
+                enable_sleep_mode=args.vllm_enable_sleep_mode,
+                model_impl=args.vllm_model_impl,
+                trust_remote_code=args.trust_remote_code,
+                temperature=args.temperature,
+                top_p=args.top_p,
+                top_k=args.top_k,
+                max_completion_length=args.max_completion_length,
+                logprobs=0 if args.distillation_objective == "iw_opd" else None,
+            )
+            self.vllm_sync_frequency = args.vllm_sync_frequency
+            self._last_vllm_sync_step = -1
+
+    @staticmethod
+    def _local_teacher_tokenizers_match(
+        student_processing_class: PreTrainedTokenizerBase,
+        teacher_processing_class: PreTrainedTokenizerBase,
+    ) -> bool:
+        """Check whether the student and local teacher tokenizers share the same vocabulary."""
+        return student_processing_class.get_vocab() == teacher_processing_class.get_vocab()
+
+    def _raise_if_local_teacher_tokenizer_mismatch(self) -> None:
+        """Guard the base local-teacher JSD path, while still allowing subclass overrides."""
+        if self.teacher_model is not None and not self._local_teacher_tokenizer_matches_student:
+            raise ValueError(
+                "IWOPDTrainer's built-in local-teacher loss only supports student/teacher pairs that use "
+                "the same tokenizer. Use a same-tokenizer local teacher, set `use_teacher_server=True`, or "
+                "override the local teacher loss path in a subclass."
+            )
+
+    def _compute_prompt_length(self, inputs: dict[str, torch.Tensor | Any]) -> int:
+        """Compute the earliest prompt boundary that still includes every completion token in the batch."""
+        if inputs.get("labels") is not None:
+            attention_mask = inputs["attention_mask"]
+            labels = inputs["labels"]
+            full_lengths = attention_mask.sum(dim=1)
+            completion_lengths = (labels != -100).sum(dim=1)
+            return int((full_lengths - completion_lengths).min().item())
+        return inputs["prompts"].shape[1]
+
+    def _get_completion_lengths(self, generated_tokens: torch.Tensor, prompt_width: int) -> torch.Tensor:
+        """Infer per-sample completion lengths from generated tokens."""
+        completion_tokens = generated_tokens[:, prompt_width:]
+        pad_token_id = self.processing_class.pad_token_id
+        eos_token_id = self.generation_config.eos_token_id
+        if eos_token_id is None:
+            eos_token_ids = set()
+        elif isinstance(eos_token_id, int):
+            eos_token_ids = {eos_token_id}
+        else:
+            eos_token_ids = set(eos_token_id)
+        pad_equals_eos = pad_token_id is not None and pad_token_id in eos_token_ids
+
+        completion_lengths = []
+        for row in completion_tokens.tolist():
+            if pad_equals_eos and eos_token_ids:
+                completion_length = len(row)
+                for idx, token_id in enumerate(row):
+                    if token_id in eos_token_ids:
+                        completion_length = idx + 1
+                        break
+            elif pad_token_id is not None:
+                completion_length = len(row)
+                while completion_length > 0 and row[completion_length - 1] == pad_token_id:
+                    completion_length -= 1
+            else:
+                completion_length = len(row)
+            completion_lengths.append(completion_length)
+
+        return torch.tensor(completion_lengths, device=generated_tokens.device, dtype=torch.long)
+
+    # ──────────────────────────────────────────────────────────────────────
+    #  Dataset / Dataloader
+    # ──────────────────────────────────────────────────────────────────────
+
+    def _set_signature_columns_if_needed(self):
+        super()._set_signature_columns_if_needed()
+        extra_columns = ["prompts", "prompt_attention_mask", "messages", "chat_template_kwargs", "tools"]
+        if self._signature_columns is None:
+            self._signature_columns = extra_columns
+        else:
+            for col in extra_columns:
+                if col not in self._signature_columns:
+                    self._signature_columns.append(col)
+
+    def _get_train_sampler(self, dataset=None):
+        if dataset is None:
+            dataset = self.train_dataset
+        return RepeatSampler(
+            data_source=dataset,
+            mini_repeat_count=self.num_generations,
+            batch_size=self.args.generation_batch_size * self.accelerator.num_processes,
+            repeat_count=1,
+            shuffle=True,
+            seed=self.args.seed,
+        )
+
+    def get_train_dataloader(self):
+        """
+        Override to load one generation batch per optimizer window.
+
+        The dataloader yields batches of size `per_device_train_batch_size * gradient_accumulation_steps`.
+        RepeatSampler ensures each generation batch is repeated `gradient_accumulation_steps` times so the Trainer's
+        loop iterates the correct number of times.
+        """
+        if self.train_dataset is None:
+            raise ValueError("Trainer: training requires a train_dataset.")
+
+        train_dataset = self.train_dataset
+        data_collator = self.data_collator
+
+        dataloader_params = {
+            "batch_size": self._train_batch_size * self.args.gradient_accumulation_steps,
+            "collate_fn": data_collator,
+            "num_workers": self.args.dataloader_num_workers,
+            "pin_memory": self.args.dataloader_pin_memory,
+            "persistent_workers": self.args.dataloader_persistent_workers,
+        }
+
+        if not isinstance(train_dataset, torch.utils.data.IterableDataset):
+            dataloader_params["sampler"] = self._get_train_sampler()
+            dataloader_params["drop_last"] = self.args.dataloader_drop_last
+            dataloader_params["worker_init_fn"] = partial(
+                seed_worker,
+                num_workers=self.args.dataloader_num_workers,
+                rank=self.args.process_index,
+            )
+            if self.args.dataloader_num_workers > 0:
+                dataloader_params["prefetch_factor"] = self.args.dataloader_prefetch_factor
+
+        base_dataloader = self.accelerator.prepare(DataLoader(train_dataset, **dataloader_params))
+        return _RepeatBatchDataLoader(base_dataloader, repeat_count=self.args.gradient_accumulation_steps)
+
+    # ──────────────────────────────────────────────────────────────────────
+    #  Buffering: on/off-policy mixing across gradient accumulation steps
+    # ──────────────────────────────────────────────────────────────────────
+
+    @profiling_decorator
+    def _prepare_inputs(self, generation_batch: dict[str, torch.Tensor | Any]) -> dict[str, torch.Tensor | Any]:
+        if not self.model.training:
+            return generation_batch
+
+        buffer_steps = self.args.gradient_accumulation_steps
+        if self._buffer_step % buffer_steps == 0 or self._buffered_inputs is None:
+            self._fill_buffer(generation_batch, buffer_steps)
+
+        slice_idx = self._buffer_step % buffer_steps
+        inputs = self._buffered_inputs[slice_idx]
+        self._buffer_step += 1
+        return inputs
+
+    @profiling_decorator
+    def _fill_buffer(self, generation_batch: dict[str, torch.Tensor | Any], buffer_steps: int):
+        """Split batch into slices and decide which are on-policy (student-generated) vs off-policy."""
+        slices = split_tensor_dict(generation_batch, buffer_steps)
+
+        # Decide on-policy flags (synchronized across processes)
+        if self.accelerator.is_main_process:
+            on_policy_flags = [random.random() <= self.lmbda for _ in range(buffer_steps)]
+        else:
+            on_policy_flags = [False] * buffer_steps
+        on_policy_flags = broadcast_object_list(on_policy_flags, from_process=0)
+
+        self._buffered_inputs = [None] * buffer_steps
+        self._buffered_on_policy_flags = on_policy_flags
+        self._buffered_text_logs = [None] * buffer_steps
+
+        # Store off-policy slices directly
+        on_policy_indices = []
+        for i, is_on_policy in enumerate(on_policy_flags):
+            if is_on_policy:
+                on_policy_indices.append(i)
+            else:
+                self._buffered_inputs[i] = slices[i]
+
+        # Generate student completions for on-policy slices
+        if on_policy_indices:
+            self._generate_student_completions(slices, on_policy_indices)
+
+        # Gather on-policy text logs once per optimizer step (all processes must participate)
+        if self.log_completions:
+            on_policy_prompts = []
+            on_policy_completions = []
+            for i in on_policy_indices:
+                if self._buffered_text_logs[i] is not None:
+                    prompts, completions = self._buffered_text_logs[i]
+                    on_policy_prompts.extend(prompts)
+                    on_policy_completions.extend(completions)
+            self._textual_logs["prompt"].extend(gather_object(on_policy_prompts))
+            self._textual_logs["completion"].extend(gather_object(on_policy_completions))
+
+    @profiling_decorator
+    def _generate_student_completions(self, slices: list[dict[str, torch.Tensor | Any]], on_policy_indices: list[int]):
+        """Generate completions from the student model for on-policy training."""
+        if not self.use_vllm:
+            self._generate_with_model(slices, on_policy_indices)
+            return
+
+        # Collect all prompts across on-policy slices, stripping left-padding so vLLM
+        # receives only real prompt token IDs (like GRPO trainer).
+        local_prompts = []
+        local_slice_indices = []
+        pad_token_id = self.processing_class.pad_token_id
+        for slice_idx in on_policy_indices:
+            prompt_mask = slices[slice_idx].get("prompt_attention_mask")
+            for i, prompt in enumerate(slices[slice_idx]["prompts"]):
+                if prompt_mask is not None:
+                    prompt = prompt[prompt_mask[i].bool()]
+                elif pad_token_id is not None:
+                    first_non_pad = (prompt != pad_token_id).nonzero(as_tuple=True)[0]
+                    if len(first_non_pad) > 0:
+                        prompt = prompt[first_non_pad[0] :]
+                local_prompts.append(prompt)
+                local_slice_indices.append(slice_idx)
+
+        # Sync student weights to vLLM if needed
+        if (
+            self.state.global_step != self._last_vllm_sync_step
+            and self.state.global_step % self.vllm_sync_frequency == 0
+        ):
+            self.vllm_generation.sync_weights()
+            self._last_vllm_sync_step = self.state.global_step
+
+        # Generate completions — pass token IDs directly, no text decoding
+        prompt_ids_list = [p.tolist() for p in local_prompts]
+        _, completion_ids, logprobs, logprob_token_ids = self.vllm_generation.generate(
+            prompts=prompt_ids_list, images=None, num_generations=self.num_generations
+        )
+
+        # Process completions into the buffer
+        self._store_completions_in_buffer(
+            slices, on_policy_indices, local_slice_indices, local_prompts, completion_ids, logprobs, logprob_token_ids
+        )
+
+    def _generate_with_model(self, slices: list[dict[str, torch.Tensor | Any]], on_policy_indices: list[int]):
+        """Fallback generation using model.generate() (no vLLM)."""
+        with unwrap_model_for_generation(
+            self.model, self.accelerator, generation_kwargs=self.generation_kwargs
+        ) as unwrapped_model:
+            for slice_idx in on_policy_indices:
+                slice_inputs = slices[slice_idx]
+                generated_outputs = unwrapped_model.generate(
+                    input_ids=slice_inputs["prompts"],
+                    attention_mask=slice_inputs.get("prompt_attention_mask", None),
+                    generation_config=self.generation_config,
+                    return_dict_in_generate=True,
+                )
+                generated_tokens = generated_outputs.sequences
+                batch_size = generated_tokens.size(0)
+                device = generated_tokens.device
+                pad_token_id = self.processing_class.pad_token_id
+                prompt_width = slice_inputs["prompts"].shape[1]
+                prompt_mask = slice_inputs.get("prompt_attention_mask")
+                if prompt_mask is not None:
+                    prompt_token_lengths = prompt_mask.sum(dim=1)
+                else:
+                    prompt_token_lengths = torch.full((batch_size,), prompt_width, dtype=torch.long, device=device)
+                completion_lengths = self._get_completion_lengths(generated_tokens, prompt_width)
+                new_attention_mask, new_labels = self._build_sequence_batch(
+                    generated_tokens, prompt_width, prompt_token_lengths, completion_lengths
+                )
+
+                # Decode for logging
+                prompt_texts = []
+                completion_texts = []
+                for idx in range(batch_size):
+                    prompt_tokens = slice_inputs["prompts"][idx]
+                    if prompt_mask is not None:
+                        prompt_tokens = prompt_tokens[prompt_mask[idx].bool()]
+                    elif pad_token_id is not None:
+                        prompt_tokens = prompt_tokens[prompt_tokens != pad_token_id]
+                    prompt_texts.append(
+                        self.processing_class.decode(prompt_tokens.tolist(), skip_special_tokens=False)
+                    )
+                    length = prompt_width
+                    completion_length = int(completion_lengths[idx].item())
+                    completion_texts.append(
+                        self.processing_class.decode(
+                            generated_tokens[idx, length : length + completion_length].tolist(),
+                            skip_special_tokens=False,
+                        )
+                    )
+
+                updated = dict(slice_inputs)
+                updated["input_ids"] = generated_tokens
+                updated["attention_mask"] = new_attention_mask
+                updated["labels"] = new_labels
+
+                self._buffered_inputs[slice_idx] = updated
+                self._buffered_text_logs[slice_idx] = (prompt_texts, completion_texts)
+
+    def _store_completions_in_buffer(
+        self,
+        slices: list[dict[str, torch.Tensor | Any]],
+        on_policy_indices: list[int],
+        local_slice_indices: list[int],
+        local_prompts: list[torch.Tensor],
+        completion_ids: list,
+        logprobs: list | None = None,
+        logprob_token_ids: list | None = None,
+    ):
+        """Process vLLM completions and store them in the buffer.
+
+        Uses original prompt token IDs directly (no decode/re-encode roundtrip), following the same approach as
+        GRPOTrainer.
+        """
+        device = self.accelerator.device
+        pad_token_id = self.processing_class.pad_token_id if self.processing_class.pad_token_id is not None else 0
+        max_completion_length = self.generation_config.max_new_tokens
+
+        # Group completions and prompt token IDs by slice
+        slice_completions = {idx: [] for idx in on_policy_indices}
+        slice_rollout_logprobs = {idx: [] for idx in on_policy_indices}
+        slice_prompt_ids = {idx: [] for idx in on_policy_indices}
+        for i, slice_idx in enumerate(local_slice_indices):
+            slice_completions[slice_idx].append(completion_ids[i])
+            if logprobs is not None:
+                sampled_logprobs = []
+                for step_lps, step_token_ids, sampled_token_id in zip(
+                    logprobs[i], logprob_token_ids[i], completion_ids[i], strict=True
+                ):
+                    sampled_idx = step_token_ids.index(sampled_token_id)
+                    sampled_logprobs.append(step_lps[sampled_idx])
+                slice_rollout_logprobs[slice_idx].append(sampled_logprobs)
+            slice_prompt_ids[slice_idx].append(local_prompts[i])
+
+        for slice_idx in on_policy_indices:
+            slice_inputs = slices[slice_idx]
+            prompt_id_tensors = slice_prompt_ids[slice_idx]
+            prompt_width = max(len(p) for p in prompt_id_tensors)
+            prompt_token_lengths = torch.tensor([len(p) for p in prompt_id_tensors], device=device, dtype=torch.long)
+            prompt_attention_mask = (
+                torch.arange(prompt_width, device=device).unsqueeze(0)
+                >= (prompt_width - prompt_token_lengths).unsqueeze(1)
+            ).long()
+
+            # Left-pad prompt token IDs to the longest prompt in this slice
+            prompt_ids = torch.stack(
+                [F.pad(p, (prompt_width - len(p), 0), value=pad_token_id) for p in prompt_id_tensors]
+            ).to(device)
+
+            # Pad/truncate completions (right-pad to max_completion_length)
+            completion_tensors = []
+            rollout_logprob_tensors = []
+            completion_ids_for_text = []
+            completion_lengths = []
+            for completion_idx, comp_ids in enumerate(slice_completions[slice_idx]):
+                t = torch.tensor(comp_ids, device=device)
+                if len(t) > max_completion_length:
+                    t = t[:max_completion_length]
+                completion_ids_for_text.append(t.tolist())
+                completion_lengths.append(len(t))
+                if logprobs is not None:
+                    rollout_logprobs = torch.tensor(
+                        slice_rollout_logprobs[slice_idx][completion_idx], device=device, dtype=torch.float32
+                    )
+                    if len(rollout_logprobs) > max_completion_length:
+                        rollout_logprobs = rollout_logprobs[:max_completion_length]
+                    if len(rollout_logprobs) < max_completion_length:
+                        rollout_logprobs = F.pad(
+                            rollout_logprobs,
+                            (0, max_completion_length - len(rollout_logprobs)),
+                            value=0.0,
+                        )
+                    rollout_logprob_tensors.append(rollout_logprobs)
+                if len(t) < max_completion_length:
+                    padding = torch.full((max_completion_length - len(t),), pad_token_id, device=device, dtype=t.dtype)
+                    t = torch.cat([t, padding])
+                completion_tensors.append(t)
+
+            completion_ids_padded = torch.stack(completion_tensors)
+            new_input_ids = torch.cat([prompt_ids, completion_ids_padded], dim=1)
+            completion_lengths = torch.tensor(completion_lengths, device=device, dtype=torch.long)
+            new_attention_mask, new_labels = self._build_sequence_batch(
+                new_input_ids, prompt_width, prompt_token_lengths, completion_lengths
+            )
+
+            # Decode for logging only
+            prompt_texts = self.processing_class.batch_decode(
+                prompt_id_tensors, skip_special_tokens=False, clean_up_tokenization_spaces=False
+            )
+            completion_texts = self.processing_class.batch_decode(
+                completion_ids_for_text, skip_special_tokens=False, clean_up_tokenization_spaces=False
+            )
+
+            updated = dict(slice_inputs)
+            updated["input_ids"] = new_input_ids
+            updated["attention_mask"] = new_attention_mask
+            updated["labels"] = new_labels
+            # Update prompts to match the new padding width so prompt_length is consistent
+            updated["prompts"] = prompt_ids
+            updated["prompt_attention_mask"] = prompt_attention_mask
+            if logprobs is not None:
+                # Full sequence width (zeros over the prompt) so compute_loss can slice it like labels
+                updated["rollout_logprobs"] = F.pad(torch.stack(rollout_logprob_tensors), (prompt_width, 0), value=0.0)
+
+            self._buffered_inputs[slice_idx] = updated
+            self._buffered_text_logs[slice_idx] = (prompt_texts, completion_texts)
+
+    @staticmethod
+    def _build_sequence_batch(
+        new_input_ids: torch.Tensor,
+        prompt_width: int,
+        prompt_token_lengths: torch.Tensor,
+        completion_lengths: torch.Tensor,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        """Build attention mask and labels from prompt/completion lengths."""
+        prompt_token_lengths = prompt_token_lengths.to(device=new_input_ids.device, dtype=torch.long)
+        completion_lengths = completion_lengths.to(device=new_input_ids.device, dtype=torch.long)
+        positions = torch.arange(new_input_ids.shape[1], device=new_input_ids.device).unsqueeze(0)
+        prompt_mask = (positions < prompt_width) & (positions >= (prompt_width - prompt_token_lengths).unsqueeze(1))
+        completion_mask = (positions >= prompt_width) & (positions < (prompt_width + completion_lengths).unsqueeze(1))
+        new_attention_mask = (prompt_mask | completion_mask).long()
+
+        new_labels = torch.full_like(new_input_ids, -100)
+        new_labels[completion_mask] = new_input_ids[completion_mask]
+
+        return new_attention_mask, new_labels
+
+    # ──────────────────────────────────────────────────────────────────────
+    #  Loss computation
+    # ──────────────────────────────────────────────────────────────────────
+
+    @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.
+
+        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
+            jsd = jsd[mask]
+
+        if num_items_in_batch is not None:
+            # Normalize by the global number of valid tokens for gradient-accumulation-correct loss.
+            jsd_sum = jsd.sum()
+            if isinstance(num_items_in_batch, torch.Tensor):
+                num_items_in_batch = num_items_in_batch.to(jsd_sum.device)
+            return jsd_sum / num_items_in_batch
+        if reduction == "batchmean":
+            # 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)
+            return jsd.sum() / denom
+        elif reduction == "sum":
+            return jsd.sum()
+        elif reduction == "mean":
+            return jsd.mean()
+        else:
+            return jsd
+
+    @staticmethod
+    def generalized_jsd_loss(
+        student_logits,
+        teacher_logits,
+        labels=None,
+        beta=0.5,
+        temperature=1.0,
+        reduction="batchmean",
+        top_k=0,
+        add_tail=True,
+        num_items_in_batch=None,
+    ):
+        """
+        Compute the generalized Jensen-Shannon Divergence loss for knowledge distillation.
+
+        Args:
+            student_logits: Tensor of shape (batch_size, sequence_length, vocab_size).
+            teacher_logits: Tensor of shape (batch_size, sequence_length, vocab_size).
+            labels: Tensor of shape (batch_size, sequence_length) with -100 for positions to ignore.
+            beta: Interpolation coefficient. 0.0 = forward KL, 1.0 = reverse KL.
+            temperature: Softmax temperature.
+            reduction: 'batchmean', 'sum', 'mean', or 'none'.
+            top_k: Number of top tokens to restrict the loss to. The support set depends on beta:
+                beta=0 (forward KL) uses teacher's top-k, beta=1 (reverse KL) uses student's top-k, 0<beta<1 (JSD) uses
+                the union of both. Distributions are re-normalized over the selected support. If 0, the full vocabulary
+                is used.
+            add_tail: Whether to append a tail bucket representing the remaining probability mass
+                outside the selected top-k support.
+
+        Returns:
+            Scalar loss tensor.
+        """
+        student_logits = student_logits / temperature
+        teacher_logits = teacher_logits / temperature
+
+        support_mask = None
+
+        if top_k > 0 and student_logits.size(-1) > top_k:
+            neg_inf = torch.full((), float("-inf"), dtype=student_logits.dtype, device=student_logits.device)
+            student_log_probs_full = F.log_softmax(student_logits, dim=-1)
+            teacher_log_probs_full = F.log_softmax(teacher_logits, dim=-1)
+
+            if beta == 0:
+                # Forward KL: teacher-selected support
+                _, support = teacher_logits.topk(top_k, dim=-1)
+                support_mask = torch.ones_like(support, dtype=torch.bool)
+            elif beta == 1:
+                # Reverse KL: student-selected support
+                _, support = student_logits.topk(top_k, dim=-1)
+                support_mask = torch.ones_like(support, dtype=torch.bool)
+            else:
+                # JSD: union of both supports (concatenate + deduplicate)
+                _, student_top = student_logits.topk(top_k, dim=-1)
+                _, teacher_top = teacher_logits.topk(top_k, dim=-1)
+                support = torch.cat([teacher_top, student_top], dim=-1)
+                support_mask = torch.ones(support.shape, dtype=torch.bool, device=support.device)
+                for i in range(1, support.shape[-1]):
+                    prev_matches = support[..., i : i + 1] == support[..., :i]
+                    prev_valid = support_mask[..., :i]
+                    support_mask[..., i] &= ~(prev_matches & prev_valid).any(dim=-1)
+                support = torch.where(support_mask, support, torch.zeros_like(support))
+
+            student_support_logps = student_log_probs_full.gather(-1, support)
+            teacher_support_logps = teacher_log_probs_full.gather(-1, support)
+
+            # Mask invalid (duplicate) positions with -inf
+            student_topk_logps = torch.where(support_mask, student_support_logps, neg_inf)
+            teacher_topk_logps = torch.where(support_mask, teacher_support_logps, neg_inf)
+
+            if add_tail:
+                # Add tail bucket: append log(1 - sum(exp(top_k_logps))) to preserve
+                # the remaining probability mass outside the top-k. This prevents trivial
+                # zero loss when top_k is small (especially top_k=1).
+                base_support_mask = support_mask
+                student_log_probs, support_mask = _add_tail_bucket(student_topk_logps, base_support_mask)
+                teacher_log_probs, _ = _add_tail_bucket(teacher_topk_logps, base_support_mask)
+            else:
+                student_log_probs = student_topk_logps - torch.logsumexp(student_topk_logps, dim=-1, keepdim=True)
+                teacher_log_probs = teacher_topk_logps - torch.logsumexp(teacher_topk_logps, dim=-1, keepdim=True)
+        else:
+            student_log_probs = F.log_softmax(student_logits, dim=-1)
+            teacher_log_probs = F.log_softmax(teacher_logits, dim=-1)
+
+        jsd = _jsd_divergence(student_log_probs, teacher_log_probs, beta, support_mask)
+        return IWOPDTrainer._reduce_divergence_loss(
+            jsd, labels=labels, reduction=reduction, num_items_in_batch=num_items_in_batch
+        )
+
+    def _get_reverse_kl_top_1_tokens(
+        self, student_scores: torch.Tensor, completion_tokens: torch.Tensor
+    ) -> torch.Tensor:
+        """Return the reverse-KL top-1 token IDs for the mixed top-1 loss path.
+
+        Args:
+            student_scores: Any (B, T, V) tensor whose argmax selects the student's top token
+                (logits or log-probs — both are order-preserving).
+            completion_tokens: (B, T) actual token IDs in the completion.
+        """
+        if self.reverse_kl_top_1_mode == "argmax":
+            return student_scores.argmax(dim=-1)
+        return completion_tokens
+
+    def _compute_iw_opd_loss(
+        self,
+        student_logits: torch.Tensor,
+        completion_tokens: torch.Tensor,
+        labels: torch.Tensor,
+        teacher_actual_logprobs: torch.Tensor,
+        rollout_logprobs: torch.Tensor | None = None,
+        num_items_in_batch=None,
+    ) -> torch.Tensor:
+        """Compute the sampled-token Importance-Weighted On-Policy Distillation loss."""
+        student_log_probs = F.log_softmax(student_logits / self.temperature, dim=-1)
+        student_actual_logprobs = student_log_probs.gather(dim=-1, index=completion_tokens.unsqueeze(-1)).squeeze(-1)
+
+        valid_mask = labels != -100
+        missing_teacher = valid_mask & ~torch.isfinite(teacher_actual_logprobs)
+        if missing_teacher.any():
+            missing_count = int(missing_teacher.sum().item())
+            total_required = int(valid_mask.sum().item())
+            raise ValueError(
+                f"Teacher logprobs are missing for required IW-OPD positions: {missing_count}/{total_required}."
+            )
+
+        safe_teacher_logprobs = torch.where(valid_mask, teacher_actual_logprobs, 0.0)
+        if rollout_logprobs is None:
+            rollout_logprobs = student_actual_logprobs
+        safe_rollout_logprobs = torch.where(valid_mask, rollout_logprobs, 0.0)
+        a_opd = (safe_teacher_logprobs - safe_rollout_logprobs).detach()
+        a_opd = torch.where(valid_mask, a_opd, torch.zeros_like(a_opd))
+
+        abs_advantages = a_opd.abs()
+        prefix_abs_advantages = abs_advantages.cumsum(dim=1) - abs_advantages
+        total_abs_advantages = abs_advantages.sum(dim=1, keepdim=True).clamp_min(self.iw_opd_epsilon)
+        weights = (1.0 + self.iw_opd_gamma * (1.0 - prefix_abs_advantages / total_abs_advantages)).detach()
+        advantages = weights * a_opd
+
+        loss = -student_actual_logprobs * advantages
+        loss = torch.where(valid_mask, loss, torch.zeros_like(loss))
+
+        with torch.no_grad():
+            mode = "train" if self.model.training else "eval"
+            if valid_mask.any():
+                valid_advantages = advantages[valid_mask]
+                valid_abs_advantages = abs_advantages[valid_mask]
+                valid_weights = weights[valid_mask]
+                self._metrics[mode]["iw_opd/mean_advantage"].append(valid_advantages.mean().item())
+                self._metrics[mode]["iw_opd/mean_abs_advantage"].append(valid_abs_advantages.mean().item())
+                self._metrics[mode]["iw_opd/mean_weight"].append(valid_weights.mean().item())
+                self._metrics[mode]["iw_opd/min_weight"].append(valid_weights.min().item())
+                self._metrics[mode]["iw_opd/max_weight"].append(valid_weights.max().item())
+
+        if num_items_in_batch is not None:
+            loss_sum = loss.sum()
+            if isinstance(num_items_in_batch, torch.Tensor):
+                num_items_in_batch = num_items_in_batch.to(loss_sum.device)
+            return loss_sum / num_items_in_batch
+        return loss.sum() / valid_mask.sum().clamp_min(1)
+
+    def _compute_sparse_top_1_divergence_loss(
+        self,
+        student_log_probs: torch.Tensor,
+        teacher_top1_token_ids: torch.Tensor,
+        teacher_top1_logprobs: torch.Tensor,
+        reverse_token_ids: torch.Tensor,
+        reverse_teacher_logprobs: torch.Tensor,
+        labels: torch.Tensor,
+        num_items_in_batch=None,
+    ) -> torch.Tensor:
+        """Compute exact generalized JSD/KL on top-1 support for the mixed beta>0 path."""
+        neg_inf = torch.full((), float("-inf"), dtype=student_log_probs.dtype, device=student_log_probs.device)
+
+        if self.beta == 1:
+            support = reverse_token_ids.unsqueeze(-1)
+            support_mask = torch.ones_like(support, dtype=torch.bool)
+            teacher_support_logprobs = reverse_teacher_logprobs.unsqueeze(-1)
+        else:
+            teacher_support = teacher_top1_token_ids.unsqueeze(-1)
+            reverse_support = reverse_token_ids.unsqueeze(-1)
+            support = torch.cat([teacher_support, reverse_support], dim=-1)
+            support_mask = torch.ones_like(support, dtype=torch.bool)
+            support_mask[..., 1] = support[..., 1] != support[..., 0]
+            teacher_support_logprobs = torch.stack([teacher_top1_logprobs, reverse_teacher_logprobs], dim=-1)
+            support = torch.where(support_mask, support, torch.zeros_like(support))
+
+        student_support_logprobs = student_log_probs.gather(-1, support)
+        student_support_logprobs = torch.where(support_mask, student_support_logprobs, neg_inf)
+        teacher_support_logprobs = torch.where(support_mask, teacher_support_logprobs, neg_inf)
+
+        if self.loss_add_tail:
+            base_support_mask = support_mask
+            student_sparse_log_probs, support_mask = _add_tail_bucket(student_support_logprobs, base_support_mask)
+            teacher_sparse_log_probs, _ = _add_tail_bucket(teacher_support_logprobs, base_support_mask)
+        else:
+            student_sparse_log_probs = student_support_logprobs - torch.logsumexp(
+                student_support_logprobs, dim=-1, keepdim=True
+            )
+            teacher_sparse_log_probs = teacher_support_logprobs - torch.logsumexp(
+                teacher_support_logprobs, dim=-1, keepdim=True
+            )
+
+        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
+        )
+
+    def _compute_local_sparse_top_1_divergence_loss(
+        self,
+        student_logits: torch.Tensor,
+        teacher_logits: torch.Tensor,
+        completion_tokens: torch.Tensor,
+        labels: torch.Tensor,
+        num_items_in_batch=None,
+    ) -> torch.Tensor:
+        """Compute the mixed top-1 loss for a local teacher using gathered full-logit probabilities."""
+        student_log_probs = F.log_softmax(student_logits / self.temperature, dim=-1)
+        teacher_log_probs = F.log_softmax(teacher_logits / self.temperature, dim=-1)
+
+        teacher_top1_token_ids = teacher_logits.argmax(dim=-1)
+        teacher_top1_logprobs = teacher_log_probs.gather(dim=-1, index=teacher_top1_token_ids.unsqueeze(-1)).squeeze(
+            -1
+        )
+        reverse_token_ids = self._get_reverse_kl_top_1_tokens(student_logits, completion_tokens)
+        reverse_teacher_logprobs = teacher_log_probs.gather(dim=-1, index=reverse_token_ids.unsqueeze(-1)).squeeze(-1)
+
+        return self._compute_sparse_top_1_divergence_loss(
+            student_log_probs=student_log_probs,
+            teacher_top1_token_ids=teacher_top1_token_ids,
+            teacher_top1_logprobs=teacher_top1_logprobs,
+            reverse_token_ids=reverse_token_ids,
+            reverse_teacher_logprobs=reverse_teacher_logprobs,
+            labels=labels,
+            num_items_in_batch=num_items_in_batch,
+        )
+
+    def _get_teacher_logits(self, inputs: dict[str, torch.Tensor | Any]) -> torch.Tensor:
+        """Get teacher logits — dispatches between local model and external server."""
+        if self.teacher_model is not None:
+            self.teacher_model.eval()
+            with torch.no_grad():
+                return self.teacher_model(
+                    input_ids=inputs["input_ids"],
+                    attention_mask=inputs["attention_mask"],
+                ).logits
+        elif self.use_teacher_server:
+            raise NotImplementedError(
+                "Fetching full teacher logits with use_teacher_server=True is not supported. "
+                "Server-backed distillation only supports per-token logprobs via "
+                "`_get_teacher_token_logprobs_from_server`."
+            )
+        else:
+            raise ValueError("No teacher model or teacher server configured.")
+
+    def _get_teacher_token_logprobs_from_server(
+        self,
+        inputs: dict[str, torch.Tensor | Any],
+        aligned_prompt_length: int,
+    ) -> dict[str, torch.Tensor]:
+        """Fetch per-token teacher logprobs from an external vLLM server.
+
+        Returns a dict with:
+            ``actual_logprobs`` – (batch, completion_length) teacher log-prob for the actual
+                                   token at each position (for reverse KL).
+            ``topk_logprobs`` – (batch, completion_length, K) teacher top-k sorted logprobs
+                                   (for forward KL).
+            ``topk_token_ids`` – (batch, completion_length, K) corresponding token IDs.
+        """
+        import numpy as np
+
+        input_ids = inputs["input_ids"]
+        batch_size = input_ids.shape[0]
+        sequences, prompt_lengths, completion_lengths = build_teacher_request_inputs(
+            input_ids,
+            inputs["attention_mask"],
+            prompt_attention_mask=inputs.get("prompt_attention_mask"),
+            labels=inputs.get("labels"),
+        )
+
+        # The pure forward server path can use the requested teacher top-k support.
+        # When beta > 0, config validation restricts the server-backed path to top-1.
+        requested_top_k = self.loss_top_k
+        result = self.teacher_client.get_sequence_logprobs(
+            sequences=sequences,
+            prompt_lengths=prompt_lengths,
+            top_logprobs=requested_top_k,
+            temperature=self.temperature,
+        )
+        K = requested_top_k
+
+        device = input_ids.device
+        labels = inputs.get("labels")
+        if labels is None:
+            raise ValueError("labels are required to align teacher-server logprobs with the student loss tensors.")
+
+        # The student loss slices tensors in padded-sequence coordinates starting at `aligned_prompt_length`.
+        # Place each teacher completion into that same coordinate system by locating the first non-masked completion
+        # token in `labels`. This works for both left-padded off-policy batches and on-policy batches where
+        # completions are right-padded after a fixed-width prompt block.
+        completion_offsets = []
+        label_mask = labels != -100
+        for sample_mask, comp_len in zip(label_mask, completion_lengths, strict=True):
+            if comp_len == 0:
+                completion_offsets.append(0)
+                continue
+            completion_start = int(torch.nonzero(sample_mask, as_tuple=False)[0].item())
+            completion_offsets.append(completion_start - aligned_prompt_length)
+
+        # Size the output tensors to tightly fit the teacher logprobs. Using the full padded
+        # sequence length would include padding positions with -inf teacher logprobs, producing
+        # +inf in the forward pass and NaN gradients in the backward pass (0 * inf = NaN).
+        # Shorter samples in variable-length batches still need the -inf sentinel at the tail;
+        # downstream loss consumers (_compute_server_sparse_top_1_divergence_loss,
+        # _compute_server_forward_kl_loss) neutralise those positions before the divergence
+        # math runs.
+        completion_length = max(
+            (offset + len(lps) for offset, lps in zip(completion_offsets, result["logprobs"], strict=True)),
+            default=0,
+        )
+
+        # actual_logprobs: (B, T) — teacher logprob for the actual token
+        def _actual_to_tensor(key):
+            arr = np.full((batch_size, completion_length), float("-inf"), dtype=np.float32)
+            for i, (offset, seq_lps) in enumerate(zip(completion_offsets, result[key], strict=True)):
+                if seq_lps:
+                    vals = np.array(seq_lps, dtype=np.float32)  # (comp_len_i, 1)
+                    arr[i, offset : offset + vals.shape[0]] = vals[:, 0]
+            return torch.from_numpy(arr).to(device)
+
+        # topk: (B, T, K)
+        def _topk_to_tensor(key, k, np_dtype, fill):
+            arr = np.full((batch_size, completion_length, k), fill, dtype=np_dtype)
+            for i, (offset, seq_vals) in enumerate(zip(completion_offsets, result[key], strict=True)):
+                if seq_vals:
+                    vals = np.array(seq_vals, dtype=np_dtype)  # (comp_len_i, k)
+                    arr[i, offset : offset + vals.shape[0], :] = vals
+            return torch.from_numpy(arr).to(device)
+
+        return {
+            "actual_logprobs": _actual_to_tensor("actual_logprobs"),
+            "topk_logprobs": _topk_to_tensor("logprobs", K, np.float32, float("-inf")),
+            "topk_token_ids": _topk_to_tensor("logprob_token_ids", K, np.int64, 0),
+        }
+
+    def _compute_server_sparse_top_1_divergence_loss(
+        self,
+        teacher_result: dict[str, torch.Tensor],
+        student_log_probs: torch.Tensor,
+        completion_tokens: torch.Tensor,
+        labels: torch.Tensor,
+    ) -> torch.Tensor:
+        """Compute exact sparse top-1 generalized JSD/KL from server-provided teacher logprobs.
+
+        Args:
+            teacher_result: dict with ``actual_logprobs`` (B, T), ``topk_logprobs`` (B, T, K),
+                ``topk_token_ids`` (B, T, K).
+            student_log_probs: (B, T, V) student log-softmax over vocabulary.
+            completion_tokens: (B, T) actual token IDs in the completion.
+            labels: (B, T) with -100 for positions to ignore.
+        """
+        topk_teacher_lps = teacher_result["topk_logprobs"]  # (B, T, 1)
+        topk_token_ids = teacher_result["topk_token_ids"]  # (B, T, 1)
+        actual_teacher_lps = teacher_result["actual_logprobs"]  # (B, T)
+        required = labels != -100
+
+        missing_actual = required & ~torch.isfinite(actual_teacher_lps)
+        if missing_actual.any():
+            missing_count = int(missing_actual.sum().item())
+            total_required = int(required.sum().item())
+            raise ValueError(
+                "Teacher server is missing actual-token logprobs for required reverse-KL positions: "
+                f"{missing_count}/{total_required}."
+            )
+        if self.beta < 1:
+            teacher_top1_logprobs = topk_teacher_lps.squeeze(-1)
+            missing_top1 = required & ~torch.isfinite(teacher_top1_logprobs)
+            if missing_top1.any():
+                missing_count = int(missing_top1.sum().item())
+                total_required = int(required.sum().item())
+                raise ValueError(
+                    "Teacher server is missing top-1 logprobs for required forward-KL positions: "
+                    f"{missing_count}/{total_required}."
+                )
+
+        # Replace -inf teacher logprobs at intra-batch padding (labels == -100) with 0 so
+        # reverse-KL's student_probs·(log_s - log_t) does not leak +inf into the backward pass.
+        pad_mask_2d = ~required
+        pad_mask_3d = pad_mask_2d.unsqueeze(-1)
+        topk_teacher_lps = torch.where(pad_mask_3d, 0.0, topk_teacher_lps)
+        actual_teacher_lps = torch.where(pad_mask_2d, 0.0, actual_teacher_lps)
+
+        # Server path only supports "sampled" mode — config validation enforces this, but we guard
+        # explicitly so future relaxations of the config check don't silently change behaviour.
+        reverse_token_ids = self._get_reverse_kl_top_1_tokens(student_log_probs, completion_tokens)
+        # The server path normalizes locally (batchmean), not by num_items_in_batch: teacher logprobs may not cover
+        # every student completion token (the loss is summed over the trimmed teacher window), so the global token
+        # count would be the wrong denominator. Gradient-accumulation normalization for the server path is left as a
+        # follow-up.
+        return self._compute_sparse_top_1_divergence_loss(
+            student_log_probs=student_log_probs,
+            teacher_top1_token_ids=topk_token_ids.squeeze(-1),
+            teacher_top1_logprobs=topk_teacher_lps.squeeze(-1),
+            reverse_token_ids=reverse_token_ids,
+            reverse_teacher_logprobs=actual_teacher_lps,
+            labels=labels,
+        )
+
+    def _compute_server_forward_kl_loss(
+        self,
+        teacher_result: dict[str, torch.Tensor],
+        student_log_probs: torch.Tensor,
+        labels: torch.Tensor,
+    ) -> torch.Tensor:
+        """Compute sparse forward KL from server-provided teacher top-k logprobs (beta==0 path).
+
+        Args:
+            teacher_result: dict with ``topk_logprobs`` (B, T, K) and ``topk_token_ids`` (B, T, K).
+            student_log_probs: (B, T, V) student log-softmax over vocabulary.
+            labels: (B, T) with -100 for positions to ignore.
+        """
+        teacher_topk_logprobs = teacher_result["topk_logprobs"]
+        teacher_topk_token_ids = teacher_result["topk_token_ids"]
+        valid = teacher_topk_logprobs > float("-inf")
+        neg_inf = torch.full((), float("-inf"), dtype=student_log_probs.dtype, device=student_log_probs.device)
+        student_topk_logprobs = student_log_probs.gather(dim=-1, index=teacher_topk_token_ids)
+        student_topk_logprobs = torch.where(valid, student_topk_logprobs, neg_inf)
+        teacher_topk_logprobs = torch.where(valid, teacher_topk_logprobs, neg_inf)
+
+        if self.loss_add_tail:
+            base_support_mask = valid
+            student_sparse_log_probs, support_mask = _add_tail_bucket(student_topk_logprobs, base_support_mask)
+            teacher_sparse_log_probs, _ = _add_tail_bucket(teacher_topk_logprobs, base_support_mask)
+        else:
+            support_mask = valid
+            student_sparse_log_probs = student_topk_logprobs - torch.logsumexp(
+                student_topk_logprobs, dim=-1, keepdim=True
+            )
+            teacher_sparse_log_probs = teacher_topk_logprobs - torch.logsumexp(
+                teacher_topk_logprobs, dim=-1, keepdim=True
+            )
+
+        jsd = _jsd_divergence(
+            student_sparse_log_probs,
+            teacher_sparse_log_probs,
+            beta=0.0,
+            support_mask=support_mask,
+        )
+        # 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")
+
+    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
+        self._raise_if_local_teacher_tokenizer_mismatch()
+
+        if self.use_liger_loss:
+            loss = self._compute_liger_loss(model, inputs, num_items_in_batch=num_items_in_batch)
+            return (loss, None) if return_outputs else loss
+
+        # Student forward pass
+        student_outputs = model(
+            input_ids=inputs["input_ids"],
+            attention_mask=inputs["attention_mask"],
+        )
+        prompt_length = self._compute_prompt_length(inputs)
+        labels = inputs["labels"][:, prompt_length:]
+        completion_tokens = inputs["input_ids"][:, prompt_length:]
+        # Cached rollout logprobs are stored at full sequence width, so this slice keeps them aligned with labels
+        rollout_logprobs = inputs.get("rollout_logprobs")
+        if rollout_logprobs is not None:
+            rollout_logprobs = rollout_logprobs[:, prompt_length:]
+
+        if self.use_teacher_server:
+            # Server path: token-level divergence using teacher logprobs.
+            # The server returns:
+            #   actual_logprobs  – (B, T)    teacher log p(x_actual)  (for reverse KL)
+            #   topk_logprobs    – (B, T, K) teacher top-k sorted logprobs (for forward KL)
+            #   topk_token_ids   – (B, T, K) corresponding token IDs
+            teacher_result = self._get_teacher_token_logprobs_from_server(inputs, prompt_length)
+
+            student_logits = student_outputs.logits[:, prompt_length - 1 : -1, :]
+            student_log_probs = F.log_softmax(student_logits / self.temperature, dim=-1)
+
+            comp_len = teacher_result["actual_logprobs"].shape[1]
+            completion_tokens = completion_tokens[:, :comp_len]
+            trimmed_labels = labels[:, :comp_len]
+
+            if self.distillation_objective == "iw_opd":
+                if rollout_logprobs is not None:
+                    rollout_logprobs = rollout_logprobs[:, :comp_len]
+                # Unlike the JSD server losses, IW-OPD raises on missing teacher coverage instead of skipping
+                # positions, so once the window is verified to cover every valid token, num_items_in_batch is
+                # the correct gradient-accumulation denominator here too.
+                if (labels[:, comp_len:] != -100).any():
+                    raise ValueError(
+                        "Teacher server returned fewer completion logprobs than the student completion length; "
+                        "IW-OPD requires teacher logprobs for every completion token."
+                    )
+                loss = self._compute_iw_opd_loss(
+                    student_logits=student_logits[:, :comp_len, :],
+                    completion_tokens=completion_tokens,
+                    labels=trimmed_labels,
+                    teacher_actual_logprobs=teacher_result["actual_logprobs"],
+                    rollout_logprobs=rollout_logprobs,
+                    num_items_in_batch=num_items_in_batch,
+                )
+            elif self.beta > 0:
+                loss = self._compute_server_sparse_top_1_divergence_loss(
+                    teacher_result=teacher_result,
+                    student_log_probs=student_log_probs[:, :comp_len, :],
+                    completion_tokens=completion_tokens,
+                    labels=trimmed_labels,
+                )
+            else:
+                loss = self._compute_server_forward_kl_loss(
+                    teacher_result=teacher_result,
+                    student_log_probs=student_log_probs[:, :comp_len, :],
+                    labels=trimmed_labels,
+                )
+        else:
+            # Local teacher: exact full-vocabulary loss except for the shared mixed top-1 path.
+            teacher_logits = self._get_teacher_logits(inputs)
+            student_logits = student_outputs.logits[:, prompt_length - 1 : -1, :]
+            teacher_logits = teacher_logits[:, prompt_length - 1 : -1, :]
+            if self.distillation_objective == "iw_opd":
+                teacher_log_probs = F.log_softmax(teacher_logits / self.temperature, dim=-1)
+                teacher_actual_logprobs = teacher_log_probs.gather(
+                    dim=-1, index=completion_tokens.unsqueeze(-1)
+                ).squeeze(-1)
+                loss = self._compute_iw_opd_loss(
+                    student_logits=student_logits,
+                    completion_tokens=completion_tokens,
+                    labels=labels,
+                    teacher_actual_logprobs=teacher_actual_logprobs,
+                    rollout_logprobs=rollout_logprobs,
+                    num_items_in_batch=num_items_in_batch,
+                )
+            elif self.beta > 0 and self.loss_top_k == 1:
+                loss = self._compute_local_sparse_top_1_divergence_loss(
+                    student_logits=student_logits,
+                    teacher_logits=teacher_logits,
+                    completion_tokens=completion_tokens,
+                    labels=labels,
+                    num_items_in_batch=num_items_in_batch,
+                )
+            else:
+                loss = self.generalized_jsd_loss(
+                    student_logits=student_logits,
+                    teacher_logits=teacher_logits,
+                    labels=labels,
+                    beta=self.beta,
+                    temperature=self.temperature,
+                    top_k=self.loss_top_k,
+                    add_tail=self.loss_add_tail,
+                    num_items_in_batch=num_items_in_batch,
+                )
+
+        return (loss, student_outputs) if return_outputs else loss
+
+    def _liger_student_forward(self, student, inputs):
+        """Decoder-only forward used by the Liger JSD path (skips lm_head to save memory)."""
+        if hasattr(student, "get_decoder") and student.get_decoder() is not None:
+            decoder = student.get_decoder()
+        else:
+            decoder = getattr(student, getattr(student, "base_model_prefix", "model"), student)
+        return decoder(
+            input_ids=inputs["input_ids"],
+            attention_mask=inputs["attention_mask"],
+            use_cache=False,
+        )
+
+    def _compute_liger_loss(self, model, inputs, num_items_in_batch=None):
+        """Memory-efficient JSD using Liger kernel (operates on hidden states, not full logits)."""
+        # Route through the DDP/FSDP wrapper via _forward_redirection so that
+        # DDP.forward() is called and prepare_for_backward() fires correctly.
+        unwrapped_student = self.accelerator.unwrap_model(model)
+        student_outputs = self._forward_redirection(
+            model, unwrapped_student, self._liger_student_forward, unwrapped_student, inputs
+        )
+
+        self.teacher_model.eval()
+        unwrapped_teacher = self.accelerator.unwrap_model(self.teacher_model)
+        if hasattr(unwrapped_teacher, "get_decoder") and unwrapped_teacher.get_decoder() is not None:
+            base_teacher = unwrapped_teacher.get_decoder()
+        else:
+            base_teacher = getattr(
+                unwrapped_teacher, getattr(unwrapped_teacher, "base_model_prefix", "model"), unwrapped_teacher
+            )
+        with torch.no_grad():
+            teacher_outputs = base_teacher(
+                input_ids=inputs["input_ids"],
+                attention_mask=inputs["attention_mask"],
+                use_cache=False,
+            )
+
+        student_hidden = student_outputs.last_hidden_state[:, :-1]
+        teacher_hidden = teacher_outputs.last_hidden_state[:, :-1]
+        del student_outputs, teacher_outputs
+
+        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))
+        true_labels = masked_input_ids[:, 1:].reshape(-1)
+
+        student_head = unwrapped_student.get_output_embeddings()
+        teacher_head = unwrapped_teacher.get_output_embeddings()
+
+        loss = self.liger_loss(
+            student_input=student_hidden,
+            student_weight=student_head.weight,
+            teacher_input=teacher_hidden,
+            teacher_weight=teacher_head.weight,
+            true_labels=true_labels,
+            student_bias=getattr(student_head, "bias", None),
+            teacher_bias=getattr(teacher_head, "bias", None),
+        )
+
+        # The Liger JSD loss normalizes by the local number of valid tokens. Under gradient accumulation we want
+        # the global normalization, so rescale by `num_valid_local / num_items_in_batch`.
+        if num_items_in_batch is not None:
+            num_valid_local = (true_labels != -100).sum().clamp_min(1)
+            if isinstance(num_items_in_batch, torch.Tensor):
+                num_items_in_batch = num_items_in_batch.to(loss.device)
+            loss = loss * num_valid_local / num_items_in_batch
+
+        del student_hidden, teacher_hidden, true_labels
+        return loss
+
+    def _get_liger_zero3_lm_head_gather_ctx(self, model: nn.Module):
+        """Context manager for gathering lm_head parameters under Liger + ZeRO-3."""
+        if not self.use_liger_loss:
+            return nullcontext()
+
+        deepspeed_plugin = self.accelerator.state.deepspeed_plugin
+        if deepspeed_plugin is None or deepspeed_plugin.zero_stage != 3:
+            return nullcontext()
+
+        import deepspeed
+
+        unwrapped_student = self.accelerator.unwrap_model(model)
+        unwrapped_teacher = self.accelerator.unwrap_model(self.teacher_model)
+        student_head = unwrapped_student.get_output_embeddings()
+        teacher_head = unwrapped_teacher.get_output_embeddings()
+        params = [student_head.weight, teacher_head.weight]
+        if student_head.bias is not None:
+            params.append(student_head.bias)
+        if teacher_head.bias is not None:
+            params.append(teacher_head.bias)
+        return deepspeed.zero.GatheredParameters(params, modifier_rank=None)
+
+    # ──────────────────────────────────────────────────────────────────────
+    #  Training step & Logging
+    # ──────────────────────────────────────────────────────────────────────
+
+    @profiling_decorator
+    def training_step(
+        self, model: nn.Module, inputs: dict[str, torch.Tensor | Any], num_items_in_batch: int | None = None
+    ) -> torch.Tensor:
+        """Training step with on/off-policy loss tracking and completion stats."""
+        buffer_steps = self.args.gradient_accumulation_steps
+
+        with self._get_liger_zero3_lm_head_gather_ctx(model):
+            loss = super().training_step(model, inputs, num_items_in_batch)
+
+        slice_idx = (self._buffer_step - 1) % buffer_steps
+
+        # Determine if this slice is on-policy
+        is_on_policy = False
+        if self._buffered_on_policy_flags is not None and slice_idx < len(self._buffered_on_policy_flags):
+            is_on_policy = self._buffered_on_policy_flags[slice_idx]
+
+        # Track completion length stats — read from buffered inputs (which reflect on-policy generation)
+        actual_inputs = self._buffered_inputs[slice_idx] if self._buffered_inputs is not None else inputs
+        labels = actual_inputs.get("labels")
+        if labels is not None:
+            completion_lengths = (labels != -100).sum(dim=1).float()
+            gathered_lengths = self.accelerator.gather(completion_lengths)
+            mode = "train"
+            prefix = "on_policy" if is_on_policy else "off_policy"
+            self._metrics[mode][f"completions/{prefix}_mean_length"].append(gathered_lengths.mean().item())
+            self._metrics[mode][f"completions/{prefix}_max_length"].append(gathered_lengths.max().item())
+            self._metrics[mode][f"completions/{prefix}_min_length"].append(gathered_lengths.min().item())
+
+            # Log fraction of completions that hit max_completion_length (truncated)
+            max_comp_len = getattr(self.generation_config, "max_new_tokens", None)
+            if is_on_policy and max_comp_len is not None:
+                truncated_frac = (gathered_lengths >= max_comp_len).float().mean().item()
+                self._metrics[mode]["completions/truncated_fraction"].append(truncated_frac)
+
+        # Track loss per policy type
+        loss_scalar = float(loss.detach())
+        step_equiv = 1.0 / self.args.gradient_accumulation_steps
+        if is_on_policy:
+            self._on_policy_loss_total += loss_scalar
+            self._on_policy_step_equiv += step_equiv
+        else:
+            self._off_policy_loss_total += loss_scalar
+            self._off_policy_step_equiv += step_equiv
+
+        return loss
+
+    def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
+        mode = "train" if self.model.training else "eval"
+        metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()}
+
+        if mode == "train":
+            # Aggregate on/off-policy losses across distributed processes
+            device = self.accelerator.device if hasattr(self.accelerator, "device") else torch.device("cpu")
+            vec = torch.tensor(
+                [
+                    self._on_policy_loss_total,
+                    self._off_policy_loss_total,
+                    self._on_policy_step_equiv,
+                    self._off_policy_step_equiv,
+                ],
+                dtype=torch.float64,
+                device=device,
+            )
+
+            if (
+                getattr(self.accelerator, "distributed_type", DistributedType.NO) != DistributedType.NO
+                and dist.is_available()
+                and dist.is_initialized()
+            ):
+                dist.all_reduce(vec, op=dist.ReduceOp.SUM)
+
+            on_sum, off_sum, on_eq, off_eq = vec.tolist()
+            if on_eq > 0:
+                logs["on_policy_loss"] = round(on_sum / on_eq, 4)
+            if off_eq > 0:
+                logs["off_policy_loss"] = round(off_sum / off_eq, 4)
+
+            self._on_policy_loss_total = self._off_policy_loss_total = 0.0
+            self._on_policy_step_equiv = self._off_policy_step_equiv = 0.0
+
+        if mode == "eval":
+            metrics = {f"eval_{key}": val for key, val in metrics.items()}
+
+        logs.update(metrics)
+        super().log(logs, start_time)
+        self._metrics[mode].clear()
+
+        # Log completions to console, wandb, and trackio
+        should_log_completions = (
+            self.log_completions
+            and self.state.global_step > 0
+            and self.state.global_step % self.log_completions_steps == 0
+        )
+
+        if should_log_completions and self.accelerator.is_main_process:
+            prompts = list(self._textual_logs["prompt"])
+            completions = list(self._textual_logs["completion"])
+
+            if prompts:
+                _print_completions_sample(prompts, completions, self.state.global_step, self.num_completions_to_print)
+
+                logging_backends = []
+                if self.args.report_to and "wandb" in self.args.report_to and wandb.run is not None:
+                    logging_backends.append(wandb)
+                if self.args.report_to and "trackio" in self.args.report_to:
+                    logging_backends.append(trackio)
+
+                if logging_backends:
+                    import pandas as pd
+
+                    table_data = {
+                        "step": [str(self.state.global_step)] * len(prompts),
+                        "prompt": prompts,
+                        "completion": completions,
+                    }
+                    df = pd.DataFrame(table_data)
+                    if self.num_completions_to_print and len(df) > self.num_completions_to_print:
+                        df = df.sample(n=self.num_completions_to_print, random_state=42)
+
+                    for logging_backend in logging_backends:
+                        logging_backend.log({"completions": logging_backend.Table(dataframe=df)})
+
+        # Clear text logs on all processes after the logging interval
+        if should_log_completions:
+            self._textual_logs["prompt"].clear()
+            self._textual_logs["completion"].clear()
diff --git a/trl/trainer/base_trainer.py b/trl/trainer/base_trainer.py
index 5b4a9457c18..19a484b9241 100644
--- a/trl/trainer/base_trainer.py
+++ b/trl/trainer/base_trainer.py
@@ -49,6 +49,7 @@
     "GMPOTrainer",
     "GOLDTrainer",
     "GRPOWithReplayBufferTrainer",
+    "IWOPDTrainer",
     "MiniLLMTrainer",
     "NashMDTrainer",
     "OnlineDPOTrainer",

From b65f8a75851e1273080972f451883dd4325044a2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Sun, 19 Jul 2026 20:15:10 +0000
Subject: [PATCH 05/44] Extract the teacher-server path into
 ServerDistillationTrainer
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The server-backed path (fetch per-token teacher logprobs from a vLLM server
instead of a local teacher forward) is a distinct use case that only ever runs
over a sparse top-k support. Move it out of the base trainer into a dedicated
experimental subclass so the base can converge on the local, full-vocabulary
objective.

New package `trl/experimental/server_distillation/`:
- `ServerDistillationTrainer(DistillationTrainer)` overriding `compute_loss`
  with the server path, plus the moved `_get_teacher_token_logprobs_from_server`,
  `_compute_server_sparse_top_1_divergence_loss`, `_compute_server_forward_kl_loss`
  and the module-level `build_teacher_request_inputs`. The shared sparse helpers
  (`_compute_sparse_top_1_divergence_loss`, `_get_reverse_kl_top_1_tokens`,
  `_add_tail_bucket`, `_jsd_divergence`, `_reduce_divergence_loss`) are inherited
  from the base, not copied — they are still used by the base local top-1 path.
- `ServerDistillationConfig(DistillationConfig)` adding `teacher_model_server_url`
  and carrying the server-only validations, now unconditional.

Base trainer:
- `__init__` drops the server branch; a `None` teacher just leaves
  `self.teacher_model = None`.
- `compute_loss` drops the server dispatch; the local path runs unconditionally.
- `_get_teacher_logits` drops the server NotImplementedError branch.
- config drops `use_teacher_server` / `teacher_model_server_url` and their checks.

Server tests move to `tests/experimental/test_server_distillation_trainer.py`.

The server-side infrastructure (`/get_sequence_logprobs/`, VLLMClient) is
untouched and still used by SDFT/SDPO.

`compute_loss` in the base goes from 5 dispatch paths to 2 (local top-1, local
full-vocab) plus the Liger path.
---
 .../experimental/test_distillation_trainer.py | 395 +--------------
 .../test_server_distillation_trainer.py       | 456 ++++++++++++++++++
 .../distillation/distillation_config.py       |  53 +-
 .../distillation/distillation_trainer.py      | 352 ++------------
 .../server_distillation/__init__.py           |  19 +
 .../server_distillation_config.py             |  66 +++
 .../server_distillation_trainer.py            | 353 ++++++++++++++
 trl/trainer/base_trainer.py                   |   1 +
 8 files changed, 932 insertions(+), 763 deletions(-)
 create mode 100644 tests/experimental/test_server_distillation_trainer.py
 create mode 100644 trl/experimental/server_distillation/__init__.py
 create mode 100644 trl/experimental/server_distillation/server_distillation_config.py
 create mode 100644 trl/experimental/server_distillation/server_distillation_trainer.py

diff --git a/tests/experimental/test_distillation_trainer.py b/tests/experimental/test_distillation_trainer.py
index d9f5b1bf89a..f9c740b4279 100644
--- a/tests/experimental/test_distillation_trainer.py
+++ b/tests/experimental/test_distillation_trainer.py
@@ -12,23 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import math
 import os
-from unittest.mock import MagicMock
 
 import pytest
 import torch
 import torch.nn.functional as F
-from datasets import Dataset, DatasetDict, IterableDatasetDict, load_dataset
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from datasets import DatasetDict, IterableDatasetDict, load_dataset
+from transformers import AutoTokenizer
 
 from trl.experimental.distillation import DistillationConfig, DistillationTrainer
-from trl.experimental.distillation.distillation_trainer import (
-    _add_tail_bucket,
-    _jsd_divergence,
-    _RepeatBatchDataLoader,
-    build_teacher_request_inputs,
-)
+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
@@ -38,290 +31,11 @@ def _make_distillation_config_kwargs(tmp_path):
     return {"output_dir": str(tmp_path), "report_to": "none", "use_cpu": True, "bf16": False}
 
 
-def _build_server_result(teacher_logits, inputs, temperature=1.0):
-    """Simulate a vLLM server response with variable-length per-sample completions."""
-    _, _, completion_lengths = build_teacher_request_inputs(
-        inputs["input_ids"],
-        inputs["attention_mask"],
-        prompt_attention_mask=inputs.get("prompt_attention_mask"),
-        labels=inputs.get("labels"),
-    )
-
-    label_mask = inputs["labels"] != -100
-    actual_logprobs = []
-    logprobs = []
-    logprob_token_ids = []
-
-    for i, comp_len in enumerate(completion_lengths):
-        if comp_len == 0:
-            actual_logprobs.append([])
-            logprobs.append([])
-            logprob_token_ids.append([])
-            continue
-
-        comp_start = int(torch.nonzero(label_mask[i], as_tuple=False)[0].item())
-        sample_logits = teacher_logits[i, comp_start - 1 : comp_start - 1 + comp_len, :]
-        sample_log_probs = F.log_softmax(sample_logits / temperature, dim=-1)
-        comp_tokens = inputs["input_ids"][i, comp_start : comp_start + comp_len]
-
-        top1_ids = sample_logits.argmax(dim=-1, keepdim=True)
-        top1_lps = sample_log_probs.gather(dim=-1, index=top1_ids)
-        actual_lps = sample_log_probs.gather(dim=-1, index=comp_tokens.unsqueeze(-1))
-
-        actual_logprobs.append(actual_lps.tolist())
-        logprobs.append(top1_lps.tolist())
-        logprob_token_ids.append(top1_ids.tolist())
-
-    return {
-        "actual_logprobs": actual_logprobs,
-        "logprobs": logprobs,
-        "logprob_token_ids": logprob_token_ids,
-    }
-
-
-class RecordingTeacherClient:
-    def __init__(self):
-        self.calls = []
-        self.result = None
-
-    def get_sequence_logprobs(self, sequences, prompt_lengths, top_logprobs, temperature):
-        self.calls.append(
-            {
-                "sequences": sequences,
-                "prompt_lengths": prompt_lengths,
-                "top_logprobs": top_logprobs,
-                "temperature": temperature,
-            }
-        )
-        return self.result
-
-
-def _ragged_server_response():
-    # Two samples with completion lengths 1 and 3 respectively; matches the wire format
-    # of VLLMClient.get_sequence_logprobs (per-sample shape (comp_len, top_k=1)).
-    return {
-        "logprobs": [[[-2.3]], [[-1.1], [-0.4], [-3.0]]],
-        "logprob_token_ids": [[[90]], [[90], [9217], [100]]],
-        "actual_logprobs": [[[-2.3]], [[-1.1], [-0.4], [-3.0]]],
-    }
-
-
-def _canned_teacher_logprobs(**kwargs):
-    # Fabricate ragged per-sample logprobs matching the requested sequence shapes.
-    sequences = kwargs["sequences"]
-    prompt_lengths = kwargs["prompt_lengths"]
-    top_k = kwargs.get("top_logprobs", 1)
-    logprobs, token_ids, actual = [], [], []
-    for seq, plen in zip(sequences, prompt_lengths, strict=True):
-        comp_len = len(seq) - plen
-        logprobs.append([[-1.0 - 0.05 * i] * top_k for i in range(comp_len)])
-        token_ids.append([[int(seq[plen + i])] * top_k for i in range(comp_len)])
-        actual.append([[-1.0 - 0.05 * i] for i in range(comp_len)])
-    return {"logprobs": logprobs, "logprob_token_ids": token_ids, "actual_logprobs": actual}
-
-
-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.",
-                    },
-                ]
-            },
-        ]
-    )
-
-
-def test_distillation_config_rejects_liger_with_teacher_server(tmp_path):
-    with pytest.raises(ValueError, match="use_liger_kernel=True is not supported with use_teacher_server=True"):
-        DistillationConfig(
-            **_make_distillation_config_kwargs(tmp_path),
-            use_teacher_server=True,
-            teacher_model_server_url="http://localhost:8000",
-            use_liger_kernel=True,
-        )
-
-
 def test_distillation_config_rejects_invalid_reverse_kl_top_1_mode(tmp_path):
     with pytest.raises(ValueError, match="reverse_kl_top_1_mode must be one of"):
         DistillationConfig(**_make_distillation_config_kwargs(tmp_path), reverse_kl_top_1_mode="invalid")
 
 
-def test_distillation_config_rejects_teacher_server_with_reverse_kl_argmax(tmp_path):
-    with pytest.raises(ValueError, match="reverse_kl_top_1_mode='argmax' is not supported"):
-        DistillationConfig(
-            **_make_distillation_config_kwargs(tmp_path),
-            use_teacher_server=True,
-            teacher_model_server_url="http://localhost:8000",
-            reverse_kl_top_1_mode="argmax",
-        )
-
-
-def test_distillation_config_rejects_teacher_server_mixed_loss_without_top_1(tmp_path):
-    with pytest.raises(ValueError, match="loss_top_k must be 1 when using use_teacher_server=True with beta>0"):
-        DistillationConfig(
-            **_make_distillation_config_kwargs(tmp_path),
-            use_teacher_server=True,
-            teacher_model_server_url="http://localhost:8000",
-            beta=0.5,
-            loss_top_k=2,
-        )
-
-
-def test_distillation_config_requires_teacher_server_url(tmp_path):
-    with pytest.raises(ValueError, match="teacher_model_server_url must be set when use_teacher_server=True"):
-        DistillationConfig(
-            **_make_distillation_config_kwargs(tmp_path),
-            use_teacher_server=True,
-            beta=0.5,
-            loss_top_k=1,
-        )
-
-
-@pytest.mark.parametrize(
-    (
-        "input_ids",
-        "attention_mask",
-        "prompt_attention_mask",
-        "labels",
-        "expected_sequences",
-        "expected_prompt_lengths",
-        "expected_completion_lengths",
-    ),
-    [
-        pytest.param(
-            torch.tensor([[0, 0, 10, 11, 12, 20, 21], [30, 31, 32, 33, 34, 40, 41]], dtype=torch.long),
-            torch.tensor([[0, 0, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1]], dtype=torch.long),
-            None,
-            torch.tensor(
-                [[-100, -100, -100, -100, -100, 20, 21], [-100, -100, -100, -100, -100, 40, 41]],
-                dtype=torch.long,
-            ),
-            [[10, 11, 12, 20, 21], [30, 31, 32, 33, 34, 40, 41]],
-            [3, 5],
-            [2, 2],
-            id="variable_prompt_lengths",
-        ),
-        pytest.param(
-            torch.tensor(
-                [[0, 0, 10, 11, 12, 20, 21, 0, 0], [30, 31, 32, 33, 34, 40, 41, 0, 0]],
-                dtype=torch.long,
-            ),
-            torch.tensor(
-                [[0, 0, 1, 1, 1, 1, 1, 0, 0], [1, 1, 1, 1, 1, 1, 1, 0, 0]],
-                dtype=torch.long,
-            ),
-            torch.tensor([[0, 0, 1, 1, 1], [1, 1, 1, 1, 1]], dtype=torch.long),
-            torch.tensor(
-                [
-                    [-100, -100, -100, -100, -100, 20, 21, -100, -100],
-                    [-100, -100, -100, -100, -100, 40, 41, -100, -100],
-                ],
-                dtype=torch.long,
-            ),
-            [[10, 11, 12, 20, 21], [30, 31, 32, 33, 34, 40, 41]],
-            [3, 5],
-            [2, 2],
-            id="padded_on_policy_completions",
-        ),
-        pytest.param(
-            torch.tensor([[10, 11, 12]], dtype=torch.long),
-            torch.tensor([[1, 1, 1]], dtype=torch.long),
-            None,
-            torch.tensor([[-100, -100, -100]], dtype=torch.long),
-            [[10, 11, 12]],
-            [3],
-            [0],
-            id="empty_completion",
-        ),
-    ],
-)
-def test_build_teacher_request_inputs(
-    input_ids,
-    attention_mask,
-    prompt_attention_mask,
-    labels,
-    expected_sequences,
-    expected_prompt_lengths,
-    expected_completion_lengths,
-):
-    sequences, prompt_lengths, completion_lengths = build_teacher_request_inputs(
-        input_ids=input_ids,
-        attention_mask=attention_mask,
-        prompt_attention_mask=prompt_attention_mask,
-        labels=labels,
-    )
-
-    assert sequences == expected_sequences
-    assert prompt_lengths == expected_prompt_lengths
-    assert completion_lengths == expected_completion_lengths
-
-
-class TestGetTeacherTokenLogprobsFromServer(TrlTestCase):
-    def test_variable_lengths_use_neg_inf_sentinel_at_padding(self):
-        mock_self = MagicMock()
-        mock_self.teacher_client.get_sequence_logprobs = MagicMock(return_value=_ragged_server_response())
-        mock_self.loss_top_k = 1
-        mock_self.temperature = 1.0
-
-        inputs = {
-            "input_ids": torch.tensor([[10, 11, 90, 0, 0], [10, 11, 90, 9217, 100]]),
-            "attention_mask": torch.tensor([[1, 1, 1, 0, 0], [1, 1, 1, 1, 1]]),
-            "labels": torch.tensor([[-100, -100, 90, -100, -100], [-100, -100, 90, 9217, 100]]),
-        }
-
-        out = DistillationTrainer._get_teacher_token_logprobs_from_server(mock_self, inputs, aligned_prompt_length=2)
-
-        assert out["actual_logprobs"].shape == (2, 3)
-        assert out["topk_logprobs"].shape == (2, 3, 1)
-
-        # Real completion positions preserved.
-        assert out["actual_logprobs"][0, 0].item() == pytest.approx(-2.3, rel=1e-5)
-        assert out["actual_logprobs"][1, 0].item() == pytest.approx(-1.1, rel=1e-5)
-        assert out["actual_logprobs"][1, 2].item() == pytest.approx(-3.0, rel=1e-5)
-
-        # Sample 0 is 1 token long; positions 1 and 2 are padded with the -inf sentinel.
-        assert out["actual_logprobs"][0, 1].item() == float("-inf")
-        assert out["actual_logprobs"][0, 2].item() == float("-inf")
-        assert out["topk_logprobs"][0, 1, 0].item() == float("-inf")
-
-        # Sample 1 is full-length and fully finite.
-        assert torch.isfinite(out["actual_logprobs"][1, :]).all()
-
-
-class TestServerReverseKLPaddingMask(TrlTestCase):
-    def test_mask_keeps_forward_and_backward_finite(self):
-        # Simulates the getter's output: sample 0 has completion length 1 (positions 1-2
-        # padded with -inf), sample 1 is full-length.
-        teacher_topk = torch.tensor(
-            [[[-2.3], [float("-inf")], [float("-inf")]], [[-1.1], [-0.4], [-3.0]]],
-            dtype=torch.float32,
-        )
-        labels = torch.tensor([[90, -100, -100], [90, 9217, 100]])
-
-        # Strategy B: neutralise -inf at labels == -100 before the divergence math.
-        pad_mask = (labels == -100).unsqueeze(-1)
-        zero = torch.zeros((), dtype=teacher_topk.dtype)
-        teacher_topk = torch.where(pad_mask, zero, teacher_topk)
-
-        valid_mask = torch.ones_like(teacher_topk, dtype=torch.bool)
-        teacher_with_tail, support_mask = _add_tail_bucket(teacher_topk, valid_mask)
-        assert torch.isfinite(teacher_with_tail).all()
-
-        raw_student = torch.randn(2, 3, 2, requires_grad=True)
-        student_log_probs = F.log_softmax(raw_student, dim=-1)
-        loss = _jsd_divergence(student_log_probs, teacher_with_tail, beta=1.0, support_mask=support_mask)
-        assert torch.isfinite(loss).all()
-
-        loss.sum().backward()
-        assert torch.isfinite(raw_student.grad).all()
-
-
 def _reference_generalized_jsd(student_logits, teacher_logits, labels=None, beta=0.5, temperature=1.0):
     """Naive reference for the generalized JSD, written straight from the definition.
 
@@ -512,16 +226,6 @@ def _make_local_trainer(self, **kwargs):
             processing_class=self.tokenizer,
         )
 
-    def _make_server_trainer(self, **kwargs):
-        dataset = load_dataset("trl-internal-testing/zen", "conversational_language_modeling", split="train")
-        return DistillationTrainer(
-            model=self.model_id,
-            teacher_model=None,
-            args=self._make_args(use_teacher_server=True, teacher_model_server_url="http://localhost:8000", **kwargs),
-            train_dataset=dataset,
-            processing_class=self.tokenizer,
-        )
-
     def _make_batch(self, trainer):
         examples = [trainer.train_dataset[i] for i in range(2)]
         return trainer.data_collator(examples)
@@ -715,99 +419,6 @@ def test_distillation_trainer_with_liger(self):
         finally:
             importlib.reload(importlib.import_module(trainer.model.__module__))
 
-    def test_sampled_mode_matches_between_local_and_external_teachers(self, monkeypatch):
-        import trl.generation.vllm_client as vllm_client_module
-
-        teacher_client = RecordingTeacherClient()
-        monkeypatch.setattr(vllm_client_module, "VLLMClient", lambda *args, **kwargs: teacher_client)
-
-        local_trainer = self._make_local_trainer(beta=0.5, loss_top_k=1, reverse_kl_top_1_mode="sampled")
-        server_trainer = self._make_server_trainer(beta=0.5, loss_top_k=1)
-
-        cpu_inputs = self._make_batch(local_trainer)
-        expected_sequences, expected_prompt_lengths, _ = build_teacher_request_inputs(
-            cpu_inputs["input_ids"],
-            cpu_inputs["attention_mask"],
-            prompt_attention_mask=cpu_inputs["prompt_attention_mask"],
-            labels=cpu_inputs["labels"],
-        )
-
-        local_inputs = self._move_batch_to_device(cpu_inputs, local_trainer.accelerator.device)
-        server_inputs = self._move_batch_to_device(cpu_inputs, server_trainer.accelerator.device)
-
-        local_trainer.teacher_model.eval()
-        with torch.no_grad():
-            teacher_logits = local_trainer.teacher_model(
-                input_ids=local_inputs["input_ids"],
-                attention_mask=local_inputs["attention_mask"],
-            ).logits
-            teacher_client.result = _build_server_result(
-                teacher_logits,
-                local_inputs,
-                temperature=local_trainer.temperature,
-            )
-            local_loss = local_trainer.compute_loss(local_trainer.model, local_inputs)
-            server_loss = server_trainer.compute_loss(server_trainer.model, server_inputs)
-
-        assert teacher_client.calls[0]["sequences"] == expected_sequences
-        assert teacher_client.calls[0]["prompt_lengths"] == expected_prompt_lengths
-        assert teacher_client.calls[0]["top_logprobs"] == 1
-        torch.testing.assert_close(local_loss, server_loss)
-
-    @classmethod
-    def setup_class(cls):
-        model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
-        cls.device = "cuda" if torch.cuda.is_available() else "cpu"
-        cls.tokenizer = AutoTokenizer.from_pretrained(model_id)
-        cls.tokenizer.pad_token = cls.tokenizer.eos_token
-        cls.model_id = model_id
-
-    def _run_one_step(self, bs, ga, monkeypatch):
-        from trl.generation import vllm_client as vllm_client_module
-
-        fake_client = MagicMock()
-        fake_client.get_sequence_logprobs.side_effect = _canned_teacher_logprobs
-        monkeypatch.setattr(vllm_client_module, "VLLMClient", lambda *a, **kw: fake_client)
-
-        config = DistillationConfig(
-            output_dir=self.tmp_dir,
-            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,
-            use_teacher_server=True,
-            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",
-            logging_steps=1,
-            use_cpu=not torch.cuda.is_available(),
-            bf16=False,
-        )
-        model = AutoModelForCausalLM.from_pretrained(self.model_id, dtype=torch.float32).to(self.device)
-        trainer = DistillationTrainer(
-            model=model,
-            args=config,
-            train_dataset=_variable_length_dataset(),
-            processing_class=self.tokenizer,
-        )
-        trainer.teacher_client = fake_client
-        trainer.train()
-        return [rec for rec in trainer.state.log_history if "grad_norm" in rec]
-
-    @pytest.mark.parametrize(("bs", "ga"), [(1, 2), (2, 1)])
-    def test_reverse_kl_finite_grad_with_ragged_batch(self, bs, ga, monkeypatch):
-        records = self._run_one_step(bs=bs, ga=ga, monkeypatch=monkeypatch)
-        assert records, "Expected at least one grad_norm log entry during training"
-        for record in records:
-            assert math.isfinite(record["grad_norm"]), f"grad_norm={record['grad_norm']} leaked -inf into backward"
-            assert math.isfinite(record["loss"])
-
 
 def test_repeat_batch_dataloader_delegates_set_epoch_via_getattr():
     class DummyDataLoader:
diff --git a/tests/experimental/test_server_distillation_trainer.py b/tests/experimental/test_server_distillation_trainer.py
new file mode 100644
index 00000000000..4334865d9d0
--- /dev/null
+++ b/tests/experimental/test_server_distillation_trainer.py
@@ -0,0 +1,456 @@
+# Copyright 2020-2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import math
+from unittest.mock import MagicMock
+
+import pytest
+import torch
+import torch.nn.functional as F
+from datasets import Dataset, load_dataset
+from transformers import AutoModelForCausalLM, AutoTokenizer
+
+from trl.experimental.distillation import DistillationConfig, DistillationTrainer
+from trl.experimental.distillation.distillation_trainer import _add_tail_bucket, _jsd_divergence
+from trl.experimental.server_distillation import (
+    ServerDistillationConfig,
+    ServerDistillationTrainer,
+    build_teacher_request_inputs,
+)
+
+from ..testing_utils import TrlTestCase
+
+
+def _make_server_config_kwargs(tmp_path):
+    return {
+        "output_dir": str(tmp_path),
+        "report_to": "none",
+        "use_cpu": True,
+        "bf16": False,
+        "teacher_model_server_url": "http://localhost:8000",
+    }
+
+
+def _build_server_result(teacher_logits, inputs, temperature=1.0):
+    """Simulate a vLLM server response with variable-length per-sample completions."""
+    _, _, completion_lengths = build_teacher_request_inputs(
+        inputs["input_ids"],
+        inputs["attention_mask"],
+        prompt_attention_mask=inputs.get("prompt_attention_mask"),
+        labels=inputs.get("labels"),
+    )
+
+    label_mask = inputs["labels"] != -100
+    actual_logprobs = []
+    logprobs = []
+    logprob_token_ids = []
+
+    for i, comp_len in enumerate(completion_lengths):
+        if comp_len == 0:
+            actual_logprobs.append([])
+            logprobs.append([])
+            logprob_token_ids.append([])
+            continue
+
+        comp_start = int(torch.nonzero(label_mask[i], as_tuple=False)[0].item())
+        sample_logits = teacher_logits[i, comp_start - 1 : comp_start - 1 + comp_len, :]
+        sample_log_probs = F.log_softmax(sample_logits / temperature, dim=-1)
+        comp_tokens = inputs["input_ids"][i, comp_start : comp_start + comp_len]
+
+        top1_ids = sample_logits.argmax(dim=-1, keepdim=True)
+        top1_lps = sample_log_probs.gather(dim=-1, index=top1_ids)
+        actual_lps = sample_log_probs.gather(dim=-1, index=comp_tokens.unsqueeze(-1))
+
+        actual_logprobs.append(actual_lps.tolist())
+        logprobs.append(top1_lps.tolist())
+        logprob_token_ids.append(top1_ids.tolist())
+
+    return {
+        "actual_logprobs": actual_logprobs,
+        "logprobs": logprobs,
+        "logprob_token_ids": logprob_token_ids,
+    }
+
+
+class RecordingTeacherClient:
+    def __init__(self):
+        self.calls = []
+        self.result = None
+
+    def get_sequence_logprobs(self, sequences, prompt_lengths, top_logprobs, temperature):
+        self.calls.append(
+            {
+                "sequences": sequences,
+                "prompt_lengths": prompt_lengths,
+                "top_logprobs": top_logprobs,
+                "temperature": temperature,
+            }
+        )
+        return self.result
+
+
+def _ragged_server_response():
+    # Two samples with completion lengths 1 and 3 respectively; matches the wire format
+    # of VLLMClient.get_sequence_logprobs (per-sample shape (comp_len, top_k=1)).
+    return {
+        "logprobs": [[[-2.3]], [[-1.1], [-0.4], [-3.0]]],
+        "logprob_token_ids": [[[90]], [[90], [9217], [100]]],
+        "actual_logprobs": [[[-2.3]], [[-1.1], [-0.4], [-3.0]]],
+    }
+
+
+def _canned_teacher_logprobs(**kwargs):
+    # Fabricate ragged per-sample logprobs matching the requested sequence shapes.
+    sequences = kwargs["sequences"]
+    prompt_lengths = kwargs["prompt_lengths"]
+    top_k = kwargs.get("top_logprobs", 1)
+    logprobs, token_ids, actual = [], [], []
+    for seq, plen in zip(sequences, prompt_lengths, strict=True):
+        comp_len = len(seq) - plen
+        logprobs.append([[-1.0 - 0.05 * i] * top_k for i in range(comp_len)])
+        token_ids.append([[int(seq[plen + i])] * top_k for i in range(comp_len)])
+        actual.append([[-1.0 - 0.05 * i] for i in range(comp_len)])
+    return {"logprobs": logprobs, "logprob_token_ids": token_ids, "actual_logprobs": actual}
+
+
+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.",
+                    },
+                ]
+            },
+        ]
+    )
+
+
+def test_config_rejects_liger(tmp_path):
+    with pytest.raises(ValueError, match="use_liger_kernel=True is not supported by ServerDistillationTrainer"):
+        ServerDistillationConfig(**_make_server_config_kwargs(tmp_path), use_liger_kernel=True)
+
+
+def test_config_rejects_reverse_kl_argmax(tmp_path):
+    with pytest.raises(ValueError, match="reverse_kl_top_1_mode='argmax' is not supported"):
+        ServerDistillationConfig(**_make_server_config_kwargs(tmp_path), reverse_kl_top_1_mode="argmax")
+
+
+def test_config_rejects_mixed_loss_without_top_1(tmp_path):
+    with pytest.raises(ValueError, match="loss_top_k must be 1 with beta>0"):
+        ServerDistillationConfig(**_make_server_config_kwargs(tmp_path), beta=0.5, loss_top_k=2)
+
+
+def test_config_requires_server_url(tmp_path):
+    kwargs = _make_server_config_kwargs(tmp_path)
+    kwargs.pop("teacher_model_server_url")
+    with pytest.raises(ValueError, match="teacher_model_server_url must be set for ServerDistillationTrainer"):
+        ServerDistillationConfig(**kwargs, beta=0.5, loss_top_k=1)
+
+
+@pytest.mark.parametrize(
+    (
+        "input_ids",
+        "attention_mask",
+        "prompt_attention_mask",
+        "labels",
+        "expected_sequences",
+        "expected_prompt_lengths",
+        "expected_completion_lengths",
+    ),
+    [
+        pytest.param(
+            torch.tensor([[0, 0, 10, 11, 12, 20, 21], [30, 31, 32, 33, 34, 40, 41]], dtype=torch.long),
+            torch.tensor([[0, 0, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1]], dtype=torch.long),
+            None,
+            torch.tensor(
+                [[-100, -100, -100, -100, -100, 20, 21], [-100, -100, -100, -100, -100, 40, 41]],
+                dtype=torch.long,
+            ),
+            [[10, 11, 12, 20, 21], [30, 31, 32, 33, 34, 40, 41]],
+            [3, 5],
+            [2, 2],
+            id="variable_prompt_lengths",
+        ),
+        pytest.param(
+            torch.tensor(
+                [[0, 0, 10, 11, 12, 20, 21, 0, 0], [30, 31, 32, 33, 34, 40, 41, 0, 0]],
+                dtype=torch.long,
+            ),
+            torch.tensor(
+                [[0, 0, 1, 1, 1, 1, 1, 0, 0], [1, 1, 1, 1, 1, 1, 1, 0, 0]],
+                dtype=torch.long,
+            ),
+            torch.tensor([[0, 0, 1, 1, 1], [1, 1, 1, 1, 1]], dtype=torch.long),
+            torch.tensor(
+                [
+                    [-100, -100, -100, -100, -100, 20, 21, -100, -100],
+                    [-100, -100, -100, -100, -100, 40, 41, -100, -100],
+                ],
+                dtype=torch.long,
+            ),
+            [[10, 11, 12, 20, 21], [30, 31, 32, 33, 34, 40, 41]],
+            [3, 5],
+            [2, 2],
+            id="padded_on_policy_completions",
+        ),
+        pytest.param(
+            torch.tensor([[10, 11, 12]], dtype=torch.long),
+            torch.tensor([[1, 1, 1]], dtype=torch.long),
+            None,
+            torch.tensor([[-100, -100, -100]], dtype=torch.long),
+            [[10, 11, 12]],
+            [3],
+            [0],
+            id="empty_completion",
+        ),
+    ],
+)
+def test_build_teacher_request_inputs(
+    input_ids,
+    attention_mask,
+    prompt_attention_mask,
+    labels,
+    expected_sequences,
+    expected_prompt_lengths,
+    expected_completion_lengths,
+):
+    sequences, prompt_lengths, completion_lengths = build_teacher_request_inputs(
+        input_ids=input_ids,
+        attention_mask=attention_mask,
+        prompt_attention_mask=prompt_attention_mask,
+        labels=labels,
+    )
+
+    assert sequences == expected_sequences
+    assert prompt_lengths == expected_prompt_lengths
+    assert completion_lengths == expected_completion_lengths
+
+
+class TestGetTeacherTokenLogprobsFromServer(TrlTestCase):
+    def test_variable_lengths_use_neg_inf_sentinel_at_padding(self):
+        mock_self = MagicMock()
+        mock_self.teacher_client.get_sequence_logprobs = MagicMock(return_value=_ragged_server_response())
+        mock_self.loss_top_k = 1
+        mock_self.temperature = 1.0
+
+        inputs = {
+            "input_ids": torch.tensor([[10, 11, 90, 0, 0], [10, 11, 90, 9217, 100]]),
+            "attention_mask": torch.tensor([[1, 1, 1, 0, 0], [1, 1, 1, 1, 1]]),
+            "labels": torch.tensor([[-100, -100, 90, -100, -100], [-100, -100, 90, 9217, 100]]),
+        }
+
+        out = ServerDistillationTrainer._get_teacher_token_logprobs_from_server(
+            mock_self, inputs, aligned_prompt_length=2
+        )
+
+        assert out["actual_logprobs"].shape == (2, 3)
+        assert out["topk_logprobs"].shape == (2, 3, 1)
+
+        # Real completion positions preserved.
+        assert out["actual_logprobs"][0, 0].item() == pytest.approx(-2.3, rel=1e-5)
+        assert out["actual_logprobs"][1, 0].item() == pytest.approx(-1.1, rel=1e-5)
+        assert out["actual_logprobs"][1, 2].item() == pytest.approx(-3.0, rel=1e-5)
+
+        # Sample 0 is 1 token long; positions 1 and 2 are padded with the -inf sentinel.
+        assert out["actual_logprobs"][0, 1].item() == float("-inf")
+        assert out["actual_logprobs"][0, 2].item() == float("-inf")
+        assert out["topk_logprobs"][0, 1, 0].item() == float("-inf")
+
+        # Sample 1 is full-length and fully finite.
+        assert torch.isfinite(out["actual_logprobs"][1, :]).all()
+
+
+class TestServerReverseKLPaddingMask(TrlTestCase):
+    def test_mask_keeps_forward_and_backward_finite(self):
+        # Simulates the getter's output: sample 0 has completion length 1 (positions 1-2
+        # padded with -inf), sample 1 is full-length.
+        teacher_topk = torch.tensor(
+            [[[-2.3], [float("-inf")], [float("-inf")]], [[-1.1], [-0.4], [-3.0]]],
+            dtype=torch.float32,
+        )
+        labels = torch.tensor([[90, -100, -100], [90, 9217, 100]])
+
+        # Strategy B: neutralise -inf at labels == -100 before the divergence math.
+        pad_mask = (labels == -100).unsqueeze(-1)
+        zero = torch.zeros((), dtype=teacher_topk.dtype)
+        teacher_topk = torch.where(pad_mask, zero, teacher_topk)
+
+        valid_mask = torch.ones_like(teacher_topk, dtype=torch.bool)
+        teacher_with_tail, support_mask = _add_tail_bucket(teacher_topk, valid_mask)
+        assert torch.isfinite(teacher_with_tail).all()
+
+        raw_student = torch.randn(2, 3, 2, requires_grad=True)
+        student_log_probs = F.log_softmax(raw_student, dim=-1)
+        loss = _jsd_divergence(student_log_probs, teacher_with_tail, beta=1.0, support_mask=support_mask)
+        assert torch.isfinite(loss).all()
+
+        loss.sum().backward()
+        assert torch.isfinite(raw_student.grad).all()
+
+
+class TestServerVsLocalTeacher(TrlTestCase):
+    def setup_method(self):
+        self.model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)
+        self.tokenizer.pad_token = self.tokenizer.eos_token
+
+    def _local_args(self, **kwargs):
+        args = {
+            "output_dir": self.tmp_dir,
+            "per_device_train_batch_size": 2,
+            "gradient_accumulation_steps": 1,
+            "max_steps": 1,
+            "save_strategy": "no",
+            "report_to": "none",
+            "use_cpu": True,
+            "bf16": False,
+            "lmbda": 0.0,
+            "max_length": 128,
+            "max_completion_length": 32,
+            "model_init_kwargs": {"dtype": "float32", "device_map": None},
+            "teacher_model_init_kwargs": {"dtype": "float32", "device_map": None},
+        }
+        args.update(kwargs)
+        return DistillationConfig(**args)
+
+    def _make_local_trainer(self, **kwargs):
+        dataset = load_dataset("trl-internal-testing/zen", "conversational_language_modeling", split="train")
+        return DistillationTrainer(
+            model=self.model_id,
+            teacher_model=self.model_id,
+            args=self._local_args(**kwargs),
+            train_dataset=dataset,
+            processing_class=self.tokenizer,
+        )
+
+    def _make_server_trainer(self, **kwargs):
+        dataset = load_dataset("trl-internal-testing/zen", "conversational_language_modeling", split="train")
+        args = self._local_args(**kwargs)
+        server_args = ServerDistillationConfig(
+            **{k: v for k, v in vars(args).items() if not k.startswith("_")},
+            teacher_model_server_url="http://localhost:8000",
+        )
+        return ServerDistillationTrainer(
+            model=self.model_id,
+            args=server_args,
+            train_dataset=dataset,
+            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_sampled_mode_matches_between_local_and_external_teachers(self, monkeypatch):
+        import trl.generation.vllm_client as vllm_client_module
+
+        teacher_client = RecordingTeacherClient()
+        monkeypatch.setattr(vllm_client_module, "VLLMClient", lambda *args, **kwargs: teacher_client)
+
+        local_trainer = self._make_local_trainer(beta=0.5, loss_top_k=1, reverse_kl_top_1_mode="sampled")
+        server_trainer = self._make_server_trainer(beta=0.5, loss_top_k=1)
+
+        cpu_inputs = self._make_batch(local_trainer)
+        expected_sequences, expected_prompt_lengths, _ = build_teacher_request_inputs(
+            cpu_inputs["input_ids"],
+            cpu_inputs["attention_mask"],
+            prompt_attention_mask=cpu_inputs["prompt_attention_mask"],
+            labels=cpu_inputs["labels"],
+        )
+
+        local_inputs = self._move_batch_to_device(cpu_inputs, local_trainer.accelerator.device)
+        server_inputs = self._move_batch_to_device(cpu_inputs, server_trainer.accelerator.device)
+
+        local_trainer.teacher_model.eval()
+        with torch.no_grad():
+            teacher_logits = local_trainer.teacher_model(
+                input_ids=local_inputs["input_ids"],
+                attention_mask=local_inputs["attention_mask"],
+            ).logits
+            teacher_client.result = _build_server_result(
+                teacher_logits,
+                local_inputs,
+                temperature=local_trainer.temperature,
+            )
+            local_loss = local_trainer.compute_loss(local_trainer.model, local_inputs)
+            server_loss = server_trainer.compute_loss(server_trainer.model, server_inputs)
+
+        assert teacher_client.calls[0]["sequences"] == expected_sequences
+        assert teacher_client.calls[0]["prompt_lengths"] == expected_prompt_lengths
+        assert teacher_client.calls[0]["top_logprobs"] == 1
+        torch.testing.assert_close(local_loss, server_loss)
+
+
+class TestServerDistillationTrainerRaggedGrad(TrlTestCase):
+    @classmethod
+    def setup_class(cls):
+        model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
+        cls.device = "cuda" if torch.cuda.is_available() else "cpu"
+        cls.tokenizer = AutoTokenizer.from_pretrained(model_id)
+        cls.tokenizer.pad_token = cls.tokenizer.eos_token
+        cls.model_id = model_id
+
+    def _run_one_step(self, bs, ga, monkeypatch):
+        from trl.generation import vllm_client as vllm_client_module
+
+        fake_client = MagicMock()
+        fake_client.get_sequence_logprobs.side_effect = _canned_teacher_logprobs
+        monkeypatch.setattr(vllm_client_module, "VLLMClient", lambda *a, **kw: fake_client)
+
+        config = ServerDistillationConfig(
+            output_dir=self.tmp_dir,
+            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,
+            beta=1.0,
+            lmbda=0.0,
+            loss_add_tail=True,
+            save_strategy="no",
+            report_to="none",
+            logging_steps=1,
+            use_cpu=not torch.cuda.is_available(),
+            bf16=False,
+        )
+        model = AutoModelForCausalLM.from_pretrained(self.model_id, dtype=torch.float32).to(self.device)
+        trainer = ServerDistillationTrainer(
+            model=model,
+            args=config,
+            train_dataset=_variable_length_dataset(),
+            processing_class=self.tokenizer,
+        )
+        trainer.teacher_client = fake_client
+        trainer.train()
+        return [rec for rec in trainer.state.log_history if "grad_norm" in rec]
+
+    @pytest.mark.parametrize(("bs", "ga"), [(1, 2), (2, 1)])
+    def test_reverse_kl_finite_grad_with_ragged_batch(self, bs, ga, monkeypatch):
+        records = self._run_one_step(bs=bs, ga=ga, monkeypatch=monkeypatch)
+        assert records, "Expected at least one grad_norm log entry during training"
+        for record in records:
+            assert math.isfinite(record["grad_norm"]), f"grad_norm={record['grad_norm']} leaked -inf into backward"
+            assert math.isfinite(record["loss"])
diff --git a/trl/experimental/distillation/distillation_config.py b/trl/experimental/distillation/distillation_config.py
index c70dd4b2d6d..5b7efc9f7c8 100644
--- a/trl/experimental/distillation/distillation_config.py
+++ b/trl/experimental/distillation/distillation_config.py
@@ -79,20 +79,12 @@ class DistillationConfig(_BaseConfig):
         teacher_model_init_kwargs (`dict[str, Any]` or `None`, *optional*):
             Keyword arguments passed to `AutoModelForCausalLM.from_pretrained` when instantiating the teacher model
             from a string.
-        use_teacher_server (`bool`, *optional*, defaults to `False`):
-            Whether to use an external vLLM teacher server instead of a local teacher model.
-        teacher_model_server_url (`str` or `None`, *optional*):
-            Base URL of a vLLM server hosting the teacher model (e.g., `"http://localhost:8000"`). When set, teacher
-            logprobs are fetched from the server instead of running a local forward pass when `use_teacher_server=True`.
         loss_top_k (`int`, *optional*, defaults to `1`):
             Number of top tokens to use when computing the JSD/KL loss. Both student and teacher distributions are
             restricted to these K tokens and re-normalized before computing divergence. If 0, the full vocabulary
             is used. For local teachers, the general support rule is teacher top-k for forward KL, student top-k for
             reverse KL, and the union for mixed JSD. When `beta > 0` and `loss_top_k == 1`, the forward support still
             uses the teacher's top-1 token, while the reverse top-1 token is controlled by `reverse_kl_top_1_mode`.
-            When `use_teacher_server=True`, the pure forward path (`beta=0`) requires this to be positive and uses the
-            teacher's top-k logprobs for the forward term. When `beta > 0`, server-backed distillation requires
-            `loss_top_k == 1` and only supports `"sampled"` reverse top-1 tokens.
         loss_add_tail (`bool`, *optional*, defaults to `True`):
             Whether to append a tail bucket that represents the remaining probability mass outside the selected top-k
             support when computing the loss.
@@ -242,19 +234,6 @@ class DistillationConfig(_BaseConfig):
             "help": "Keyword arguments for `AutoModelForCausalLM.from_pretrained` when instantiating the teacher."
         },
     )
-
-    # Teacher model (external vLLM server)
-    use_teacher_server: bool = field(
-        default=False,
-        metadata={"help": "Whether to use an external vLLM teacher server instead of a local teacher model."},
-    )
-    teacher_model_server_url: str | None = field(
-        default=None,
-        metadata={
-            "help": 'Base URL of a vLLM server hosting the teacher model (e.g., "http://localhost:8000"). '
-            "Required when use_teacher_server=True."
-        },
-    )
     loss_top_k: int = field(
         default=1,
         metadata={
@@ -264,10 +243,7 @@ class DistillationConfig(_BaseConfig):
             "union of both for JSD) and re-normalized before computing divergence. "
             "If 0, the full vocabulary is used (slower but exact). "
             "When beta > 0 and loss_top_k == 1, the forward support still uses the teacher's top-1 token, "
-            "while the reverse top-1 token is controlled by reverse_kl_top_1_mode. "
-            "When use_teacher_server=True, beta=0 requires loss_top_k > 0 and uses the teacher's top-k "
-            "logprobs for the forward term. When beta > 0, server-backed distillation requires loss_top_k == 1 "
-            "and only supports 'sampled' reverse top-1 tokens."
+            "while the reverse top-1 token is controlled by reverse_kl_top_1_mode."
         },
     )
     loss_add_tail: bool = field(
@@ -410,33 +386,6 @@ def __post_init__(self):
                 f"{self.per_device_train_batch_size} * {self.gradient_accumulation_steps}."
             )
 
-        if self.use_teacher_server and self.use_liger_kernel:
-            raise ValueError(
-                "use_liger_kernel=True is not supported with use_teacher_server=True because the Liger loss path "
-                "requires a local teacher model."
-            )
-        if self.use_teacher_server and (
-            self.teacher_model_server_url is None or not self.teacher_model_server_url.strip()
-        ):
-            raise ValueError("teacher_model_server_url must be set when use_teacher_server=True.")
-
-        if self.use_teacher_server and self.beta == 0 and self.loss_top_k < 1:
-            raise ValueError(
-                f"loss_top_k must be positive when using use_teacher_server=True with beta=0 "
-                f"(got loss_top_k={self.loss_top_k}). The pure forward server path only has access to the "
-                f"teacher's top-k logprobs, so it cannot compute the exact full-vocabulary loss when loss_top_k=0."
-            )
-        if self.use_teacher_server and self.reverse_kl_top_1_mode == "argmax":
-            raise ValueError(
-                "reverse_kl_top_1_mode='argmax' is not supported with use_teacher_server=True because the server "
-                "cannot provide teacher logprobs for arbitrary student-selected tokens."
-            )
-        if self.use_teacher_server and self.beta > 0 and self.loss_top_k != 1:
-            raise ValueError(
-                f"loss_top_k must be 1 when using use_teacher_server=True with beta>0 "
-                f"(got loss_top_k={self.loss_top_k}). Mixed forward/reverse distillation with an external teacher "
-                "is only implemented for top-1 support."
-            )
         if self.reverse_kl_top_1_mode != "sampled" and (self.beta == 0 or self.loss_top_k != 1):
             warnings.warn(
                 f"reverse_kl_top_1_mode='{self.reverse_kl_top_1_mode}' has no effect when beta={self.beta} "
diff --git a/trl/experimental/distillation/distillation_trainer.py b/trl/experimental/distillation/distillation_trainer.py
index be6a1632e95..dd50fc96243 100644
--- a/trl/experimental/distillation/distillation_trainer.py
+++ b/trl/experimental/distillation/distillation_trainer.py
@@ -160,51 +160,6 @@ def _jsd_divergence(student_log_probs, teacher_log_probs, beta, support_mask=Non
             return beta_t * kl_teacher + (1 - beta_t) * kl_student
 
 
-def build_teacher_request_inputs(
-    input_ids: torch.Tensor,
-    attention_mask: torch.Tensor,
-    prompt_attention_mask: torch.Tensor | None = None,
-    labels: torch.Tensor | None = None,
-) -> tuple[list[list[int]], list[int], list[int]]:
-    """Trim padded batch tensors into per-sample sequences for teacher-server requests."""
-
-    if input_ids.shape != attention_mask.shape:
-        raise ValueError(
-            f"input_ids and attention_mask must have the same shape, got {input_ids.shape} and {attention_mask.shape}."
-        )
-
-    input_ids_cpu = input_ids.detach().cpu()
-    attention_mask_cpu = attention_mask.detach().cpu().bool()
-
-    if prompt_attention_mask is not None:
-        prompt_lengths = prompt_attention_mask.detach().cpu().sum(dim=1).to(torch.long)
-    else:
-        if labels is None:
-            raise ValueError("labels are required when prompt_attention_mask is not provided.")
-        if labels.shape != input_ids.shape:
-            raise ValueError(f"labels must match input_ids shape, got {labels.shape} and {input_ids.shape}.")
-        full_lengths = attention_mask_cpu.sum(dim=1).to(torch.long)
-        completion_lengths = (labels.detach().cpu() != -100).sum(dim=1).to(torch.long)
-        prompt_lengths = full_lengths - completion_lengths
-
-    trimmed_input_ids: list[list[int]] = []
-    prompt_lengths_list: list[int] = []
-    completion_lengths_list: list[int] = []
-
-    for row, mask, prompt_length in zip(input_ids_cpu, attention_mask_cpu, prompt_lengths, strict=True):
-        trimmed_row = row[mask]
-        prompt_len = int(prompt_length.item())
-        if prompt_len < 0 or prompt_len > trimmed_row.numel():
-            raise ValueError(
-                f"Invalid prompt length {prompt_len} for trimmed sequence of length {trimmed_row.numel()}."
-            )
-        trimmed_input_ids.append(trimmed_row.tolist())
-        prompt_lengths_list.append(prompt_len)
-        completion_lengths_list.append(int(trimmed_row.numel()) - prompt_len)
-
-    return trimmed_input_ids, prompt_lengths_list, completion_lengths_list
-
-
 class _DistillationCollator:
     """Data collator for the distillation trainer with independent prompt/completion budgets.
 
@@ -493,16 +448,8 @@ def __init__(
             self._forward_redirection = _ForwardRedirection()
 
         # ── Teacher model setup ──
-        self.teacher_client = None
-        self.use_teacher_server = args.use_teacher_server
-        self.teacher_model_server_url = args.teacher_model_server_url
         self._local_teacher_tokenizer_matches_student = True
-        if self.use_teacher_server:
-            from ...generation.vllm_client import VLLMClient
-
-            self.teacher_client = VLLMClient(base_url=self.teacher_model_server_url, connection_timeout=60.0)
-            teacher_model = None
-        elif teacher_model is not None:
+        if teacher_model is not None:
             if args.teacher_model_init_kwargs is not None and not isinstance(teacher_model, str):
                 raise ValueError(
                     "You passed teacher_model_init_kwargs to the config, but your teacher_model is already "
@@ -686,8 +633,8 @@ def _raise_if_local_teacher_tokenizer_mismatch(self) -> None:
         if self.teacher_model is not None and not self._local_teacher_tokenizer_matches_student:
             raise ValueError(
                 "DistillationTrainer's built-in local-teacher loss only supports student/teacher pairs that use "
-                "the same tokenizer. Use a same-tokenizer local teacher, set `use_teacher_server=True`, or "
-                "override the local teacher loss path in a subclass."
+                "the same tokenizer. Use a same-tokenizer local teacher, use ServerDistillationTrainer with an "
+                "external teacher server, or override the local teacher loss path in a subclass."
             )
 
     def _compute_prompt_length(self, inputs: dict[str, torch.Tensor | Any]) -> int:
@@ -1267,218 +1214,14 @@ def _compute_local_sparse_top_1_divergence_loss(
 
     def _get_teacher_logits(self, inputs: dict[str, torch.Tensor | Any]) -> torch.Tensor:
         """Get teacher logits — dispatches between local model and external server."""
-        if self.teacher_model is not None:
-            self.teacher_model.eval()
-            with torch.no_grad():
-                return self.teacher_model(
-                    input_ids=inputs["input_ids"],
-                    attention_mask=inputs["attention_mask"],
-                ).logits
-        elif self.use_teacher_server:
-            raise NotImplementedError(
-                "Fetching full teacher logits with use_teacher_server=True is not supported. "
-                "Server-backed distillation only supports per-token logprobs via "
-                "`_get_teacher_token_logprobs_from_server`."
-            )
-        else:
-            raise ValueError("No teacher model or teacher server configured.")
-
-    def _get_teacher_token_logprobs_from_server(
-        self,
-        inputs: dict[str, torch.Tensor | Any],
-        aligned_prompt_length: int,
-    ) -> dict[str, torch.Tensor]:
-        """Fetch per-token teacher logprobs from an external vLLM server.
-
-        Returns a dict with:
-            ``actual_logprobs`` – (batch, completion_length) teacher log-prob for the actual
-                                   token at each position (for reverse KL).
-            ``topk_logprobs`` – (batch, completion_length, K) teacher top-k sorted logprobs
-                                   (for forward KL).
-            ``topk_token_ids`` – (batch, completion_length, K) corresponding token IDs.
-        """
-        import numpy as np
-
-        input_ids = inputs["input_ids"]
-        batch_size = input_ids.shape[0]
-        sequences, prompt_lengths, completion_lengths = build_teacher_request_inputs(
-            input_ids,
-            inputs["attention_mask"],
-            prompt_attention_mask=inputs.get("prompt_attention_mask"),
-            labels=inputs.get("labels"),
-        )
-
-        # The pure forward server path can use the requested teacher top-k support.
-        # When beta > 0, config validation restricts the server-backed path to top-1.
-        requested_top_k = self.loss_top_k
-        result = self.teacher_client.get_sequence_logprobs(
-            sequences=sequences,
-            prompt_lengths=prompt_lengths,
-            top_logprobs=requested_top_k,
-            temperature=self.temperature,
-        )
-        K = requested_top_k
-
-        device = input_ids.device
-        labels = inputs.get("labels")
-        if labels is None:
-            raise ValueError("labels are required to align teacher-server logprobs with the student loss tensors.")
-
-        # The student loss slices tensors in padded-sequence coordinates starting at `aligned_prompt_length`.
-        # Place each teacher completion into that same coordinate system by locating the first non-masked completion
-        # token in `labels`. This works for both left-padded off-policy batches and on-policy batches where
-        # completions are right-padded after a fixed-width prompt block.
-        completion_offsets = []
-        label_mask = labels != -100
-        for sample_mask, comp_len in zip(label_mask, completion_lengths, strict=True):
-            if comp_len == 0:
-                completion_offsets.append(0)
-                continue
-            completion_start = int(torch.nonzero(sample_mask, as_tuple=False)[0].item())
-            completion_offsets.append(completion_start - aligned_prompt_length)
-
-        # Size the output tensors to tightly fit the teacher logprobs. Using the full padded
-        # sequence length would include padding positions with -inf teacher logprobs, producing
-        # +inf in the forward pass and NaN gradients in the backward pass (0 * inf = NaN).
-        # Shorter samples in variable-length batches still need the -inf sentinel at the tail;
-        # downstream loss consumers (_compute_server_sparse_top_1_divergence_loss,
-        # _compute_server_forward_kl_loss) neutralise those positions before the divergence
-        # math runs.
-        completion_length = max(
-            (offset + len(lps) for offset, lps in zip(completion_offsets, result["logprobs"], strict=True)),
-            default=0,
-        )
-
-        # actual_logprobs: (B, T) — teacher logprob for the actual token
-        def _actual_to_tensor(key):
-            arr = np.full((batch_size, completion_length), float("-inf"), dtype=np.float32)
-            for i, (offset, seq_lps) in enumerate(zip(completion_offsets, result[key], strict=True)):
-                if seq_lps:
-                    vals = np.array(seq_lps, dtype=np.float32)  # (comp_len_i, 1)
-                    arr[i, offset : offset + vals.shape[0]] = vals[:, 0]
-            return torch.from_numpy(arr).to(device)
-
-        # topk: (B, T, K)
-        def _topk_to_tensor(key, k, np_dtype, fill):
-            arr = np.full((batch_size, completion_length, k), fill, dtype=np_dtype)
-            for i, (offset, seq_vals) in enumerate(zip(completion_offsets, result[key], strict=True)):
-                if seq_vals:
-                    vals = np.array(seq_vals, dtype=np_dtype)  # (comp_len_i, k)
-                    arr[i, offset : offset + vals.shape[0], :] = vals
-            return torch.from_numpy(arr).to(device)
-
-        return {
-            "actual_logprobs": _actual_to_tensor("actual_logprobs"),
-            "topk_logprobs": _topk_to_tensor("logprobs", K, np.float32, float("-inf")),
-            "topk_token_ids": _topk_to_tensor("logprob_token_ids", K, np.int64, 0),
-        }
-
-    def _compute_server_sparse_top_1_divergence_loss(
-        self,
-        teacher_result: dict[str, torch.Tensor],
-        student_log_probs: torch.Tensor,
-        completion_tokens: torch.Tensor,
-        labels: torch.Tensor,
-    ) -> torch.Tensor:
-        """Compute exact sparse top-1 generalized JSD/KL from server-provided teacher logprobs.
-
-        Args:
-            teacher_result: dict with ``actual_logprobs`` (B, T), ``topk_logprobs`` (B, T, K),
-                ``topk_token_ids`` (B, T, K).
-            student_log_probs: (B, T, V) student log-softmax over vocabulary.
-            completion_tokens: (B, T) actual token IDs in the completion.
-            labels: (B, T) with -100 for positions to ignore.
-        """
-        topk_teacher_lps = teacher_result["topk_logprobs"]  # (B, T, 1)
-        topk_token_ids = teacher_result["topk_token_ids"]  # (B, T, 1)
-        actual_teacher_lps = teacher_result["actual_logprobs"]  # (B, T)
-        required = labels != -100
-
-        missing_actual = required & ~torch.isfinite(actual_teacher_lps)
-        if missing_actual.any():
-            missing_count = int(missing_actual.sum().item())
-            total_required = int(required.sum().item())
-            raise ValueError(
-                "Teacher server is missing actual-token logprobs for required reverse-KL positions: "
-                f"{missing_count}/{total_required}."
-            )
-        if self.beta < 1:
-            teacher_top1_logprobs = topk_teacher_lps.squeeze(-1)
-            missing_top1 = required & ~torch.isfinite(teacher_top1_logprobs)
-            if missing_top1.any():
-                missing_count = int(missing_top1.sum().item())
-                total_required = int(required.sum().item())
-                raise ValueError(
-                    "Teacher server is missing top-1 logprobs for required forward-KL positions: "
-                    f"{missing_count}/{total_required}."
-                )
-
-        # Replace -inf teacher logprobs at intra-batch padding (labels == -100) with 0 so
-        # reverse-KL's student_probs·(log_s - log_t) does not leak +inf into the backward pass.
-        pad_mask_2d = ~required
-        pad_mask_3d = pad_mask_2d.unsqueeze(-1)
-        topk_teacher_lps = torch.where(pad_mask_3d, 0.0, topk_teacher_lps)
-        actual_teacher_lps = torch.where(pad_mask_2d, 0.0, actual_teacher_lps)
-
-        # Server path only supports "sampled" mode — config validation enforces this, but we guard
-        # explicitly so future relaxations of the config check don't silently change behaviour.
-        reverse_token_ids = self._get_reverse_kl_top_1_tokens(student_log_probs, completion_tokens)
-        # The server path normalizes locally (batchmean), not by num_items_in_batch: teacher logprobs may not cover
-        # every student completion token (the loss is summed over the trimmed teacher window), so the global token
-        # count would be the wrong denominator. Gradient-accumulation normalization for the server path is left as a
-        # follow-up.
-        return self._compute_sparse_top_1_divergence_loss(
-            student_log_probs=student_log_probs,
-            teacher_top1_token_ids=topk_token_ids.squeeze(-1),
-            teacher_top1_logprobs=topk_teacher_lps.squeeze(-1),
-            reverse_token_ids=reverse_token_ids,
-            reverse_teacher_logprobs=actual_teacher_lps,
-            labels=labels,
-        )
-
-    def _compute_server_forward_kl_loss(
-        self,
-        teacher_result: dict[str, torch.Tensor],
-        student_log_probs: torch.Tensor,
-        labels: torch.Tensor,
-    ) -> torch.Tensor:
-        """Compute sparse forward KL from server-provided teacher top-k logprobs (beta==0 path).
-
-        Args:
-            teacher_result: dict with ``topk_logprobs`` (B, T, K) and ``topk_token_ids`` (B, T, K).
-            student_log_probs: (B, T, V) student log-softmax over vocabulary.
-            labels: (B, T) with -100 for positions to ignore.
-        """
-        teacher_topk_logprobs = teacher_result["topk_logprobs"]
-        teacher_topk_token_ids = teacher_result["topk_token_ids"]
-        valid = teacher_topk_logprobs > float("-inf")
-        neg_inf = torch.full((), float("-inf"), dtype=student_log_probs.dtype, device=student_log_probs.device)
-        student_topk_logprobs = student_log_probs.gather(dim=-1, index=teacher_topk_token_ids)
-        student_topk_logprobs = torch.where(valid, student_topk_logprobs, neg_inf)
-        teacher_topk_logprobs = torch.where(valid, teacher_topk_logprobs, neg_inf)
-
-        if self.loss_add_tail:
-            base_support_mask = valid
-            student_sparse_log_probs, support_mask = _add_tail_bucket(student_topk_logprobs, base_support_mask)
-            teacher_sparse_log_probs, _ = _add_tail_bucket(teacher_topk_logprobs, base_support_mask)
-        else:
-            support_mask = valid
-            student_sparse_log_probs = student_topk_logprobs - torch.logsumexp(
-                student_topk_logprobs, dim=-1, keepdim=True
-            )
-            teacher_sparse_log_probs = teacher_topk_logprobs - torch.logsumexp(
-                teacher_topk_logprobs, dim=-1, keepdim=True
-            )
-
-        jsd = _jsd_divergence(
-            student_sparse_log_probs,
-            teacher_sparse_log_probs,
-            beta=0.0,
-            support_mask=support_mask,
-        )
-        # 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")
+        if self.teacher_model is None:
+            raise ValueError("No teacher model configured.")
+        self.teacher_model.eval()
+        with torch.no_grad():
+            return self.teacher_model(
+                input_ids=inputs["input_ids"],
+                attention_mask=inputs["attention_mask"],
+            ).logits
 
     def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
         self._raise_if_local_teacher_tokenizer_mismatch()
@@ -1496,58 +1239,29 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
         labels = inputs["labels"][:, prompt_length:]
         completion_tokens = inputs["input_ids"][:, prompt_length:]
 
-        if self.use_teacher_server:
-            # Server path: token-level divergence using teacher logprobs.
-            # The server returns:
-            #   actual_logprobs  – (B, T)    teacher log p(x_actual)  (for reverse KL)
-            #   topk_logprobs    – (B, T, K) teacher top-k sorted logprobs (for forward KL)
-            #   topk_token_ids   – (B, T, K) corresponding token IDs
-            teacher_result = self._get_teacher_token_logprobs_from_server(inputs, prompt_length)
-
-            student_logits = student_outputs.logits[:, prompt_length - 1 : -1, :]
-            student_log_probs = F.log_softmax(student_logits / self.temperature, dim=-1)
-
-            comp_len = teacher_result["actual_logprobs"].shape[1]
-            completion_tokens = completion_tokens[:, :comp_len]
-            trimmed_labels = labels[:, :comp_len]
-
-            if self.beta > 0:
-                loss = self._compute_server_sparse_top_1_divergence_loss(
-                    teacher_result=teacher_result,
-                    student_log_probs=student_log_probs[:, :comp_len, :],
-                    completion_tokens=completion_tokens,
-                    labels=trimmed_labels,
-                )
-            else:
-                loss = self._compute_server_forward_kl_loss(
-                    teacher_result=teacher_result,
-                    student_log_probs=student_log_probs[:, :comp_len, :],
-                    labels=trimmed_labels,
-                )
+        # Local teacher: exact full-vocabulary loss except for the shared mixed top-1 path.
+        teacher_logits = self._get_teacher_logits(inputs)
+        student_logits = student_outputs.logits[:, prompt_length - 1 : -1, :]
+        teacher_logits = teacher_logits[:, prompt_length - 1 : -1, :]
+        if self.beta > 0 and self.loss_top_k == 1:
+            loss = self._compute_local_sparse_top_1_divergence_loss(
+                student_logits=student_logits,
+                teacher_logits=teacher_logits,
+                completion_tokens=completion_tokens,
+                labels=labels,
+                num_items_in_batch=num_items_in_batch,
+            )
         else:
-            # Local teacher: exact full-vocabulary loss except for the shared mixed top-1 path.
-            teacher_logits = self._get_teacher_logits(inputs)
-            student_logits = student_outputs.logits[:, prompt_length - 1 : -1, :]
-            teacher_logits = teacher_logits[:, prompt_length - 1 : -1, :]
-            if self.beta > 0 and self.loss_top_k == 1:
-                loss = self._compute_local_sparse_top_1_divergence_loss(
-                    student_logits=student_logits,
-                    teacher_logits=teacher_logits,
-                    completion_tokens=completion_tokens,
-                    labels=labels,
-                    num_items_in_batch=num_items_in_batch,
-                )
-            else:
-                loss = self.generalized_jsd_loss(
-                    student_logits=student_logits,
-                    teacher_logits=teacher_logits,
-                    labels=labels,
-                    beta=self.beta,
-                    temperature=self.temperature,
-                    top_k=self.loss_top_k,
-                    add_tail=self.loss_add_tail,
-                    num_items_in_batch=num_items_in_batch,
-                )
+            loss = self.generalized_jsd_loss(
+                student_logits=student_logits,
+                teacher_logits=teacher_logits,
+                labels=labels,
+                beta=self.beta,
+                temperature=self.temperature,
+                top_k=self.loss_top_k,
+                add_tail=self.loss_add_tail,
+                num_items_in_batch=num_items_in_batch,
+            )
 
         return (loss, student_outputs) if return_outputs else loss
 
diff --git a/trl/experimental/server_distillation/__init__.py b/trl/experimental/server_distillation/__init__.py
new file mode 100644
index 00000000000..59cb9347339
--- /dev/null
+++ b/trl/experimental/server_distillation/__init__.py
@@ -0,0 +1,19 @@
+# Copyright 2020-2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .server_distillation_config import ServerDistillationConfig
+from .server_distillation_trainer import ServerDistillationTrainer, build_teacher_request_inputs
+
+
+__all__ = ["ServerDistillationConfig", "ServerDistillationTrainer", "build_teacher_request_inputs"]
diff --git a/trl/experimental/server_distillation/server_distillation_config.py b/trl/experimental/server_distillation/server_distillation_config.py
new file mode 100644
index 00000000000..4298a38b2f1
--- /dev/null
+++ b/trl/experimental/server_distillation/server_distillation_config.py
@@ -0,0 +1,66 @@
+# Copyright 2020-2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from dataclasses import dataclass, field
+
+from ..distillation.distillation_config import DistillationConfig
+
+
+@dataclass
+class ServerDistillationConfig(DistillationConfig):
+    r"""
+    Configuration class for the [`ServerDistillationTrainer`].
+
+    Extends [`DistillationConfig`] with the address of an external vLLM server that scores the student's completions.
+    The teacher is never held locally: instead of a local forward pass, per-token teacher logprobs are fetched from
+    the server, so only the teacher's top-k logprobs are available and the loss is restricted to a sparse support.
+
+    Parameters:
+        teacher_model_server_url (`str` or `None`, *optional*):
+            Base URL of a vLLM server hosting the teacher model (e.g., `"http://localhost:8000"`). Required.
+    """
+
+    teacher_model_server_url: str | None = field(
+        default=None,
+        metadata={
+            "help": 'Base URL of a vLLM server hosting the teacher model (e.g., "http://localhost:8000"). Required.'
+        },
+    )
+
+    def __post_init__(self):
+        super().__post_init__()
+
+        if self.use_liger_kernel:
+            raise ValueError(
+                "use_liger_kernel=True is not supported by ServerDistillationTrainer because the Liger loss path "
+                "requires a local teacher model."
+            )
+        if self.teacher_model_server_url is None or not self.teacher_model_server_url.strip():
+            raise ValueError("teacher_model_server_url must be set for ServerDistillationTrainer.")
+        if self.beta == 0 and self.loss_top_k < 1:
+            raise ValueError(
+                f"loss_top_k must be positive with beta=0 (got loss_top_k={self.loss_top_k}). The pure forward "
+                f"server path only has access to the teacher's top-k logprobs, so it cannot compute the exact "
+                f"full-vocabulary loss when loss_top_k=0."
+            )
+        if self.reverse_kl_top_1_mode == "argmax":
+            raise ValueError(
+                "reverse_kl_top_1_mode='argmax' is not supported by ServerDistillationTrainer because the server "
+                "cannot provide teacher logprobs for arbitrary student-selected tokens."
+            )
+        if self.beta > 0 and self.loss_top_k != 1:
+            raise ValueError(
+                f"loss_top_k must be 1 with beta>0 (got loss_top_k={self.loss_top_k}). Mixed forward/reverse "
+                "distillation with an external teacher is only implemented for top-1 support."
+            )
diff --git a/trl/experimental/server_distillation/server_distillation_trainer.py b/trl/experimental/server_distillation/server_distillation_trainer.py
new file mode 100644
index 00000000000..acf6e021131
--- /dev/null
+++ b/trl/experimental/server_distillation/server_distillation_trainer.py
@@ -0,0 +1,353 @@
+# Copyright 2020-2026 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any
+
+import torch
+import torch.nn.functional as F
+
+from ..distillation.distillation_trainer import DistillationTrainer, _add_tail_bucket, _jsd_divergence
+from .server_distillation_config import ServerDistillationConfig
+
+
+def build_teacher_request_inputs(
+    input_ids: torch.Tensor,
+    attention_mask: torch.Tensor,
+    prompt_attention_mask: torch.Tensor | None = None,
+    labels: torch.Tensor | None = None,
+) -> tuple[list[list[int]], list[int], list[int]]:
+    """Trim padded batch tensors into per-sample sequences for teacher-server requests."""
+
+    if input_ids.shape != attention_mask.shape:
+        raise ValueError(
+            f"input_ids and attention_mask must have the same shape, got {input_ids.shape} and {attention_mask.shape}."
+        )
+
+    input_ids_cpu = input_ids.detach().cpu()
+    attention_mask_cpu = attention_mask.detach().cpu().bool()
+
+    if prompt_attention_mask is not None:
+        prompt_lengths = prompt_attention_mask.detach().cpu().sum(dim=1).to(torch.long)
+    else:
+        if labels is None:
+            raise ValueError("labels are required when prompt_attention_mask is not provided.")
+        if labels.shape != input_ids.shape:
+            raise ValueError(f"labels must match input_ids shape, got {labels.shape} and {input_ids.shape}.")
+        full_lengths = attention_mask_cpu.sum(dim=1).to(torch.long)
+        completion_lengths = (labels.detach().cpu() != -100).sum(dim=1).to(torch.long)
+        prompt_lengths = full_lengths - completion_lengths
+
+    trimmed_input_ids: list[list[int]] = []
+    prompt_lengths_list: list[int] = []
+    completion_lengths_list: list[int] = []
+
+    for row, mask, prompt_length in zip(input_ids_cpu, attention_mask_cpu, prompt_lengths, strict=True):
+        trimmed_row = row[mask]
+        prompt_len = int(prompt_length.item())
+        if prompt_len < 0 or prompt_len > trimmed_row.numel():
+            raise ValueError(
+                f"Invalid prompt length {prompt_len} for trimmed sequence of length {trimmed_row.numel()}."
+            )
+        trimmed_input_ids.append(trimmed_row.tolist())
+        prompt_lengths_list.append(prompt_len)
+        completion_lengths_list.append(int(trimmed_row.numel()) - prompt_len)
+
+    return trimmed_input_ids, prompt_lengths_list, completion_lengths_list
+
+
+class ServerDistillationTrainer(DistillationTrainer):
+    """Distillation from a teacher hosted on an external vLLM server.
+
+    Instead of running a local teacher forward pass, per-token teacher logprobs are fetched from a vLLM server via
+    [`~generation.vllm_client.VLLMClient`]. The server only returns the teacher's top-k logprobs, so the divergence is
+    restricted to a sparse support (top-1 for `beta > 0`, top-k for the pure forward path `beta = 0`). Everything else
+    — the student forward, generation, buffering, metrics — is inherited from [`experimental.distillation.DistillationTrainer`].
+    """
+
+    _tag_names = ["trl", "server-distillation"]
+    _name = "Server Distillation"
+
+    def __init__(
+        self,
+        model=None,
+        args: ServerDistillationConfig | None = None,
+        data_collator=None,
+        train_dataset=None,
+        eval_dataset=None,
+        processing_class=None,
+        compute_metrics=None,
+        callbacks=None,
+        optimizers=(None, None),
+        preprocess_logits_for_metrics=None,
+        peft_config=None,
+    ):
+        if args is None:
+            args = ServerDistillationConfig(output_dir="tmp_server_distillation")
+
+        # No local teacher: the base sets `self.teacher_model = None`, then we attach the server client below.
+        super().__init__(
+            model=model,
+            teacher_model=None,
+            args=args,
+            data_collator=data_collator,
+            train_dataset=train_dataset,
+            eval_dataset=eval_dataset,
+            processing_class=processing_class,
+            compute_metrics=compute_metrics,
+            callbacks=callbacks,
+            optimizers=optimizers,
+            preprocess_logits_for_metrics=preprocess_logits_for_metrics,
+            peft_config=peft_config,
+        )
+
+        from ...generation.vllm_client import VLLMClient
+
+        self.teacher_client = VLLMClient(base_url=args.teacher_model_server_url, connection_timeout=60.0)
+
+    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
+        # Student forward pass
+        student_outputs = model(
+            input_ids=inputs["input_ids"],
+            attention_mask=inputs["attention_mask"],
+        )
+        prompt_length = self._compute_prompt_length(inputs)
+        labels = inputs["labels"][:, prompt_length:]
+        completion_tokens = inputs["input_ids"][:, prompt_length:]
+
+        # Server path: token-level divergence using teacher logprobs.
+        # The server returns:
+        #   actual_logprobs  – (B, T)    teacher log p(x_actual)  (for reverse KL)
+        #   topk_logprobs    – (B, T, K) teacher top-k sorted logprobs (for forward KL)
+        #   topk_token_ids   – (B, T, K) corresponding token IDs
+        teacher_result = self._get_teacher_token_logprobs_from_server(inputs, prompt_length)
+
+        student_logits = student_outputs.logits[:, prompt_length - 1 : -1, :]
+        student_log_probs = F.log_softmax(student_logits / self.temperature, dim=-1)
+
+        comp_len = teacher_result["actual_logprobs"].shape[1]
+        completion_tokens = completion_tokens[:, :comp_len]
+        trimmed_labels = labels[:, :comp_len]
+
+        if self.beta > 0:
+            loss = self._compute_server_sparse_top_1_divergence_loss(
+                teacher_result=teacher_result,
+                student_log_probs=student_log_probs[:, :comp_len, :],
+                completion_tokens=completion_tokens,
+                labels=trimmed_labels,
+            )
+        else:
+            loss = self._compute_server_forward_kl_loss(
+                teacher_result=teacher_result,
+                student_log_probs=student_log_probs[:, :comp_len, :],
+                labels=trimmed_labels,
+            )
+
+        return (loss, student_outputs) if return_outputs else loss
+
+    def _get_teacher_token_logprobs_from_server(
+        self,
+        inputs: dict[str, torch.Tensor | Any],
+        aligned_prompt_length: int,
+    ) -> dict[str, torch.Tensor]:
+        """Fetch per-token teacher logprobs from an external vLLM server.
+
+        Returns a dict with:
+            ``actual_logprobs`` - (batch, completion_length) teacher log-prob for the actual
+                                   token at each position (for reverse KL).
+            ``topk_logprobs`` - (batch, completion_length, K) teacher top-k sorted logprobs
+                                   (for forward KL).
+            ``topk_token_ids`` - (batch, completion_length, K) corresponding token IDs.
+        """
+        import numpy as np
+
+        input_ids = inputs["input_ids"]
+        batch_size = input_ids.shape[0]
+        sequences, prompt_lengths, completion_lengths = build_teacher_request_inputs(
+            input_ids,
+            inputs["attention_mask"],
+            prompt_attention_mask=inputs.get("prompt_attention_mask"),
+            labels=inputs.get("labels"),
+        )
+
+        # The pure forward server path can use the requested teacher top-k support.
+        # When beta > 0, config validation restricts the server-backed path to top-1.
+        requested_top_k = self.loss_top_k
+        result = self.teacher_client.get_sequence_logprobs(
+            sequences=sequences,
+            prompt_lengths=prompt_lengths,
+            top_logprobs=requested_top_k,
+            temperature=self.temperature,
+        )
+        K = requested_top_k
+
+        device = input_ids.device
+        labels = inputs.get("labels")
+        if labels is None:
+            raise ValueError("labels are required to align teacher-server logprobs with the student loss tensors.")
+
+        # The student loss slices tensors in padded-sequence coordinates starting at `aligned_prompt_length`.
+        # Place each teacher completion into that same coordinate system by locating the first non-masked completion
+        # token in `labels`. This works for both left-padded off-policy batches and on-policy batches where
+        # completions are right-padded after a fixed-width prompt block.
+        completion_offsets = []
+        label_mask = labels != -100
+        for sample_mask, comp_len in zip(label_mask, completion_lengths, strict=True):
+            if comp_len == 0:
+                completion_offsets.append(0)
+                continue
+            completion_start = int(torch.nonzero(sample_mask, as_tuple=False)[0].item())
+            completion_offsets.append(completion_start - aligned_prompt_length)
+
+        # Size the output tensors to tightly fit the teacher logprobs. Using the full padded
+        # sequence length would include padding positions with -inf teacher logprobs, producing
+        # +inf in the forward pass and NaN gradients in the backward pass (0 * inf = NaN).
+        # Shorter samples in variable-length batches still need the -inf sentinel at the tail;
+        # downstream loss consumers (_compute_server_sparse_top_1_divergence_loss,
+        # _compute_server_forward_kl_loss) neutralise those positions before the divergence
+        # math runs.
+        completion_length = max(
+            (offset + len(lps) for offset, lps in zip(completion_offsets, result["logprobs"], strict=True)),
+            default=0,
+        )
+
+        # actual_logprobs: (B, T) - teacher logprob for the actual token
+        def _actual_to_tensor(key):
+            arr = np.full((batch_size, completion_length), float("-inf"), dtype=np.float32)
+            for i, (offset, seq_lps) in enumerate(zip(completion_offsets, result[key], strict=True)):
+                if seq_lps:
+                    vals = np.array(seq_lps, dtype=np.float32)  # (comp_len_i, 1)
+                    arr[i, offset : offset + vals.shape[0]] = vals[:, 0]
+            return torch.from_numpy(arr).to(device)
+
+        # topk: (B, T, K)
+        def _topk_to_tensor(key, k, np_dtype, fill):
+            arr = np.full((batch_size, completion_length, k), fill, dtype=np_dtype)
+            for i, (offset, seq_vals) in enumerate(zip(completion_offsets, result[key], strict=True)):
+                if seq_vals:
+                    vals = np.array(seq_vals, dtype=np_dtype)  # (comp_len_i, k)
+                    arr[i, offset : offset + vals.shape[0], :] = vals
+            return torch.from_numpy(arr).to(device)
+
+        return {
+            "actual_logprobs": _actual_to_tensor("actual_logprobs"),
+            "topk_logprobs": _topk_to_tensor("logprobs", K, np.float32, float("-inf")),
+            "topk_token_ids": _topk_to_tensor("logprob_token_ids", K, np.int64, 0),
+        }
+
+    def _compute_server_sparse_top_1_divergence_loss(
+        self,
+        teacher_result: dict[str, torch.Tensor],
+        student_log_probs: torch.Tensor,
+        completion_tokens: torch.Tensor,
+        labels: torch.Tensor,
+    ) -> torch.Tensor:
+        """Compute exact sparse top-1 generalized JSD/KL from server-provided teacher logprobs.
+
+        Args:
+            teacher_result: dict with ``actual_logprobs`` (B, T), ``topk_logprobs`` (B, T, K),
+                ``topk_token_ids`` (B, T, K).
+            student_log_probs: (B, T, V) student log-softmax over vocabulary.
+            completion_tokens: (B, T) actual token IDs in the completion.
+            labels: (B, T) with -100 for positions to ignore.
+        """
+        topk_teacher_lps = teacher_result["topk_logprobs"]  # (B, T, 1)
+        topk_token_ids = teacher_result["topk_token_ids"]  # (B, T, 1)
+        actual_teacher_lps = teacher_result["actual_logprobs"]  # (B, T)
+        required = labels != -100
+
+        missing_actual = required & ~torch.isfinite(actual_teacher_lps)
+        if missing_actual.any():
+            missing_count = int(missing_actual.sum().item())
+            total_required = int(required.sum().item())
+            raise ValueError(
+                "Teacher server is missing actual-token logprobs for required reverse-KL positions: "
+                f"{missing_count}/{total_required}."
+            )
+        if self.beta < 1:
+            teacher_top1_logprobs = topk_teacher_lps.squeeze(-1)
+            missing_top1 = required & ~torch.isfinite(teacher_top1_logprobs)
+            if missing_top1.any():
+                missing_count = int(missing_top1.sum().item())
+                total_required = int(required.sum().item())
+                raise ValueError(
+                    "Teacher server is missing top-1 logprobs for required forward-KL positions: "
+                    f"{missing_count}/{total_required}."
+                )
+
+        # Replace -inf teacher logprobs at intra-batch padding (labels == -100) with 0 so
+        # reverse-KL's student_probs·(log_s - log_t) does not leak +inf into the backward pass.
+        pad_mask_2d = ~required
+        pad_mask_3d = pad_mask_2d.unsqueeze(-1)
+        topk_teacher_lps = torch.where(pad_mask_3d, 0.0, topk_teacher_lps)
+        actual_teacher_lps = torch.where(pad_mask_2d, 0.0, actual_teacher_lps)
+
+        # Server path only supports "sampled" mode — config validation enforces this, but we guard
+        # explicitly so future relaxations of the config check don't silently change behaviour.
+        reverse_token_ids = self._get_reverse_kl_top_1_tokens(student_log_probs, completion_tokens)
+        # The server path normalizes locally (batchmean), not by num_items_in_batch: teacher logprobs may not cover
+        # every student completion token (the loss is summed over the trimmed teacher window), so the global token
+        # count would be the wrong denominator. Gradient-accumulation normalization for the server path is left as a
+        # follow-up.
+        return self._compute_sparse_top_1_divergence_loss(
+            student_log_probs=student_log_probs,
+            teacher_top1_token_ids=topk_token_ids.squeeze(-1),
+            teacher_top1_logprobs=topk_teacher_lps.squeeze(-1),
+            reverse_token_ids=reverse_token_ids,
+            reverse_teacher_logprobs=actual_teacher_lps,
+            labels=labels,
+        )
+
+    def _compute_server_forward_kl_loss(
+        self,
+        teacher_result: dict[str, torch.Tensor],
+        student_log_probs: torch.Tensor,
+        labels: torch.Tensor,
+    ) -> torch.Tensor:
+        """Compute sparse forward KL from server-provided teacher top-k logprobs (beta==0 path).
+
+        Args:
+            teacher_result: dict with ``topk_logprobs`` (B, T, K) and ``topk_token_ids`` (B, T, K).
+            student_log_probs: (B, T, V) student log-softmax over vocabulary.
+            labels: (B, T) with -100 for positions to ignore.
+        """
+        teacher_topk_logprobs = teacher_result["topk_logprobs"]
+        teacher_topk_token_ids = teacher_result["topk_token_ids"]
+        valid = teacher_topk_logprobs > float("-inf")
+        neg_inf = torch.full((), float("-inf"), dtype=student_log_probs.dtype, device=student_log_probs.device)
+        student_topk_logprobs = student_log_probs.gather(dim=-1, index=teacher_topk_token_ids)
+        student_topk_logprobs = torch.where(valid, student_topk_logprobs, neg_inf)
+        teacher_topk_logprobs = torch.where(valid, teacher_topk_logprobs, neg_inf)
+
+        if self.loss_add_tail:
+            base_support_mask = valid
+            student_sparse_log_probs, support_mask = _add_tail_bucket(student_topk_logprobs, base_support_mask)
+            teacher_sparse_log_probs, _ = _add_tail_bucket(teacher_topk_logprobs, base_support_mask)
+        else:
+            support_mask = valid
+            student_sparse_log_probs = student_topk_logprobs - torch.logsumexp(
+                student_topk_logprobs, dim=-1, keepdim=True
+            )
+            teacher_sparse_log_probs = teacher_topk_logprobs - torch.logsumexp(
+                teacher_topk_logprobs, dim=-1, keepdim=True
+            )
+
+        jsd = _jsd_divergence(
+            student_sparse_log_probs,
+            teacher_sparse_log_probs,
+            beta=0.0,
+            support_mask=support_mask,
+        )
+        # 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")
diff --git a/trl/trainer/base_trainer.py b/trl/trainer/base_trainer.py
index 19a484b9241..f844c2bb36f 100644
--- a/trl/trainer/base_trainer.py
+++ b/trl/trainer/base_trainer.py
@@ -58,6 +58,7 @@
     "PRMTrainer",
     "SDFTTrainer",
     "SDPOTrainer",
+    "ServerDistillationTrainer",
     "SSDTrainer",
     "TPOTrainer",
     "XPOTrainer",

From f75cd7ee6c19e5cd63b481dccee1387e6fb58b11 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Sun, 19 Jul 2026 20:24:44 +0000
Subject: [PATCH 06/44] Remove the local sparse top-1 loss path from the base
 trainer

With the server path extracted, the base trainer no longer needs the mixed
top-1 support path: the local teacher always has full logits, so it can compute
the exact generalized JSD/KL directly.

- `compute_loss` local branch now always calls `generalized_jsd_loss` (which
  still honours `loss_top_k` for the top-k approximation); the
  `beta > 0 and loss_top_k == 1` special case and its `completion_tokens` slice
  are gone.
- `_compute_local_sparse_top_1_divergence_loss` is deleted.
- `_compute_sparse_top_1_divergence_loss` and `_get_reverse_kl_top_1_tokens`
  become server-only, so they move to `ServerDistillationTrainer`. The subclass
  now defines them directly rather than inheriting them.
- `reverse_kl_top_1_mode` only drove those helpers, so it moves from
  `DistillationConfig` to `ServerDistillationConfig` (field + the
  "must be one of" validation), and the `self.reverse_kl_top_1_mode` assignment
  moves from the base `__init__` to the subclass `__init__`. Its stale mentions
  in the `loss_top_k` docs are trimmed.
- The corresponding config test moves to the server test file.

Verified on CPU that the server sparse loss is numerically unchanged by the
relocation.
---
 .../experimental/test_distillation_trainer.py |   9 -
 .../test_server_distillation_trainer.py       | 163 +-----------------
 .../distillation/distillation_config.py       |  30 +---
 .../distillation/distillation_trainer.py      | 123 ++-----------
 .../server_distillation_config.py             |  15 ++
 .../server_distillation_trainer.py            |  63 +++++++
 6 files changed, 97 insertions(+), 306 deletions(-)

diff --git a/tests/experimental/test_distillation_trainer.py b/tests/experimental/test_distillation_trainer.py
index f9c740b4279..674bb70ebff 100644
--- a/tests/experimental/test_distillation_trainer.py
+++ b/tests/experimental/test_distillation_trainer.py
@@ -27,15 +27,6 @@
 from ..testing_utils import TrlTestCase, require_liger_kernel, require_torch_accelerator
 
 
-def _make_distillation_config_kwargs(tmp_path):
-    return {"output_dir": str(tmp_path), "report_to": "none", "use_cpu": True, "bf16": False}
-
-
-def test_distillation_config_rejects_invalid_reverse_kl_top_1_mode(tmp_path):
-    with pytest.raises(ValueError, match="reverse_kl_top_1_mode must be one of"):
-        DistillationConfig(**_make_distillation_config_kwargs(tmp_path), reverse_kl_top_1_mode="invalid")
-
-
 def _reference_generalized_jsd(student_logits, teacher_logits, labels=None, beta=0.5, temperature=1.0):
     """Naive reference for the generalized JSD, written straight from the definition.
 
diff --git a/tests/experimental/test_server_distillation_trainer.py b/tests/experimental/test_server_distillation_trainer.py
index 4334865d9d0..b85f7c689fd 100644
--- a/tests/experimental/test_server_distillation_trainer.py
+++ b/tests/experimental/test_server_distillation_trainer.py
@@ -18,10 +18,9 @@
 import pytest
 import torch
 import torch.nn.functional as F
-from datasets import Dataset, load_dataset
+from datasets import Dataset
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-from trl.experimental.distillation import DistillationConfig, DistillationTrainer
 from trl.experimental.distillation.distillation_trainer import _add_tail_bucket, _jsd_divergence
 from trl.experimental.server_distillation import (
     ServerDistillationConfig,
@@ -42,64 +41,6 @@ def _make_server_config_kwargs(tmp_path):
     }
 
 
-def _build_server_result(teacher_logits, inputs, temperature=1.0):
-    """Simulate a vLLM server response with variable-length per-sample completions."""
-    _, _, completion_lengths = build_teacher_request_inputs(
-        inputs["input_ids"],
-        inputs["attention_mask"],
-        prompt_attention_mask=inputs.get("prompt_attention_mask"),
-        labels=inputs.get("labels"),
-    )
-
-    label_mask = inputs["labels"] != -100
-    actual_logprobs = []
-    logprobs = []
-    logprob_token_ids = []
-
-    for i, comp_len in enumerate(completion_lengths):
-        if comp_len == 0:
-            actual_logprobs.append([])
-            logprobs.append([])
-            logprob_token_ids.append([])
-            continue
-
-        comp_start = int(torch.nonzero(label_mask[i], as_tuple=False)[0].item())
-        sample_logits = teacher_logits[i, comp_start - 1 : comp_start - 1 + comp_len, :]
-        sample_log_probs = F.log_softmax(sample_logits / temperature, dim=-1)
-        comp_tokens = inputs["input_ids"][i, comp_start : comp_start + comp_len]
-
-        top1_ids = sample_logits.argmax(dim=-1, keepdim=True)
-        top1_lps = sample_log_probs.gather(dim=-1, index=top1_ids)
-        actual_lps = sample_log_probs.gather(dim=-1, index=comp_tokens.unsqueeze(-1))
-
-        actual_logprobs.append(actual_lps.tolist())
-        logprobs.append(top1_lps.tolist())
-        logprob_token_ids.append(top1_ids.tolist())
-
-    return {
-        "actual_logprobs": actual_logprobs,
-        "logprobs": logprobs,
-        "logprob_token_ids": logprob_token_ids,
-    }
-
-
-class RecordingTeacherClient:
-    def __init__(self):
-        self.calls = []
-        self.result = None
-
-    def get_sequence_logprobs(self, sequences, prompt_lengths, top_logprobs, temperature):
-        self.calls.append(
-            {
-                "sequences": sequences,
-                "prompt_lengths": prompt_lengths,
-                "top_logprobs": top_logprobs,
-                "temperature": temperature,
-            }
-        )
-        return self.result
-
-
 def _ragged_server_response():
     # Two samples with completion lengths 1 and 3 respectively; matches the wire format
     # of VLLMClient.get_sequence_logprobs (per-sample shape (comp_len, top_k=1)).
@@ -151,6 +92,11 @@ def test_config_rejects_reverse_kl_argmax(tmp_path):
         ServerDistillationConfig(**_make_server_config_kwargs(tmp_path), reverse_kl_top_1_mode="argmax")
 
 
+def test_config_rejects_invalid_reverse_kl_top_1_mode(tmp_path):
+    with pytest.raises(ValueError, match="reverse_kl_top_1_mode must be one of"):
+        ServerDistillationConfig(**_make_server_config_kwargs(tmp_path), reverse_kl_top_1_mode="invalid")
+
+
 def test_config_rejects_mixed_loss_without_top_1(tmp_path):
     with pytest.raises(ValueError, match="loss_top_k must be 1 with beta>0"):
         ServerDistillationConfig(**_make_server_config_kwargs(tmp_path), beta=0.5, loss_top_k=2)
@@ -304,103 +250,6 @@ def test_mask_keeps_forward_and_backward_finite(self):
         assert torch.isfinite(raw_student.grad).all()
 
 
-class TestServerVsLocalTeacher(TrlTestCase):
-    def setup_method(self):
-        self.model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
-        self.tokenizer = AutoTokenizer.from_pretrained(self.model_id)
-        self.tokenizer.pad_token = self.tokenizer.eos_token
-
-    def _local_args(self, **kwargs):
-        args = {
-            "output_dir": self.tmp_dir,
-            "per_device_train_batch_size": 2,
-            "gradient_accumulation_steps": 1,
-            "max_steps": 1,
-            "save_strategy": "no",
-            "report_to": "none",
-            "use_cpu": True,
-            "bf16": False,
-            "lmbda": 0.0,
-            "max_length": 128,
-            "max_completion_length": 32,
-            "model_init_kwargs": {"dtype": "float32", "device_map": None},
-            "teacher_model_init_kwargs": {"dtype": "float32", "device_map": None},
-        }
-        args.update(kwargs)
-        return DistillationConfig(**args)
-
-    def _make_local_trainer(self, **kwargs):
-        dataset = load_dataset("trl-internal-testing/zen", "conversational_language_modeling", split="train")
-        return DistillationTrainer(
-            model=self.model_id,
-            teacher_model=self.model_id,
-            args=self._local_args(**kwargs),
-            train_dataset=dataset,
-            processing_class=self.tokenizer,
-        )
-
-    def _make_server_trainer(self, **kwargs):
-        dataset = load_dataset("trl-internal-testing/zen", "conversational_language_modeling", split="train")
-        args = self._local_args(**kwargs)
-        server_args = ServerDistillationConfig(
-            **{k: v for k, v in vars(args).items() if not k.startswith("_")},
-            teacher_model_server_url="http://localhost:8000",
-        )
-        return ServerDistillationTrainer(
-            model=self.model_id,
-            args=server_args,
-            train_dataset=dataset,
-            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_sampled_mode_matches_between_local_and_external_teachers(self, monkeypatch):
-        import trl.generation.vllm_client as vllm_client_module
-
-        teacher_client = RecordingTeacherClient()
-        monkeypatch.setattr(vllm_client_module, "VLLMClient", lambda *args, **kwargs: teacher_client)
-
-        local_trainer = self._make_local_trainer(beta=0.5, loss_top_k=1, reverse_kl_top_1_mode="sampled")
-        server_trainer = self._make_server_trainer(beta=0.5, loss_top_k=1)
-
-        cpu_inputs = self._make_batch(local_trainer)
-        expected_sequences, expected_prompt_lengths, _ = build_teacher_request_inputs(
-            cpu_inputs["input_ids"],
-            cpu_inputs["attention_mask"],
-            prompt_attention_mask=cpu_inputs["prompt_attention_mask"],
-            labels=cpu_inputs["labels"],
-        )
-
-        local_inputs = self._move_batch_to_device(cpu_inputs, local_trainer.accelerator.device)
-        server_inputs = self._move_batch_to_device(cpu_inputs, server_trainer.accelerator.device)
-
-        local_trainer.teacher_model.eval()
-        with torch.no_grad():
-            teacher_logits = local_trainer.teacher_model(
-                input_ids=local_inputs["input_ids"],
-                attention_mask=local_inputs["attention_mask"],
-            ).logits
-            teacher_client.result = _build_server_result(
-                teacher_logits,
-                local_inputs,
-                temperature=local_trainer.temperature,
-            )
-            local_loss = local_trainer.compute_loss(local_trainer.model, local_inputs)
-            server_loss = server_trainer.compute_loss(server_trainer.model, server_inputs)
-
-        assert teacher_client.calls[0]["sequences"] == expected_sequences
-        assert teacher_client.calls[0]["prompt_lengths"] == expected_prompt_lengths
-        assert teacher_client.calls[0]["top_logprobs"] == 1
-        torch.testing.assert_close(local_loss, server_loss)
-
-
 class TestServerDistillationTrainerRaggedGrad(TrlTestCase):
     @classmethod
     def setup_class(cls):
diff --git a/trl/experimental/distillation/distillation_config.py b/trl/experimental/distillation/distillation_config.py
index 5b7efc9f7c8..b9bd0cdb03a 100644
--- a/trl/experimental/distillation/distillation_config.py
+++ b/trl/experimental/distillation/distillation_config.py
@@ -57,11 +57,6 @@ class DistillationConfig(_BaseConfig):
             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
             JSD.
-        reverse_kl_top_1_mode (`str`, *optional*, defaults to `"sampled"`):
-            Selection rule for the reverse-KL top-1 token when `beta > 0` and `loss_top_k == 1`. `"sampled"` uses the
-            actual completion token in the batch. `"argmax"` uses the student's highest-probability token. This
-            setting does not affect the forward-KL support, which always uses the teacher's top-1 token. Ignored when
-            `beta == 0` or `loss_top_k != 1`.
         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*):
@@ -83,8 +78,7 @@ class DistillationConfig(_BaseConfig):
             Number of top tokens to use when computing the JSD/KL loss. Both student and teacher distributions are
             restricted to these K tokens and re-normalized before computing divergence. If 0, the full vocabulary
             is used. For local teachers, the general support rule is teacher top-k for forward KL, student top-k for
-            reverse KL, and the union for mixed JSD. When `beta > 0` and `loss_top_k == 1`, the forward support still
-            uses the teacher's top-1 token, while the reverse top-1 token is controlled by `reverse_kl_top_1_mode`.
+            reverse KL, and the union for mixed JSD.
         loss_add_tail (`bool`, *optional*, defaults to `True`):
             Whether to append a tail bucket that represents the remaining probability mass outside the selected top-k
             support when computing the loss.
@@ -194,14 +188,6 @@ class DistillationConfig(_BaseConfig):
             "0.0 = forward KL, 0.5 = JSD, 1.0 = reverse KL."
         },
     )
-    reverse_kl_top_1_mode: str = field(
-        default="sampled",
-        metadata={
-            "help": "Reverse-KL top-1 token selection when beta > 0 and loss_top_k == 1. "
-            "Use 'sampled' for the actual completion token or 'argmax' for the student's top-1 token. "
-            "The forward-KL support always uses the teacher's top-1 token. Ignored when beta == 0 or loss_top_k != 1."
-        },
-    )
     max_completion_length: int = field(
         default=512,
         metadata={"help": "Maximum number of tokens to generate per completion."},
@@ -241,9 +227,7 @@ class DistillationConfig(_BaseConfig):
             "Both student and teacher distributions are restricted to these K tokens "
             "(selected based on beta: teacher's top-k for forward KL, student's top-k for reverse KL, "
             "union of both for JSD) and re-normalized before computing divergence. "
-            "If 0, the full vocabulary is used (slower but exact). "
-            "When beta > 0 and loss_top_k == 1, the forward support still uses the teacher's top-1 token, "
-            "while the reverse top-1 token is controlled by reverse_kl_top_1_mode."
+            "If 0, the full vocabulary is used (slower but exact)."
         },
     )
     loss_add_tail: bool = field(
@@ -359,8 +343,6 @@ def __post_init__(self):
             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}.")
-        if self.reverse_kl_top_1_mode not in {"sampled", "argmax"}:
-            raise ValueError("reverse_kl_top_1_mode must be one of: 'sampled', 'argmax'")
 
         if self.max_length is not None and self.max_completion_length >= self.max_length:
             raise ValueError(
@@ -386,14 +368,6 @@ def __post_init__(self):
                 f"{self.per_device_train_batch_size} * {self.gradient_accumulation_steps}."
             )
 
-        if self.reverse_kl_top_1_mode != "sampled" and (self.beta == 0 or self.loss_top_k != 1):
-            warnings.warn(
-                f"reverse_kl_top_1_mode='{self.reverse_kl_top_1_mode}' has no effect when beta={self.beta} "
-                f"and loss_top_k={self.loss_top_k}. It only applies when beta > 0 and loss_top_k == 1.",
-                UserWarning,
-                stacklevel=2,
-            )
-
         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 "
diff --git a/trl/experimental/distillation/distillation_trainer.py b/trl/experimental/distillation/distillation_trainer.py
index dd50fc96243..7b96114df00 100644
--- a/trl/experimental/distillation/distillation_trainer.py
+++ b/trl/experimental/distillation/distillation_trainer.py
@@ -540,7 +540,6 @@ def __init__(
         self.temperature = args.temperature
         self.top_p = args.top_p
         self.num_generations = args.num_generations
-        self.reverse_kl_top_1_mode = args.reverse_kl_top_1_mode
         self.loss_top_k = args.loss_top_k
         self.loss_add_tail = args.loss_add_tail
 
@@ -1122,96 +1121,6 @@ def generalized_jsd_loss(
             jsd, labels=labels, reduction=reduction, num_items_in_batch=num_items_in_batch
         )
 
-    def _get_reverse_kl_top_1_tokens(
-        self, student_scores: torch.Tensor, completion_tokens: torch.Tensor
-    ) -> torch.Tensor:
-        """Return the reverse-KL top-1 token IDs for the mixed top-1 loss path.
-
-        Args:
-            student_scores: Any (B, T, V) tensor whose argmax selects the student's top token
-                (logits or log-probs — both are order-preserving).
-            completion_tokens: (B, T) actual token IDs in the completion.
-        """
-        if self.reverse_kl_top_1_mode == "argmax":
-            return student_scores.argmax(dim=-1)
-        return completion_tokens
-
-    def _compute_sparse_top_1_divergence_loss(
-        self,
-        student_log_probs: torch.Tensor,
-        teacher_top1_token_ids: torch.Tensor,
-        teacher_top1_logprobs: torch.Tensor,
-        reverse_token_ids: torch.Tensor,
-        reverse_teacher_logprobs: torch.Tensor,
-        labels: torch.Tensor,
-        num_items_in_batch=None,
-    ) -> torch.Tensor:
-        """Compute exact generalized JSD/KL on top-1 support for the mixed beta>0 path."""
-        neg_inf = torch.full((), float("-inf"), dtype=student_log_probs.dtype, device=student_log_probs.device)
-
-        if self.beta == 1:
-            support = reverse_token_ids.unsqueeze(-1)
-            support_mask = torch.ones_like(support, dtype=torch.bool)
-            teacher_support_logprobs = reverse_teacher_logprobs.unsqueeze(-1)
-        else:
-            teacher_support = teacher_top1_token_ids.unsqueeze(-1)
-            reverse_support = reverse_token_ids.unsqueeze(-1)
-            support = torch.cat([teacher_support, reverse_support], dim=-1)
-            support_mask = torch.ones_like(support, dtype=torch.bool)
-            support_mask[..., 1] = support[..., 1] != support[..., 0]
-            teacher_support_logprobs = torch.stack([teacher_top1_logprobs, reverse_teacher_logprobs], dim=-1)
-            support = torch.where(support_mask, support, torch.zeros_like(support))
-
-        student_support_logprobs = student_log_probs.gather(-1, support)
-        student_support_logprobs = torch.where(support_mask, student_support_logprobs, neg_inf)
-        teacher_support_logprobs = torch.where(support_mask, teacher_support_logprobs, neg_inf)
-
-        if self.loss_add_tail:
-            base_support_mask = support_mask
-            student_sparse_log_probs, support_mask = _add_tail_bucket(student_support_logprobs, base_support_mask)
-            teacher_sparse_log_probs, _ = _add_tail_bucket(teacher_support_logprobs, base_support_mask)
-        else:
-            student_sparse_log_probs = student_support_logprobs - torch.logsumexp(
-                student_support_logprobs, dim=-1, keepdim=True
-            )
-            teacher_sparse_log_probs = teacher_support_logprobs - torch.logsumexp(
-                teacher_support_logprobs, dim=-1, keepdim=True
-            )
-
-        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
-        )
-
-    def _compute_local_sparse_top_1_divergence_loss(
-        self,
-        student_logits: torch.Tensor,
-        teacher_logits: torch.Tensor,
-        completion_tokens: torch.Tensor,
-        labels: torch.Tensor,
-        num_items_in_batch=None,
-    ) -> torch.Tensor:
-        """Compute the mixed top-1 loss for a local teacher using gathered full-logit probabilities."""
-        student_log_probs = F.log_softmax(student_logits / self.temperature, dim=-1)
-        teacher_log_probs = F.log_softmax(teacher_logits / self.temperature, dim=-1)
-
-        teacher_top1_token_ids = teacher_logits.argmax(dim=-1)
-        teacher_top1_logprobs = teacher_log_probs.gather(dim=-1, index=teacher_top1_token_ids.unsqueeze(-1)).squeeze(
-            -1
-        )
-        reverse_token_ids = self._get_reverse_kl_top_1_tokens(student_logits, completion_tokens)
-        reverse_teacher_logprobs = teacher_log_probs.gather(dim=-1, index=reverse_token_ids.unsqueeze(-1)).squeeze(-1)
-
-        return self._compute_sparse_top_1_divergence_loss(
-            student_log_probs=student_log_probs,
-            teacher_top1_token_ids=teacher_top1_token_ids,
-            teacher_top1_logprobs=teacher_top1_logprobs,
-            reverse_token_ids=reverse_token_ids,
-            reverse_teacher_logprobs=reverse_teacher_logprobs,
-            labels=labels,
-            num_items_in_batch=num_items_in_batch,
-        )
-
     def _get_teacher_logits(self, inputs: dict[str, torch.Tensor | Any]) -> torch.Tensor:
         """Get teacher logits — dispatches between local model and external server."""
         if self.teacher_model is None:
@@ -1237,31 +1146,21 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
         )
         prompt_length = self._compute_prompt_length(inputs)
         labels = inputs["labels"][:, prompt_length:]
-        completion_tokens = inputs["input_ids"][:, prompt_length:]
 
-        # Local teacher: exact full-vocabulary loss except for the shared mixed top-1 path.
+        # 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, :]
-        if self.beta > 0 and self.loss_top_k == 1:
-            loss = self._compute_local_sparse_top_1_divergence_loss(
-                student_logits=student_logits,
-                teacher_logits=teacher_logits,
-                completion_tokens=completion_tokens,
-                labels=labels,
-                num_items_in_batch=num_items_in_batch,
-            )
-        else:
-            loss = self.generalized_jsd_loss(
-                student_logits=student_logits,
-                teacher_logits=teacher_logits,
-                labels=labels,
-                beta=self.beta,
-                temperature=self.temperature,
-                top_k=self.loss_top_k,
-                add_tail=self.loss_add_tail,
-                num_items_in_batch=num_items_in_batch,
-            )
+        loss = self.generalized_jsd_loss(
+            student_logits=student_logits,
+            teacher_logits=teacher_logits,
+            labels=labels,
+            beta=self.beta,
+            temperature=self.temperature,
+            top_k=self.loss_top_k,
+            add_tail=self.loss_add_tail,
+            num_items_in_batch=num_items_in_batch,
+        )
 
         return (loss, student_outputs) if return_outputs else loss
 
diff --git a/trl/experimental/server_distillation/server_distillation_config.py b/trl/experimental/server_distillation/server_distillation_config.py
index 4298a38b2f1..dfea92faf65 100644
--- a/trl/experimental/server_distillation/server_distillation_config.py
+++ b/trl/experimental/server_distillation/server_distillation_config.py
@@ -29,6 +29,11 @@ class ServerDistillationConfig(DistillationConfig):
     Parameters:
         teacher_model_server_url (`str` or `None`, *optional*):
             Base URL of a vLLM server hosting the teacher model (e.g., `"http://localhost:8000"`). Required.
+        reverse_kl_top_1_mode (`str`, *optional*, defaults to `"sampled"`):
+            Selection rule for the reverse-KL top-1 token when `beta > 0` and `loss_top_k == 1`. `"sampled"` uses the
+            actual completion token in the batch. `"argmax"` uses the student's highest-probability token (not
+            supported by the server, which cannot score arbitrary student-selected tokens). This setting does not
+            affect the forward-KL support, which always uses the teacher's top-1 token.
     """
 
     teacher_model_server_url: str | None = field(
@@ -37,10 +42,20 @@ class ServerDistillationConfig(DistillationConfig):
             "help": 'Base URL of a vLLM server hosting the teacher model (e.g., "http://localhost:8000"). Required.'
         },
     )
+    reverse_kl_top_1_mode: str = field(
+        default="sampled",
+        metadata={
+            "help": "Reverse-KL top-1 token selection when beta > 0 and loss_top_k == 1. "
+            "Use 'sampled' for the actual completion token or 'argmax' for the student's top-1 token. "
+            "The forward-KL support always uses the teacher's top-1 token."
+        },
+    )
 
     def __post_init__(self):
         super().__post_init__()
 
+        if self.reverse_kl_top_1_mode not in {"sampled", "argmax"}:
+            raise ValueError("reverse_kl_top_1_mode must be one of: 'sampled', 'argmax'")
         if self.use_liger_kernel:
             raise ValueError(
                 "use_liger_kernel=True is not supported by ServerDistillationTrainer because the Liger loss path "
diff --git a/trl/experimental/server_distillation/server_distillation_trainer.py b/trl/experimental/server_distillation/server_distillation_trainer.py
index acf6e021131..db8be2f1e52 100644
--- a/trl/experimental/server_distillation/server_distillation_trainer.py
+++ b/trl/experimental/server_distillation/server_distillation_trainer.py
@@ -111,10 +111,73 @@ def __init__(
             peft_config=peft_config,
         )
 
+        self.reverse_kl_top_1_mode = args.reverse_kl_top_1_mode
+
         from ...generation.vllm_client import VLLMClient
 
         self.teacher_client = VLLMClient(base_url=args.teacher_model_server_url, connection_timeout=60.0)
 
+    def _get_reverse_kl_top_1_tokens(
+        self, student_scores: torch.Tensor, completion_tokens: torch.Tensor
+    ) -> torch.Tensor:
+        """Return the reverse-KL top-1 token IDs for the mixed top-1 loss path.
+
+        Args:
+            student_scores: Any (B, T, V) tensor whose argmax selects the student's top token
+                (logits or log-probs — both are order-preserving).
+            completion_tokens: (B, T) actual token IDs in the completion.
+        """
+        if self.reverse_kl_top_1_mode == "argmax":
+            return student_scores.argmax(dim=-1)
+        return completion_tokens
+
+    def _compute_sparse_top_1_divergence_loss(
+        self,
+        student_log_probs: torch.Tensor,
+        teacher_top1_token_ids: torch.Tensor,
+        teacher_top1_logprobs: torch.Tensor,
+        reverse_token_ids: torch.Tensor,
+        reverse_teacher_logprobs: torch.Tensor,
+        labels: torch.Tensor,
+        num_items_in_batch=None,
+    ) -> torch.Tensor:
+        """Compute exact generalized JSD/KL on top-1 support for the mixed beta>0 path."""
+        neg_inf = torch.full((), float("-inf"), dtype=student_log_probs.dtype, device=student_log_probs.device)
+
+        if self.beta == 1:
+            support = reverse_token_ids.unsqueeze(-1)
+            support_mask = torch.ones_like(support, dtype=torch.bool)
+            teacher_support_logprobs = reverse_teacher_logprobs.unsqueeze(-1)
+        else:
+            teacher_support = teacher_top1_token_ids.unsqueeze(-1)
+            reverse_support = reverse_token_ids.unsqueeze(-1)
+            support = torch.cat([teacher_support, reverse_support], dim=-1)
+            support_mask = torch.ones_like(support, dtype=torch.bool)
+            support_mask[..., 1] = support[..., 1] != support[..., 0]
+            teacher_support_logprobs = torch.stack([teacher_top1_logprobs, reverse_teacher_logprobs], dim=-1)
+            support = torch.where(support_mask, support, torch.zeros_like(support))
+
+        student_support_logprobs = student_log_probs.gather(-1, support)
+        student_support_logprobs = torch.where(support_mask, student_support_logprobs, neg_inf)
+        teacher_support_logprobs = torch.where(support_mask, teacher_support_logprobs, neg_inf)
+
+        if self.loss_add_tail:
+            base_support_mask = support_mask
+            student_sparse_log_probs, support_mask = _add_tail_bucket(student_support_logprobs, base_support_mask)
+            teacher_sparse_log_probs, _ = _add_tail_bucket(teacher_support_logprobs, base_support_mask)
+        else:
+            student_sparse_log_probs = student_support_logprobs - torch.logsumexp(
+                student_support_logprobs, dim=-1, keepdim=True
+            )
+            teacher_sparse_log_probs = teacher_support_logprobs - torch.logsumexp(
+                teacher_support_logprobs, dim=-1, keepdim=True
+            )
+
+        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
+        )
+
     def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
         # Student forward pass
         student_outputs = model(

From 6458636831bafcb6711c533b3d9de15993225078 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Sun, 19 Jul 2026 20:42:13 +0000
Subject: [PATCH 07/44] Remove top-k support from the base generalized JSD loss
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

The base trainer always has full teacher logits, so it never needs the top-k
approximation — only the extracted server path (which sees just the teacher's
top-k logprobs) does.

- `generalized_jsd_loss` drops the `top_k` / `add_tail` parameters and the whole
  top-k support-selection block; it is now the pure full-vocabulary path.
- config `loss_top_k` / `loss_add_tail` move from `DistillationConfig` to
  `ServerDistillationConfig`; the `self.loss_top_k` / `self.loss_add_tail`
  assignments move from the base `__init__` to the subclass `__init__`; the
  base `compute_loss` call drops the two kwargs.
- `_add_tail_bucket` was only reachable from the top-k block and the server
  methods, so it moves to the server trainer module. `_jsd_divergence` stays
  shared in the base module (the base uses its dense branch, the server its
  masked branch).
- `test_loss_normalizes_by_num_items_in_batch` drops its `loss_top_k`
  parametrization (base is full-vocab only); the `_add_tail_bucket` import in
  the server test now comes from the server module.

Verified on CPU that `generalized_jsd_loss` still matches a naive full-vocab
reference for beta in {0, 0.5, 1} with and without labels, and that the server
sparse loss is numerically unchanged.
---
 .../experimental/test_distillation_trainer.py | 14 ++--
 .../test_server_distillation_trainer.py       |  3 +-
 .../distillation/distillation_config.py       | 25 ------
 .../distillation/distillation_trainer.py      | 78 +------------------
 .../server_distillation_config.py             | 21 +++++
 .../server_distillation_trainer.py            | 17 +++-
 6 files changed, 50 insertions(+), 108 deletions(-)

diff --git a/tests/experimental/test_distillation_trainer.py b/tests/experimental/test_distillation_trainer.py
index 674bb70ebff..ee1c40387ac 100644
--- a/tests/experimental/test_distillation_trainer.py
+++ b/tests/experimental/test_distillation_trainer.py
@@ -249,7 +249,9 @@ def test_distillation_trainer_train_runs_with_local_teacher(self):
 
         assert trainer.state.log_history[-1]["train_loss"] is not None
         assert trainer.state.log_history[0]["eval_loss"] is not None
-        assert train_result.metrics["train_loss"] >= 0.0
+        # Self-distillation (teacher == student), so the divergence is ~0; allow tiny floating-point noise below zero
+        # while still catching a genuinely negative loss.
+        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])
@@ -356,13 +358,11 @@ def test_init_with_eval_dataset(self, eval_dataset_type):
         else:
             assert trainer.eval_dataset is eval_dataset
 
-    @pytest.mark.parametrize("loss_top_k", [0, 1])
-    def test_loss_normalizes_by_num_items_in_batch(self, loss_top_k):
+    def test_loss_normalizes_by_num_items_in_batch(self):
         # When `num_items_in_batch` is passed (as under gradient accumulation), the divergence loss must be reduced as
-        # sum / num_items_in_batch rather than the local per-microbatch mean. See issue #4719. Both the full-vocabulary
-        # JSD path (loss_top_k=0) and the default mixed top-1 path (loss_top_k=1) route through
-        # `_reduce_divergence_loss`, so both must honor `num_items_in_batch`.
-        trainer = self._make_local_trainer(beta=0.5, loss_top_k=loss_top_k)
+        # sum / num_items_in_batch rather than the local per-microbatch mean. See issue #4719. The full-vocabulary JSD
+        # path routes through `_reduce_divergence_loss`, which must honor `num_items_in_batch`.
+        trainer = self._make_local_trainer(beta=0.5)
 
         # Diverge the teacher from the student so the divergence is well above fp noise (else the loss is ~0).
         torch.manual_seed(0)
diff --git a/tests/experimental/test_server_distillation_trainer.py b/tests/experimental/test_server_distillation_trainer.py
index b85f7c689fd..ac2c816189c 100644
--- a/tests/experimental/test_server_distillation_trainer.py
+++ b/tests/experimental/test_server_distillation_trainer.py
@@ -21,12 +21,13 @@
 from datasets import Dataset
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
-from trl.experimental.distillation.distillation_trainer import _add_tail_bucket, _jsd_divergence
+from trl.experimental.distillation.distillation_trainer import _jsd_divergence
 from trl.experimental.server_distillation import (
     ServerDistillationConfig,
     ServerDistillationTrainer,
     build_teacher_request_inputs,
 )
+from trl.experimental.server_distillation.server_distillation_trainer import _add_tail_bucket
 
 from ..testing_utils import TrlTestCase
 
diff --git a/trl/experimental/distillation/distillation_config.py b/trl/experimental/distillation/distillation_config.py
index b9bd0cdb03a..8dc9274f444 100644
--- a/trl/experimental/distillation/distillation_config.py
+++ b/trl/experimental/distillation/distillation_config.py
@@ -74,15 +74,6 @@ class DistillationConfig(_BaseConfig):
         teacher_model_init_kwargs (`dict[str, Any]` or `None`, *optional*):
             Keyword arguments passed to `AutoModelForCausalLM.from_pretrained` when instantiating the teacher model
             from a string.
-        loss_top_k (`int`, *optional*, defaults to `1`):
-            Number of top tokens to use when computing the JSD/KL loss. Both student and teacher distributions are
-            restricted to these K tokens and re-normalized before computing divergence. If 0, the full vocabulary
-            is used. For local teachers, the general support rule is teacher top-k for forward KL, student top-k for
-            reverse KL, and the union for mixed JSD.
-        loss_add_tail (`bool`, *optional*, defaults to `True`):
-            Whether to append a tail bucket that represents the remaining probability mass outside the selected top-k
-            support when computing the loss.
-
         > Parameters that control on-policy generation
 
         num_generations (`int`, *optional*, defaults to `1`):
@@ -220,22 +211,6 @@ class DistillationConfig(_BaseConfig):
             "help": "Keyword arguments for `AutoModelForCausalLM.from_pretrained` when instantiating the teacher."
         },
     )
-    loss_top_k: int = field(
-        default=1,
-        metadata={
-            "help": "Number of top tokens to use when computing the JSD/KL loss. "
-            "Both student and teacher distributions are restricted to these K tokens "
-            "(selected based on beta: teacher's top-k for forward KL, student's top-k for reverse KL, "
-            "union of both for JSD) and re-normalized before computing divergence. "
-            "If 0, the full vocabulary is used (slower but exact)."
-        },
-    )
-    loss_add_tail: bool = field(
-        default=True,
-        metadata={
-            "help": "Whether to append a tail bucket representing the remaining probability mass outside the selected top-k support."
-        },
-    )
 
     # On-policy generation
     num_generations: int = field(
diff --git a/trl/experimental/distillation/distillation_trainer.py b/trl/experimental/distillation/distillation_trainer.py
index 7b96114df00..3716611a9ed 100644
--- a/trl/experimental/distillation/distillation_trainer.py
+++ b/trl/experimental/distillation/distillation_trainer.py
@@ -103,19 +103,6 @@ def _print_completions_sample(prompts: list[str], completions: list[str], step:
     console.print(panel)
 
 
-def _add_tail_bucket(log_probs, valid_mask):
-    """Append a (K+1)-th tail element: log(1 - sum(exp(top_k_logps))).
-
-    This creates a proper probability distribution over K+1 elements, preventing trivial zero loss when top_k is small
-    (especially top_k=1).
-    """
-    log_sum = torch.logsumexp(log_probs, dim=-1, keepdim=True)
-    log_sum = torch.clamp(log_sum, max=-1e-7)  # ensure sum < 1
-    tail = torch.log(-torch.expm1(log_sum))  # log(1 - exp(log_sum))
-    tail_mask = torch.ones_like(valid_mask[..., :1], dtype=torch.bool)
-    return torch.cat([log_probs, tail], dim=-1), torch.cat([valid_mask, tail_mask], dim=-1)
-
-
 def _jsd_divergence(student_log_probs, teacher_log_probs, beta, support_mask=None):
     """Compute JSD (or forward/reverse KL) from log-probability tensors.
 
@@ -540,8 +527,6 @@ def __init__(
         self.temperature = args.temperature
         self.top_p = args.top_p
         self.num_generations = args.num_generations
-        self.loss_top_k = args.loss_top_k
-        self.loss_add_tail = args.loss_add_tail
 
         # ── Buffer state ──
         self._buffered_inputs = None
@@ -1041,12 +1026,10 @@ def generalized_jsd_loss(
         beta=0.5,
         temperature=1.0,
         reduction="batchmean",
-        top_k=0,
-        add_tail=True,
         num_items_in_batch=None,
     ):
         """
-        Compute the generalized Jensen-Shannon Divergence loss for knowledge distillation.
+        Compute the generalized Jensen-Shannon Divergence loss for knowledge distillation over the full vocabulary.
 
         Args:
             student_logits: Tensor of shape (batch_size, sequence_length, vocab_size).
@@ -1055,12 +1038,6 @@ def generalized_jsd_loss(
             beta: Interpolation coefficient. 0.0 = forward KL, 1.0 = reverse KL.
             temperature: Softmax temperature.
             reduction: 'batchmean', 'sum', 'mean', or 'none'.
-            top_k: Number of top tokens to restrict the loss to. The support set depends on beta:
-                beta=0 (forward KL) uses teacher's top-k, beta=1 (reverse KL) uses student's top-k, 0<beta<1 (JSD) uses
-                the union of both. Distributions are re-normalized over the selected support. If 0, the full vocabulary
-                is used.
-            add_tail: Whether to append a tail bucket representing the remaining probability mass
-                outside the selected top-k support.
 
         Returns:
             Scalar loss tensor.
@@ -1068,55 +1045,10 @@ def generalized_jsd_loss(
         student_logits = student_logits / temperature
         teacher_logits = teacher_logits / temperature
 
-        support_mask = None
-
-        if top_k > 0 and student_logits.size(-1) > top_k:
-            neg_inf = torch.full((), float("-inf"), dtype=student_logits.dtype, device=student_logits.device)
-            student_log_probs_full = F.log_softmax(student_logits, dim=-1)
-            teacher_log_probs_full = F.log_softmax(teacher_logits, dim=-1)
-
-            if beta == 0:
-                # Forward KL: teacher-selected support
-                _, support = teacher_logits.topk(top_k, dim=-1)
-                support_mask = torch.ones_like(support, dtype=torch.bool)
-            elif beta == 1:
-                # Reverse KL: student-selected support
-                _, support = student_logits.topk(top_k, dim=-1)
-                support_mask = torch.ones_like(support, dtype=torch.bool)
-            else:
-                # JSD: union of both supports (concatenate + deduplicate)
-                _, student_top = student_logits.topk(top_k, dim=-1)
-                _, teacher_top = teacher_logits.topk(top_k, dim=-1)
-                support = torch.cat([teacher_top, student_top], dim=-1)
-                support_mask = torch.ones(support.shape, dtype=torch.bool, device=support.device)
-                for i in range(1, support.shape[-1]):
-                    prev_matches = support[..., i : i + 1] == support[..., :i]
-                    prev_valid = support_mask[..., :i]
-                    support_mask[..., i] &= ~(prev_matches & prev_valid).any(dim=-1)
-                support = torch.where(support_mask, support, torch.zeros_like(support))
-
-            student_support_logps = student_log_probs_full.gather(-1, support)
-            teacher_support_logps = teacher_log_probs_full.gather(-1, support)
-
-            # Mask invalid (duplicate) positions with -inf
-            student_topk_logps = torch.where(support_mask, student_support_logps, neg_inf)
-            teacher_topk_logps = torch.where(support_mask, teacher_support_logps, neg_inf)
-
-            if add_tail:
-                # Add tail bucket: append log(1 - sum(exp(top_k_logps))) to preserve
-                # the remaining probability mass outside the top-k. This prevents trivial
-                # zero loss when top_k is small (especially top_k=1).
-                base_support_mask = support_mask
-                student_log_probs, support_mask = _add_tail_bucket(student_topk_logps, base_support_mask)
-                teacher_log_probs, _ = _add_tail_bucket(teacher_topk_logps, base_support_mask)
-            else:
-                student_log_probs = student_topk_logps - torch.logsumexp(student_topk_logps, dim=-1, keepdim=True)
-                teacher_log_probs = teacher_topk_logps - torch.logsumexp(teacher_topk_logps, dim=-1, keepdim=True)
-        else:
-            student_log_probs = F.log_softmax(student_logits, dim=-1)
-            teacher_log_probs = F.log_softmax(teacher_logits, dim=-1)
+        student_log_probs = F.log_softmax(student_logits, dim=-1)
+        teacher_log_probs = F.log_softmax(teacher_logits, dim=-1)
 
-        jsd = _jsd_divergence(student_log_probs, teacher_log_probs, beta, support_mask)
+        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
         )
@@ -1157,8 +1089,6 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
             labels=labels,
             beta=self.beta,
             temperature=self.temperature,
-            top_k=self.loss_top_k,
-            add_tail=self.loss_add_tail,
             num_items_in_batch=num_items_in_batch,
         )
 
diff --git a/trl/experimental/server_distillation/server_distillation_config.py b/trl/experimental/server_distillation/server_distillation_config.py
index dfea92faf65..808d85824a4 100644
--- a/trl/experimental/server_distillation/server_distillation_config.py
+++ b/trl/experimental/server_distillation/server_distillation_config.py
@@ -29,6 +29,13 @@ class ServerDistillationConfig(DistillationConfig):
     Parameters:
         teacher_model_server_url (`str` or `None`, *optional*):
             Base URL of a vLLM server hosting the teacher model (e.g., `"http://localhost:8000"`). Required.
+        loss_top_k (`int`, *optional*, defaults to `1`):
+            Number of top tokens the teacher server returns and over which the divergence is computed. For `beta == 0`
+            (pure forward KL) this must be positive; for `beta > 0` it must be `1` (top-1 support). `0` (full
+            vocabulary) is not available with a server teacher, which only exposes top-k logprobs.
+        loss_add_tail (`bool`, *optional*, defaults to `True`):
+            Whether to append a tail bucket representing the remaining probability mass outside the selected top-k
+            support when computing the loss.
         reverse_kl_top_1_mode (`str`, *optional*, defaults to `"sampled"`):
             Selection rule for the reverse-KL top-1 token when `beta > 0` and `loss_top_k == 1`. `"sampled"` uses the
             actual completion token in the batch. `"argmax"` uses the student's highest-probability token (not
@@ -42,6 +49,20 @@ class ServerDistillationConfig(DistillationConfig):
             "help": 'Base URL of a vLLM server hosting the teacher model (e.g., "http://localhost:8000"). Required.'
         },
     )
+    loss_top_k: int = field(
+        default=1,
+        metadata={
+            "help": "Number of top tokens the teacher server returns and over which the divergence is computed. "
+            "For beta == 0 this must be positive; for beta > 0 it must be 1."
+        },
+    )
+    loss_add_tail: bool = field(
+        default=True,
+        metadata={
+            "help": "Whether to append a tail bucket representing the remaining probability mass outside the selected "
+            "top-k support."
+        },
+    )
     reverse_kl_top_1_mode: str = field(
         default="sampled",
         metadata={
diff --git a/trl/experimental/server_distillation/server_distillation_trainer.py b/trl/experimental/server_distillation/server_distillation_trainer.py
index db8be2f1e52..2c3d45a3e63 100644
--- a/trl/experimental/server_distillation/server_distillation_trainer.py
+++ b/trl/experimental/server_distillation/server_distillation_trainer.py
@@ -17,10 +17,23 @@
 import torch
 import torch.nn.functional as F
 
-from ..distillation.distillation_trainer import DistillationTrainer, _add_tail_bucket, _jsd_divergence
+from ..distillation.distillation_trainer import DistillationTrainer, _jsd_divergence
 from .server_distillation_config import ServerDistillationConfig
 
 
+def _add_tail_bucket(log_probs, valid_mask):
+    """Append a (K+1)-th tail element: log(1 - sum(exp(top_k_logps))).
+
+    This creates a proper probability distribution over K+1 elements, preventing trivial zero loss when top_k is small
+    (especially top_k=1).
+    """
+    log_sum = torch.logsumexp(log_probs, dim=-1, keepdim=True)
+    log_sum = torch.clamp(log_sum, max=-1e-7)  # ensure sum < 1
+    tail = torch.log(-torch.expm1(log_sum))  # log(1 - exp(log_sum))
+    tail_mask = torch.ones_like(valid_mask[..., :1], dtype=torch.bool)
+    return torch.cat([log_probs, tail], dim=-1), torch.cat([valid_mask, tail_mask], dim=-1)
+
+
 def build_teacher_request_inputs(
     input_ids: torch.Tensor,
     attention_mask: torch.Tensor,
@@ -112,6 +125,8 @@ def __init__(
         )
 
         self.reverse_kl_top_1_mode = args.reverse_kl_top_1_mode
+        self.loss_top_k = args.loss_top_k
+        self.loss_add_tail = args.loss_add_tail
 
         from ...generation.vllm_client import VLLMClient
 

From 166808b6fcb99f75816ce3c67d5abb21c17469f5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Sun, 19 Jul 2026 23:20:41 +0000
Subject: [PATCH 08/44] Validate the teacher by vocab size instead of tokenizer
 identity

The base local-teacher loss compares full next-token distributions, so the only
real requirement is a shared vocabulary. Replace the tokenizer-loading /
`get_vocab()` comparison + warn + `resize_token_embeddings` dance with a strict
`config.vocab_size` equality check that raises (pointing at GOLD for the
cross-tokenizer case), mirroring GKD.

- `__init__` teacher setup no longer loads the teacher tokenizer, no longer
  requires `config._name_or_path`, and no longer resizes the teacher embeddings.
  The check is gated on `teacher_model is not None`, so subclasses that supply
  the teacher another way (ServerDistillationTrainer) are unaffected.
- Remove `_local_teacher_tokenizers_match`, `_raise_if_local_teacher_tokenizer_mismatch`,
  the `self._local_teacher_tokenizer_matches_student` flag, and the guard call at
  the top of `compute_loss`. Drop the now-unused `warnings` import.
- Add tests for the vocab-size mismatch (Qwen2 student + Llama teacher) and for
  passing `teacher_model_init_kwargs` alongside an already-instantiated teacher.

Note: `teacher_model_name_or_path` remains an unused config field for now; it is
removed later when the teacher becomes a constructor-only argument.
---
 .../experimental/test_distillation_trainer.py | 28 +++++-
 .../distillation/distillation_trainer.py      | 91 +++++--------------
 2 files changed, 48 insertions(+), 71 deletions(-)

diff --git a/tests/experimental/test_distillation_trainer.py b/tests/experimental/test_distillation_trainer.py
index ee1c40387ac..59d5318367c 100644
--- a/tests/experimental/test_distillation_trainer.py
+++ b/tests/experimental/test_distillation_trainer.py
@@ -18,7 +18,7 @@
 import torch
 import torch.nn.functional as F
 from datasets import DatasetDict, IterableDatasetDict, load_dataset
-from transformers import AutoTokenizer
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 from trl.experimental.distillation import DistillationConfig, DistillationTrainer
 from trl.experimental.distillation.distillation_trainer import _RepeatBatchDataLoader
@@ -410,6 +410,32 @@ def test_distillation_trainer_with_liger(self):
         finally:
             importlib.reload(importlib.import_module(trainer.model.__module__))
 
+    def test_teacher_vocab_size_mismatch_raises(self):
+        # The local-teacher loss compares full next-token distributions, so student and teacher must share a
+        # vocabulary. A teacher with a different vocab_size is rejected (use GOLD for cross-tokenizer distillation).
+        dataset = load_dataset("trl-internal-testing/zen", "conversational_language_modeling", split="train")
+        with pytest.raises(ValueError, match="vocab_size"):
+            DistillationTrainer(
+                model=self.model_id,
+                teacher_model="trl-internal-testing/tiny-LlamaForCausalLM-3.2",
+                args=self._make_args(),
+                train_dataset=dataset,
+                processing_class=self.tokenizer,
+            )
+
+    def test_teacher_model_init_kwargs_with_instantiated_teacher_raises(self):
+        # `teacher_model_init_kwargs` only applies when the teacher is a model id; passing it alongside an already
+        # instantiated teacher is a mistake worth surfacing.
+        dataset = load_dataset("trl-internal-testing/zen", "conversational_language_modeling", split="train")
+        with pytest.raises(ValueError, match="teacher_model_init_kwargs"):
+            DistillationTrainer(
+                model=self.model_id,
+                teacher_model=AutoModelForCausalLM.from_pretrained(self.model_id),
+                args=self._make_args(),
+                train_dataset=dataset,
+                processing_class=self.tokenizer,
+            )
+
 
 def test_repeat_batch_dataloader_delegates_set_epoch_via_getattr():
     class DummyDataLoader:
diff --git a/trl/experimental/distillation/distillation_trainer.py b/trl/experimental/distillation/distillation_trainer.py
index 3716611a9ed..59d56044170 100644
--- a/trl/experimental/distillation/distillation_trainer.py
+++ b/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
@@ -435,61 +434,23 @@ def __init__(
             self._forward_redirection = _ForwardRedirection()
 
         # ── Teacher model setup ──
-        self._local_teacher_tokenizer_matches_student = True
+        # `teacher_model` may be None: subclasses (e.g. ServerDistillationTrainer) supply the teacher another way.
         if teacher_model is not None:
-            if args.teacher_model_init_kwargs is not None and not isinstance(teacher_model, str):
-                raise ValueError(
-                    "You passed teacher_model_init_kwargs to the config, but your teacher_model is already "
-                    "instantiated."
-                )
-
-            teacher_model_name_or_path = (
-                teacher_model
-                if isinstance(teacher_model, str)
-                else getattr(getattr(teacher_model, "config", None), "_name_or_path", None)
-            )
-            if teacher_model_name_or_path is None:
-                raise ValueError(
-                    "DistillationTrainer requires a local teacher model with `config._name_or_path` set so its "
-                    "tokenizer can be validated against the student tokenizer."
-                )
-
-            teacher_model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code)
-            teacher_tokenizer_kwargs = {}
-            teacher_revision = teacher_model_init_kwargs.get("revision", args.teacher_model_revision)
-            if teacher_revision is not None:
-                teacher_tokenizer_kwargs["revision"] = teacher_revision
-            teacher_tokenizer_kwargs["trust_remote_code"] = teacher_model_init_kwargs["trust_remote_code"]
-            teacher_processing_class = AutoTokenizer.from_pretrained(
-                teacher_model_name_or_path, **teacher_tokenizer_kwargs
-            )
-            if getattr(teacher_processing_class, "pad_token", None) is None:
-                teacher_processing_class.pad_token = teacher_processing_class.eos_token
-            self._local_teacher_tokenizer_matches_student = self._local_teacher_tokenizers_match(
-                processing_class, teacher_processing_class
-            )
-            if not self._local_teacher_tokenizer_matches_student:
-                warnings.warn(
-                    "DistillationTrainer's built-in local-teacher loss assumes the student and teacher share the "
-                    "same tokenizer. The loaded local teacher tokenizer does not match the student tokenizer, so "
-                    "the teacher weights will be left unchanged for subclass overrides. Direct use of the base "
-                    "DistillationTrainer with this local teacher remains unsupported.",
-                    UserWarning,
-                    stacklevel=2,
-                )
-
             if isinstance(teacher_model, str):
+                teacher_model_init_kwargs.setdefault("trust_remote_code", args.trust_remote_code)
                 dtype = teacher_model_init_kwargs.get("dtype")
                 teacher_model_init_kwargs["dtype"] = dtype if dtype in ["auto", None] else getattr(torch, dtype)
-
-            if isinstance(teacher_model, str):
-                init_kwargs = dict(teacher_model_init_kwargs)
                 if args.teacher_model_revision is not None:
-                    init_kwargs.setdefault("revision", args.teacher_model_revision)
+                    teacher_model_init_kwargs.setdefault("revision", args.teacher_model_revision)
                 # Distributed training requires device_map=None ("auto" fails)
                 if args.distributed_state.distributed_type in ["MULTI_GPU", "DEEPSPEED"]:
-                    init_kwargs["device_map"] = None
-                teacher_model = create_model_from_path(teacher_model, **init_kwargs)
+                    teacher_model_init_kwargs["device_map"] = None
+                teacher_model = create_model_from_path(teacher_model, **teacher_model_init_kwargs)
+            elif args.teacher_model_init_kwargs is not None:
+                raise ValueError(
+                    "You passed teacher_model_init_kwargs to the config, but your teacher_model is already "
+                    "instantiated."
+                )
 
         # Trainer does not need to remove unused columns — the collator handles raw data
         args.remove_unused_columns = False
@@ -509,8 +470,17 @@ def __init__(
 
         # ── Prepare teacher model (after super().__init__ so accelerator is ready) ──
         if teacher_model is not None:
-            if self._local_teacher_tokenizer_matches_student:
-                teacher_model.resize_token_embeddings(self.model.config.get_text_config().vocab_size)
+            # The divergence compares the full next-token distribution of the student against the teacher's, so both
+            # must be defined over the same vocabulary.
+            student_vocab_size = self.model.config.get_text_config().vocab_size
+            teacher_vocab_size = teacher_model.config.get_text_config().vocab_size
+            if student_vocab_size != teacher_vocab_size:
+                raise ValueError(
+                    f"The student model has vocab_size {student_vocab_size} but the teacher model has vocab_size "
+                    f"{teacher_vocab_size}. Distillation compares the teacher's full next-token distribution, which "
+                    f"requires a shared vocabulary. Use a teacher with the same vocab_size, or GOLD for "
+                    f"cross-tokenizer distillation."
+                )
             if self.is_deepspeed_enabled:
                 self.teacher_model = prepare_deepspeed(teacher_model, self.accelerator)
             else:
@@ -604,23 +574,6 @@ def __init__(
             self.vllm_sync_frequency = args.vllm_sync_frequency
             self._last_vllm_sync_step = -1
 
-    @staticmethod
-    def _local_teacher_tokenizers_match(
-        student_processing_class: PreTrainedTokenizerBase,
-        teacher_processing_class: PreTrainedTokenizerBase,
-    ) -> bool:
-        """Check whether the student and local teacher tokenizers share the same vocabulary."""
-        return student_processing_class.get_vocab() == teacher_processing_class.get_vocab()
-
-    def _raise_if_local_teacher_tokenizer_mismatch(self) -> None:
-        """Guard the base local-teacher JSD path, while still allowing subclass overrides."""
-        if self.teacher_model is not None and not self._local_teacher_tokenizer_matches_student:
-            raise ValueError(
-                "DistillationTrainer's built-in local-teacher loss only supports student/teacher pairs that use "
-                "the same tokenizer. Use a same-tokenizer local teacher, use ServerDistillationTrainer with an "
-                "external teacher server, or override the local teacher loss path in a subclass."
-            )
-
     def _compute_prompt_length(self, inputs: dict[str, torch.Tensor | Any]) -> int:
         """Compute the earliest prompt boundary that still includes every completion token in the batch."""
         if inputs.get("labels") is not None:
@@ -1065,8 +1018,6 @@ def _get_teacher_logits(self, inputs: dict[str, torch.Tensor | Any]) -> torch.Te
             ).logits
 
     def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
-        self._raise_if_local_teacher_tokenizer_mismatch()
-
         if self.use_liger_loss:
             loss = self._compute_liger_loss(model, inputs, num_items_in_batch=num_items_in_batch)
             return (loss, None) if return_outputs else loss

From 952b0c321526fe027fdfbdd9ccac0049b3098685 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Sun, 19 Jul 2026 23:43:12 +0000
Subject: [PATCH 09/44] Deprecate `lmbda` on/off-policy mixing

The distillation trainer is converging on always-on-policy training (the
GRPO-shaped stable target generates every batch). `lmbda`, which mixes
on-policy (student-generated) and off-policy (dataset) completions per
gradient-accumulation slice, is being removed.

Warn with FutureWarning when `lmbda != 1.0`, pointing at `GKDTrainer`, which
exposes the same knob for users who need on/off-policy mixing. The default
(`lmbda=1.0`, fully on-policy) does not warn.

Deprecate-before-remove: the removal lands in a later PR. Both policy modes are
already pinned end to end (see the on/off-policy training test), so that removal
will be a visible, guarded change.
---
 tests/experimental/test_distillation_trainer.py      | 12 ++++++++++++
 trl/experimental/distillation/distillation_config.py |  8 ++++++++
 2 files changed, 20 insertions(+)

diff --git a/tests/experimental/test_distillation_trainer.py b/tests/experimental/test_distillation_trainer.py
index 59d5318367c..24472bbb4d4 100644
--- a/tests/experimental/test_distillation_trainer.py
+++ b/tests/experimental/test_distillation_trainer.py
@@ -27,6 +27,18 @@
 from ..testing_utils import TrlTestCase, require_liger_kernel, require_torch_accelerator
 
 
+def test_lmbda_deprecation_warns(tmp_path):
+    # `lmbda` (on/off-policy mixing) is on its way out; the trainer is becoming always on-policy. Setting it to
+    # anything other than 1.0 warns and points at GKDTrainer.
+    with pytest.warns(FutureWarning, match="lmbda"):
+        DistillationConfig(output_dir=str(tmp_path), use_cpu=True, bf16=False, lmbda=0.5)
+
+
+def test_lmbda_default_does_not_warn(tmp_path, recwarn):
+    DistillationConfig(output_dir=str(tmp_path), use_cpu=True, bf16=False)
+    assert not [w for w in recwarn.list if issubclass(w.category, FutureWarning) and "lmbda" in str(w.message)]
+
+
 def _reference_generalized_jsd(student_logits, teacher_logits, labels=None, beta=0.5, temperature=1.0):
     """Naive reference for the generalized JSD, written straight from the definition.
 
diff --git a/trl/experimental/distillation/distillation_config.py b/trl/experimental/distillation/distillation_config.py
index 8dc9274f444..b11afc79f03 100644
--- a/trl/experimental/distillation/distillation_config.py
+++ b/trl/experimental/distillation/distillation_config.py
@@ -316,6 +316,14 @@ def __post_init__(self):
 
         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.lmbda != 1.0:
+            warnings.warn(
+                "`lmbda` (on/off-policy mixing) is deprecated and will be removed: the distillation trainer is "
+                "becoming always on-policy. For on/off-policy mixing, use `GKDTrainer`, which exposes the same "
+                "`lmbda` knob.",
+                FutureWarning,
+                stacklevel=3,
+            )
         if self.beta < 0.0 or self.beta > 1.0:
             raise ValueError(f"beta must be in [0.0, 1.0], got {self.beta}.")
 

From 55e5973fc8f681c38a2a28eba2727d78ac7c5ea9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Sun, 19 Jul 2026 23:44:14 +0000
Subject: [PATCH 10/44] Point the GKD paper reproduction at GKDTrainer

The GKD (2306.13649) reproduction sets `lmbda=0.5` for the paper's on/off-policy
mixing. `DistillationTrainer` is dropping `lmbda` (becoming always on-policy),
so move the reproduction snippet to `GKDConfig` / `GKDTrainer`, which keep that
knob. Translate `max_completion_length` to GKD's `max_new_tokens`.

The prose now frames GKDTrainer as the home for the paper's `lmbda` mixing and
DistillationTrainer as the same generalized-JSD objective for the always-on-policy
case. Must land before `lmbda` is removed from DistillationConfig, so the index
never references a removed argument.
---
 docs/source/paper_index.md | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/docs/source/paper_index.md b/docs/source/paper_index.md
index 5e815a8a86f..4967b8cb413 100644
--- a/docs/source/paper_index.md
+++ b/docs/source/paper_index.md
@@ -1665,13 +1665,13 @@ 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)
@@ -1679,7 +1679,7 @@ training_args = DistillationConfig(
     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)
 )
 ```
 

From 0b72a4e2ea109548d5856ed25445719d533e5653 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Sun, 19 Jul 2026 23:53:23 +0000
Subject: [PATCH 11/44] Remove `lmbda` and the off-policy training branch

The trainer is now always on-policy: every generation-batch slice is generated
by the student, and there is no dataset-completion (off-policy) path.

Trainer:
- `_fill_buffer` generates for every slice; the per-slice on/off-policy coin flip
  (`random.random() <= lmbda`, broadcast across ranks) and the off-policy
  "store the dataset slice directly" branch are gone.
- Remove `self.lmbda`, `self._buffered_on_policy_flags`, and the four on/off-policy
  loss accumulators. `training_step` drops the on/off-policy attribution (its
  completion-length metrics lose the `on_policy_`/`off_policy_` prefix), and `log`
  drops the `on_policy_loss` / `off_policy_loss` emission and the all-reduce of
  the accumulator vector.
- Drop the now-unused `random` coin flip, `broadcast_object_list`, and
  `torch.distributed` / `DistributedType` imports (rich still uses `random`).

Config:
- Remove `lmbda`, its range check, the deprecation warning added last PR, and the
  `num_generations > 1 and lmbda < 1.0` warning.

Tests:
- `_make_args` drops `lmbda=0.0`; the on/off-policy training test becomes a single
  always-on-policy `test_train_updates_params`; the `lmbda` deprecation tests are
  removed; the server ragged-grad test drops `lmbda=0.0`.

Both policy modes were pinned before this (PR 02 / PR 10), so the off-policy
removal is a deliberate, test-visible deletion. Full experimental suite green.
---
 .../experimental/test_distillation_trainer.py |  25 +---
 .../test_server_distillation_trainer.py       |   1 -
 .../distillation/distillation_config.py       |  29 -----
 .../distillation/distillation_trainer.py      | 118 ++++--------------
 4 files changed, 27 insertions(+), 146 deletions(-)

diff --git a/tests/experimental/test_distillation_trainer.py b/tests/experimental/test_distillation_trainer.py
index 24472bbb4d4..78af1c861a7 100644
--- a/tests/experimental/test_distillation_trainer.py
+++ b/tests/experimental/test_distillation_trainer.py
@@ -27,18 +27,6 @@
 from ..testing_utils import TrlTestCase, require_liger_kernel, require_torch_accelerator
 
 
-def test_lmbda_deprecation_warns(tmp_path):
-    # `lmbda` (on/off-policy mixing) is on its way out; the trainer is becoming always on-policy. Setting it to
-    # anything other than 1.0 warns and points at GKDTrainer.
-    with pytest.warns(FutureWarning, match="lmbda"):
-        DistillationConfig(output_dir=str(tmp_path), use_cpu=True, bf16=False, lmbda=0.5)
-
-
-def test_lmbda_default_does_not_warn(tmp_path, recwarn):
-    DistillationConfig(output_dir=str(tmp_path), use_cpu=True, bf16=False)
-    assert not [w for w in recwarn.list if issubclass(w.category, FutureWarning) and "lmbda" in str(w.message)]
-
-
 def _reference_generalized_jsd(student_logits, teacher_logits, labels=None, beta=0.5, temperature=1.0):
     """Naive reference for the generalized JSD, written straight from the definition.
 
@@ -210,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},
@@ -266,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)
         previous_params = {name: param.clone() for name, param in trainer.model.named_parameters()}
 
         trainer.train()
diff --git a/tests/experimental/test_server_distillation_trainer.py b/tests/experimental/test_server_distillation_trainer.py
index ac2c816189c..87639f550ef 100644
--- a/tests/experimental/test_server_distillation_trainer.py
+++ b/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",
diff --git a/trl/experimental/distillation/distillation_config.py b/trl/experimental/distillation/distillation_config.py
index b11afc79f03..6babe3d43f4 100644
--- a/trl/experimental/distillation/distillation_config.py
+++ b/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,16 +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.lmbda != 1.0:
-            warnings.warn(
-                "`lmbda` (on/off-policy mixing) is deprecated and will be removed: the distillation trainer is "
-                "becoming always on-policy. For on/off-policy mixing, use `GKDTrainer`, which exposes the same "
-                "`lmbda` knob.",
-                FutureWarning,
-                stacklevel=3,
-            )
         if self.beta < 0.0 or self.beta > 1.0:
             raise ValueError(f"beta must be in [0.0, 1.0], got {self.beta}.")
 
@@ -350,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,
-            )
diff --git a/trl/experimental/distillation/distillation_trainer.py b/trl/experimental/distillation/distillation_trainer.py
index 59d56044170..fd06f8c45a7 100644
--- a/trl/experimental/distillation/distillation_trainer.py
+++ b/trl/experimental/distillation/distillation_trainer.py
@@ -21,10 +21,9 @@
 from typing import Any, Optional
 
 import torch
-import torch.distributed as dist
 import torch.nn as nn
 import torch.nn.functional as F
-from accelerate.utils import DistributedType, broadcast_object_list, gather_object
+from accelerate.utils import gather_object
 from datasets import Dataset
 from packaging.version import Version
 from torch.utils.data import DataLoader
@@ -150,8 +149,8 @@ class _DistillationCollator:
     """Data collator for the distillation trainer with independent prompt/completion budgets.
 
     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) for
-    pure on-policy distillation (``lmbda=1``).
+    completions can never truncate the prompt to empty. It also handles prompt-only data (no assistant completions),
+    which is what on-policy distillation uses.
     """
 
     def __init__(
@@ -300,8 +299,8 @@ class DistillationTrainer(_BaseTrainer):
 
     Supports:
     - Generalized JSD loss (forward KL, reverse KL, or interpolated JSD via `beta`)
-    - On-policy / off-policy mixing via `lmbda` (buffered across gradient accumulation)
-    - Local teacher model or external teacher via vLLM server
+    - On-policy distillation: the student generates completions, the teacher scores them
+    - Local teacher model
     - Student on-policy generation via vLLM or model.generate()
     - Liger kernel for memory-efficient fused JSD loss
     """
@@ -492,7 +491,6 @@ def __init__(
             disable_dropout_in_model(self.model)
 
         # ── Store config values ──
-        self.lmbda = args.lmbda
         self.beta = args.beta
         self.temperature = args.temperature
         self.top_p = args.top_p
@@ -500,15 +498,10 @@ def __init__(
 
         # ── Buffer state ──
         self._buffered_inputs = None
-        self._buffered_on_policy_flags = None
         self._buffered_text_logs = None
         self._buffer_step = 0
 
         # ── Loss tracking ──
-        self._on_policy_loss_total = 0.0
-        self._off_policy_loss_total = 0.0
-        self._on_policy_step_equiv = 0.0
-        self._off_policy_step_equiv = 0.0
 
         # ── Generation config ──
         generation_kwargs = {
@@ -697,43 +690,26 @@ def _prepare_inputs(self, generation_batch: dict[str, torch.Tensor | Any]) -> di
 
     @profiling_decorator
     def _fill_buffer(self, generation_batch: dict[str, torch.Tensor | Any], buffer_steps: int):
-        """Split batch into slices and decide which are on-policy (student-generated) vs off-policy."""
+        """Split the batch into slices and generate student completions for each (always on-policy)."""
         slices = split_tensor_dict(generation_batch, buffer_steps)
 
-        # Decide on-policy flags (synchronized across processes)
-        if self.accelerator.is_main_process:
-            on_policy_flags = [random.random() <= self.lmbda for _ in range(buffer_steps)]
-        else:
-            on_policy_flags = [False] * buffer_steps
-        on_policy_flags = broadcast_object_list(on_policy_flags, from_process=0)
-
         self._buffered_inputs = [None] * buffer_steps
-        self._buffered_on_policy_flags = on_policy_flags
         self._buffered_text_logs = [None] * buffer_steps
 
-        # Store off-policy slices directly
-        on_policy_indices = []
-        for i, is_on_policy in enumerate(on_policy_flags):
-            if is_on_policy:
-                on_policy_indices.append(i)
-            else:
-                self._buffered_inputs[i] = slices[i]
-
-        # Generate student completions for on-policy slices
-        if on_policy_indices:
-            self._generate_student_completions(slices, on_policy_indices)
+        # Generate student completions for every slice
+        self._generate_student_completions(slices, list(range(buffer_steps)))
 
-        # Gather on-policy text logs once per optimizer step (all processes must participate)
+        # Gather text logs once per optimizer step (all processes must participate)
         if self.log_completions:
-            on_policy_prompts = []
-            on_policy_completions = []
-            for i in on_policy_indices:
-                if self._buffered_text_logs[i] is not None:
-                    prompts, completions = self._buffered_text_logs[i]
-                    on_policy_prompts.extend(prompts)
-                    on_policy_completions.extend(completions)
-            self._textual_logs["prompt"].extend(gather_object(on_policy_prompts))
-            self._textual_logs["completion"].extend(gather_object(on_policy_completions))
+            prompts_all = []
+            completions_all = []
+            for entry in self._buffered_text_logs:
+                if entry is not None:
+                    prompts, completions = entry
+                    prompts_all.extend(prompts)
+                    completions_all.extend(completions)
+            self._textual_logs["prompt"].extend(gather_object(prompts_all))
+            self._textual_logs["completion"].extend(gather_object(completions_all))
 
     @profiling_decorator
     def _generate_student_completions(self, slices: list[dict[str, torch.Tensor | Any]], on_policy_indices: list[int]):
@@ -1146,7 +1122,7 @@ def _get_liger_zero3_lm_head_gather_ctx(self, model: nn.Module):
     def training_step(
         self, model: nn.Module, inputs: dict[str, torch.Tensor | Any], num_items_in_batch: int | None = None
     ) -> torch.Tensor:
-        """Training step with on/off-policy loss tracking and completion stats."""
+        """Training step: generate on-policy, then run the parent step, tracking completion stats."""
         buffer_steps = self.args.gradient_accumulation_steps
 
         with self._get_liger_zero3_lm_head_gather_ctx(model):
@@ -1154,75 +1130,29 @@ def training_step(
 
         slice_idx = (self._buffer_step - 1) % buffer_steps
 
-        # Determine if this slice is on-policy
-        is_on_policy = False
-        if self._buffered_on_policy_flags is not None and slice_idx < len(self._buffered_on_policy_flags):
-            is_on_policy = self._buffered_on_policy_flags[slice_idx]
-
-        # Track completion length stats — read from buffered inputs (which reflect on-policy generation)
+        # Track completion length stats — read from buffered inputs (which reflect the generated completions)
         actual_inputs = self._buffered_inputs[slice_idx] if self._buffered_inputs is not None else inputs
         labels = actual_inputs.get("labels")
         if labels is not None:
             completion_lengths = (labels != -100).sum(dim=1).float()
             gathered_lengths = self.accelerator.gather(completion_lengths)
             mode = "train"
-            prefix = "on_policy" if is_on_policy else "off_policy"
-            self._metrics[mode][f"completions/{prefix}_mean_length"].append(gathered_lengths.mean().item())
-            self._metrics[mode][f"completions/{prefix}_max_length"].append(gathered_lengths.max().item())
-            self._metrics[mode][f"completions/{prefix}_min_length"].append(gathered_lengths.min().item())
+            self._metrics[mode]["completions/mean_length"].append(gathered_lengths.mean().item())
+            self._metrics[mode]["completions/max_length"].append(gathered_lengths.max().item())
+            self._metrics[mode]["completions/min_length"].append(gathered_lengths.min().item())
 
             # Log fraction of completions that hit max_completion_length (truncated)
             max_comp_len = getattr(self.generation_config, "max_new_tokens", None)
-            if is_on_policy and max_comp_len is not None:
+            if max_comp_len is not None:
                 truncated_frac = (gathered_lengths >= max_comp_len).float().mean().item()
                 self._metrics[mode]["completions/truncated_fraction"].append(truncated_frac)
 
-        # Track loss per policy type
-        loss_scalar = float(loss.detach())
-        step_equiv = 1.0 / self.args.gradient_accumulation_steps
-        if is_on_policy:
-            self._on_policy_loss_total += loss_scalar
-            self._on_policy_step_equiv += step_equiv
-        else:
-            self._off_policy_loss_total += loss_scalar
-            self._off_policy_step_equiv += step_equiv
-
         return loss
 
     def log(self, logs: dict[str, float], start_time: float | None = None) -> None:
         mode = "train" if self.model.training else "eval"
         metrics = {key: sum(val) / len(val) for key, val in self._metrics[mode].items()}
 
-        if mode == "train":
-            # Aggregate on/off-policy losses across distributed processes
-            device = self.accelerator.device if hasattr(self.accelerator, "device") else torch.device("cpu")
-            vec = torch.tensor(
-                [
-                    self._on_policy_loss_total,
-                    self._off_policy_loss_total,
-                    self._on_policy_step_equiv,
-                    self._off_policy_step_equiv,
-                ],
-                dtype=torch.float64,
-                device=device,
-            )
-
-            if (
-                getattr(self.accelerator, "distributed_type", DistributedType.NO) != DistributedType.NO
-                and dist.is_available()
-                and dist.is_initialized()
-            ):
-                dist.all_reduce(vec, op=dist.ReduceOp.SUM)
-
-            on_sum, off_sum, on_eq, off_eq = vec.tolist()
-            if on_eq > 0:
-                logs["on_policy_loss"] = round(on_sum / on_eq, 4)
-            if off_eq > 0:
-                logs["off_policy_loss"] = round(off_sum / off_eq, 4)
-
-            self._on_policy_loss_total = self._off_policy_loss_total = 0.0
-            self._on_policy_step_equiv = self._off_policy_step_equiv = 0.0
-
         if mode == "eval":
             metrics = {f"eval_{key}": val for key, val in metrics.items()}
 

From 2527d3d1944abb9431911c373996011b181cd186 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Mon, 20 Jul 2026 01:47:01 +0000
Subject: [PATCH 12/44] style

---
 trl/experimental/iw_opd/iw_opd_trainer.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/trl/experimental/iw_opd/iw_opd_trainer.py b/trl/experimental/iw_opd/iw_opd_trainer.py
index 1acaa8f32c4..685defce7a9 100644
--- a/trl/experimental/iw_opd/iw_opd_trainer.py
+++ b/trl/experimental/iw_opd/iw_opd_trainer.py
@@ -362,8 +362,8 @@ class IWOPDTrainer(_BaseTrainer):
     whose supervision has drifted out of distribution. Select it with `distillation_objective="iw_opd"`.
 
     .. note::
-        This trainer is a frozen snapshot of the pre-refactor ``DistillationTrainer``, kept as the home for IW-OPD
-        (a teacher-guided policy-gradient method that does not fit the stable, full-vocabulary ``DistillationTrainer``).
+        This trainer is a frozen snapshot of the pre-refactor ``DistillationTrainer``, kept as the home for IW-OPD (a
+        teacher-guided policy-gradient method that does not fit the stable, full-vocabulary ``DistillationTrainer``).
         It is not actively maintained and will not track that trainer's later improvements.
 
     Supports:

From e0587640b5c2d0d0e0c532589fd51672bb936b69 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Mon, 20 Jul 2026 02:48:31 +0000
Subject: [PATCH 13/44] style

---
 .../server_distillation/server_distillation_config.py         | 4 ++--
 .../server_distillation/server_distillation_trainer.py        | 3 ++-
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/trl/experimental/server_distillation/server_distillation_config.py b/trl/experimental/server_distillation/server_distillation_config.py
index 4298a38b2f1..d8d8bb452a0 100644
--- a/trl/experimental/server_distillation/server_distillation_config.py
+++ b/trl/experimental/server_distillation/server_distillation_config.py
@@ -23,8 +23,8 @@ class ServerDistillationConfig(DistillationConfig):
     Configuration class for the [`ServerDistillationTrainer`].
 
     Extends [`DistillationConfig`] with the address of an external vLLM server that scores the student's completions.
-    The teacher is never held locally: instead of a local forward pass, per-token teacher logprobs are fetched from
-    the server, so only the teacher's top-k logprobs are available and the loss is restricted to a sparse support.
+    The teacher is never held locally: instead of a local forward pass, per-token teacher logprobs are fetched from the
+    server, so only the teacher's top-k logprobs are available and the loss is restricted to a sparse support.
 
     Parameters:
         teacher_model_server_url (`str` or `None`, *optional*):
diff --git a/trl/experimental/server_distillation/server_distillation_trainer.py b/trl/experimental/server_distillation/server_distillation_trainer.py
index acf6e021131..2f8cd33518a 100644
--- a/trl/experimental/server_distillation/server_distillation_trainer.py
+++ b/trl/experimental/server_distillation/server_distillation_trainer.py
@@ -72,7 +72,8 @@ class ServerDistillationTrainer(DistillationTrainer):
     Instead of running a local teacher forward pass, per-token teacher logprobs are fetched from a vLLM server via
     [`~generation.vllm_client.VLLMClient`]. The server only returns the teacher's top-k logprobs, so the divergence is
     restricted to a sparse support (top-1 for `beta > 0`, top-k for the pure forward path `beta = 0`). Everything else
-    — the student forward, generation, buffering, metrics — is inherited from [`experimental.distillation.DistillationTrainer`].
+    — the student forward, generation, buffering, metrics — is inherited from
+    [`experimental.distillation.DistillationTrainer`].
     """
 
     _tag_names = ["trl", "server-distillation"]

From 456df918c68331e37fb0b0de9fc6365c933d946f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Mon, 20 Jul 2026 02:53:56 +0000
Subject: [PATCH 14/44] Fix IW-OPD paper metadata used in the model card

The model card built by IWOPDTrainer attributed runs to the OPD paper
(2306.13649) carried over from the fork source. Point _paper at the
IW-OPD paper (2606.22600, "On the Position Bias of On-Policy
Distillation") so published checkpoints cite the correct work.
---
 trl/experimental/iw_opd/iw_opd_trainer.py | 16 +++++++---------
 1 file changed, 7 insertions(+), 9 deletions(-)

diff --git a/trl/experimental/iw_opd/iw_opd_trainer.py b/trl/experimental/iw_opd/iw_opd_trainer.py
index 685defce7a9..b884b49b7d8 100644
--- a/trl/experimental/iw_opd/iw_opd_trainer.py
+++ b/trl/experimental/iw_opd/iw_opd_trainer.py
@@ -378,17 +378,15 @@ class IWOPDTrainer(_BaseTrainer):
     _tag_names = ["trl", "iw-opd"]
     _name = "IW-OPD"
     _paper = {
-        "title": "On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes",
-        "id": "2306.13649",
+        "title": "On the Position Bias of On-Policy Distillation",
+        "id": "2606.22600",
         # docstyle-ignore
         "citation": textwrap.dedent("""\
-            @inproceedings{agarwal2024on-policy,
-                title        = {{On-Policy Distillation of Language Models: Learning from Self-Generated Mistakes}},
-                author       = {Rishabh Agarwal and Nino Vieillard and Yongchao Zhou and Piotr Stanczyk and Sabela Ramos Garea and Matthieu Geist and Olivier Bachem},
-                year         = 2024,
-                booktitle    = {The Twelfth International Conference on Learning Representations, {ICLR} 2024, Vienna, Austria, May 7-11, 2024},
-                publisher    = {OpenReview.net},
-                url          = {https://openreview.net/forum?id=3zKtaqxLhW},
+            @article{xie2026position,
+                title        = {{On the Position Bias of On-Policy Distillation}},
+                author       = {Yan Xie and Sijie Zhu and Tiansheng Wen and Bo Chen and Yifei Wang},
+                year         = 2026,
+                eprint       = {arXiv:2606.22600},
             }"""),
     }
 

From 26038ddee73f4c1b7b47e1043a20f56bd18b7bb8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Mon, 20 Jul 2026 03:21:27 +0000
Subject: [PATCH 15/44] Remove stale on/off-policy comments after the
 off-policy removal

Drop the now-empty 'Loss tracking' banner and retitle the buffering
section header, both left over from removing the on/off-policy path.
---
 trl/experimental/distillation/distillation_trainer.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/trl/experimental/distillation/distillation_trainer.py b/trl/experimental/distillation/distillation_trainer.py
index fd06f8c45a7..fe9575d0f5f 100644
--- a/trl/experimental/distillation/distillation_trainer.py
+++ b/trl/experimental/distillation/distillation_trainer.py
@@ -501,8 +501,6 @@ def __init__(
         self._buffered_text_logs = None
         self._buffer_step = 0
 
-        # ── Loss tracking ──
-
         # ── Generation config ──
         generation_kwargs = {
             "max_new_tokens": args.max_completion_length,
@@ -671,7 +669,7 @@ def get_train_dataloader(self):
         return _RepeatBatchDataLoader(base_dataloader, repeat_count=self.args.gradient_accumulation_steps)
 
     # ──────────────────────────────────────────────────────────────────────
-    #  Buffering: on/off-policy mixing across gradient accumulation steps
+    #  Buffering across gradient accumulation steps
     # ──────────────────────────────────────────────────────────────────────
 
     @profiling_decorator

From 3515b22e14c352ab36df980437ff03a0031d131f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Mon, 20 Jul 2026 03:41:03 +0000
Subject: [PATCH 16/44] Accept a `prompt` column alongside `messages`

Extend the distillation collator to accept the forward-looking prompt-only
dataset format (a `prompt` column of conversational messages, as used by GRPO)
in addition to the existing `messages` language-modeling format. A prompt-only
example carries no completion, so the student generates one on-policy; the batch
is tokenized prompt-only with every label masked.

Register `prompt` in the trainer's signature columns.

Tests: the collator yields a fully-masked batch for a prompt-only dataset;
end-to-end prompt-only training is added as xfail (it hits the known
num_items_in_batch=0 -> NaN denominator bug, un-xfailed at plan 5.6).
---
 .../experimental/test_distillation_trainer.py | 43 +++++++++++++++++++
 .../distillation/distillation_trainer.py      | 27 ++++++++----
 2 files changed, 62 insertions(+), 8 deletions(-)

diff --git a/tests/experimental/test_distillation_trainer.py b/tests/experimental/test_distillation_trainer.py
index 78af1c861a7..769374fca66 100644
--- a/tests/experimental/test_distillation_trainer.py
+++ b/tests/experimental/test_distillation_trainer.py
@@ -266,6 +266,49 @@ 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."
 
+    def test_collator_accepts_prompt_only_dataset(self):
+        """A `prompt` column carries no completion, so the collated batch is entirely masked (nothing to train on)."""
+        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(),
+            train_dataset=dataset,
+            processing_class=self.tokenizer,
+        )
+        batch = trainer.data_collator([dataset[i] for i in range(2)])
+        assert (batch["labels"] == -100).all()
+        assert torch.equal(batch["input_ids"], batch["prompts"])
+
+    @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."""
+        losses = []
+
+        class _RecordingTrainer(DistillationTrainer):
+            def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
+                out = super().compute_loss(model, inputs, return_outputs, num_items_in_batch)
+                losses.append(out[0] if isinstance(out, tuple) else out)
+                return out
+
+        dataset = load_dataset("trl-internal-testing/zen", "conversational_prompt_only", split="train")
+        trainer = _RecordingTrainer(
+            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(loss).item() for loss in losses)
+
     @pytest.mark.xfail(
         reason="num_items_in_batch is computed from the raw dataloader batches, before _prepare_inputs runs. "
         "_RepeatBatchDataLoader yields the same generation batch once per accumulation step, so the count is "
diff --git a/trl/experimental/distillation/distillation_trainer.py b/trl/experimental/distillation/distillation_trainer.py
index fe9575d0f5f..37b2ae3d3d0 100644
--- a/trl/experimental/distillation/distillation_trainer.py
+++ b/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:

From 4e5d06ffbf8ad7e2c6c6befe0ad7d3411c482a85 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Mon, 20 Jul 2026 03:58:59 +0000
Subject: [PATCH 17/44] Accept a `prompt` column alongside `messages`
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

Extend the distillation collator to accept the forward-looking prompt-only
dataset format (a `prompt` column of conversational messages, as used by GRPO)
in addition to the existing `messages` language-modeling format. A prompt-only
example carries no completion, so the student generates one on-policy; the batch
is tokenized prompt-only with every label masked.

Register `prompt` in the trainer's signature columns.

Test: end-to-end prompt-only training, added as xfail — it hits the known
num_items_in_batch=0 -> NaN denominator bug (un-xfailed at plan 5.6).
---
 .../experimental/test_distillation_trainer.py | 21 +++++++++++++++++++
 1 file changed, 21 insertions(+)

diff --git a/tests/experimental/test_distillation_trainer.py b/tests/experimental/test_distillation_trainer.py
index 769374fca66..c580d1ace2e 100644
--- a/tests/experimental/test_distillation_trainer.py
+++ b/tests/experimental/test_distillation_trainer.py
@@ -309,6 +309,27 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
 
         assert all(torch.isfinite(loss).item() for loss in losses)
 
+    @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="num_items_in_batch is computed from the raw dataloader batches, before _prepare_inputs runs. "
         "_RepeatBatchDataLoader yields the same generation batch once per accumulation step, so the count is "

From dfbf0e5d897684b7330baa25d198ae7a355ffad2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Mon, 20 Jul 2026 04:01:03 +0000
Subject: [PATCH 18/44] style

---
 .../experimental/test_distillation_trainer.py | 43 -------------------
 .../distillation/distillation_trainer.py      |  4 +-
 2 files changed, 2 insertions(+), 45 deletions(-)

diff --git a/tests/experimental/test_distillation_trainer.py b/tests/experimental/test_distillation_trainer.py
index c580d1ace2e..a409424222a 100644
--- a/tests/experimental/test_distillation_trainer.py
+++ b/tests/experimental/test_distillation_trainer.py
@@ -266,49 +266,6 @@ 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."
 
-    def test_collator_accepts_prompt_only_dataset(self):
-        """A `prompt` column carries no completion, so the collated batch is entirely masked (nothing to train on)."""
-        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(),
-            train_dataset=dataset,
-            processing_class=self.tokenizer,
-        )
-        batch = trainer.data_collator([dataset[i] for i in range(2)])
-        assert (batch["labels"] == -100).all()
-        assert torch.equal(batch["input_ids"], batch["prompts"])
-
-    @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."""
-        losses = []
-
-        class _RecordingTrainer(DistillationTrainer):
-            def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
-                out = super().compute_loss(model, inputs, return_outputs, num_items_in_batch)
-                losses.append(out[0] if isinstance(out, tuple) else out)
-                return out
-
-        dataset = load_dataset("trl-internal-testing/zen", "conversational_prompt_only", split="train")
-        trainer = _RecordingTrainer(
-            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(loss).item() for loss in losses)
-
     @pytest.mark.xfail(
         strict=True,
         reason="Prompt-only datasets carry no dataset-completion tokens, so num_items_in_batch (counted from the raw "
diff --git a/trl/experimental/distillation/distillation_trainer.py b/trl/experimental/distillation/distillation_trainer.py
index 37b2ae3d3d0..c0e9e64e811 100644
--- a/trl/experimental/distillation/distillation_trainer.py
+++ b/trl/experimental/distillation/distillation_trainer.py
@@ -152,8 +152,8 @@ class _DistillationCollator:
 
     - 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.
+    - 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.

From 89494f2ba92d4c6bfc03dc7e17637d8f16094f39 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Mon, 20 Jul 2026 04:50:11 +0000
Subject: [PATCH 19/44] Relax self-distillation train-loss assertion to
 tolerate fp noise

With the local sparse path removed, the full-vocabulary loss is the default. For
self-distillation (teacher == student) the divergence is ~0 and rounds a hair
below zero (~1e-12), so `>= 0.0` is too strict. Assert `>= -1e-4` instead, which
still catches a genuinely negative loss.
---
 tests/experimental/test_distillation_trainer.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/tests/experimental/test_distillation_trainer.py b/tests/experimental/test_distillation_trainer.py
index 674bb70ebff..c79379e4225 100644
--- a/tests/experimental/test_distillation_trainer.py
+++ b/tests/experimental/test_distillation_trainer.py
@@ -249,7 +249,9 @@ def test_distillation_trainer_train_runs_with_local_teacher(self):
 
         assert trainer.state.log_history[-1]["train_loss"] is not None
         assert trainer.state.log_history[0]["eval_loss"] is not None
-        assert train_result.metrics["train_loss"] >= 0.0
+        # Self-distillation (teacher == student), so the divergence is ~0; allow tiny floating-point noise below zero
+        # while still catching a genuinely negative loss.
+        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])

From 9b5ad813afe57bd051a2c70b6ff98b01862f3efe Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Mon, 20 Jul 2026 04:53:42 +0000
Subject: [PATCH 20/44] Fix server-config construction to only copy dataclass
 init fields

The test built ServerDistillationConfig by splatting vars() of a DistillationConfig,
which also carries post-init-derived attributes (mixed_precision, distributed_state,
...) that __init__ rejects. Restrict the copy to actual init fields.
---
 tests/experimental/test_server_distillation_trainer.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/tests/experimental/test_server_distillation_trainer.py b/tests/experimental/test_server_distillation_trainer.py
index 4334865d9d0..699b6e1938d 100644
--- a/tests/experimental/test_server_distillation_trainer.py
+++ b/tests/experimental/test_server_distillation_trainer.py
@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import dataclasses
 import math
 from unittest.mock import MagicMock
 
@@ -342,8 +343,11 @@ def _make_local_trainer(self, **kwargs):
     def _make_server_trainer(self, **kwargs):
         dataset = load_dataset("trl-internal-testing/zen", "conversational_language_modeling", split="train")
         args = self._local_args(**kwargs)
+        # Copy only the dataclass init fields: vars() also carries post-init-derived attributes
+        # (mixed_precision, distributed_state, ...) that are not accepted by __init__.
+        init_fields = {f.name for f in dataclasses.fields(ServerDistillationConfig) if f.init}
         server_args = ServerDistillationConfig(
-            **{k: v for k, v in vars(args).items() if not k.startswith("_")},
+            **{k: v for k, v in vars(args).items() if k in init_fields},
             teacher_model_server_url="http://localhost:8000",
         )
         return ServerDistillationTrainer(

From c978997a6f1a2152f3919063166334d4f060bdda Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Mon, 20 Jul 2026 16:03:27 +0000
Subject: [PATCH 21/44] 0.5 -> 0.25

---
 tests/experimental/test_distillation_trainer.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/experimental/test_distillation_trainer.py b/tests/experimental/test_distillation_trainer.py
index 7642b112773..3dca332d771 100644
--- a/tests/experimental/test_distillation_trainer.py
+++ b/tests/experimental/test_distillation_trainer.py
@@ -408,7 +408,7 @@ def setup_method(self):
         self.teacher_logits = torch.randn(2, 3, 5, generator=generator)
         self.labels = torch.tensor([[-100, 1, 2], [-100, -100, 3]])
 
-    @pytest.mark.parametrize("beta", [0.0, 0.5, 1.0])
+    @pytest.mark.parametrize("beta", [0.0, 0.25, 1.0])
     @pytest.mark.parametrize("use_labels", [False, True])
     def test_matches_reference_implementation(self, beta, use_labels):
         labels = self.labels if use_labels else None
@@ -418,7 +418,7 @@ def test_matches_reference_implementation(self, beta, use_labels):
         expected = _reference_generalized_jsd(self.student_logits, self.teacher_logits, labels=labels, beta=beta)
         torch.testing.assert_close(loss, expected)
 
-    @pytest.mark.parametrize("beta", [0.0, 0.5, 1.0])
+    @pytest.mark.parametrize("beta", [0.0, 0.25, 1.0])
     @pytest.mark.parametrize("use_labels", [False, True])
     def test_matches_gkd(self, beta, use_labels):
         # GKD implements the same objective. Keeping the two in lockstep is the cross-trainer contract: if this breaks,
@@ -430,7 +430,7 @@ def test_matches_gkd(self, beta, use_labels):
         gkd_loss = GKDTrainer.generalized_jsd_loss(self.student_logits, self.teacher_logits, labels=labels, beta=beta)
         torch.testing.assert_close(loss, gkd_loss)
 
-    @pytest.mark.parametrize("beta", [0.0, 0.5, 1.0])
+    @pytest.mark.parametrize("beta", [0.0, 0.25, 1.0])
     def test_temperature_matches_reference(self, beta):
         # `temperature` is applied to the loss today. It is scheduled to become sampling-only, so pin it explicitly:
         # that change must be a visible diff here, not a silent drift.

From b9f99f13a3f673ea0c41e864b7bdfb197130ea33 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Mon, 20 Jul 2026 16:18:58 +0000
Subject: [PATCH 22/44] Diverge the teacher in the on/off-policy param-update
 pin

With teacher == student the divergence is ~0, so the every-parameter-changed
assertion could pass on floating-point noise. Perturb the teacher first (as
test_loss_normalizes_by_num_items_in_batch already does) so the training signal
is real.
---
 tests/experimental/test_distillation_trainer.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/tests/experimental/test_distillation_trainer.py b/tests/experimental/test_distillation_trainer.py
index 8039cbdd735..deb5dead653 100644
--- a/tests/experimental/test_distillation_trainer.py
+++ b/tests/experimental/test_distillation_trainer.py
@@ -614,6 +614,14 @@ def test_train_updates_params_on_and_off_policy(self, lmbda):
         # 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)
+
+        # 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.
+        torch.manual_seed(0)
+        with torch.no_grad():
+            for p in trainer.teacher_model.parameters():
+                p.add_(0.5 * torch.randn_like(p))
+
         previous_params = {name: param.clone() for name, param in trainer.model.named_parameters()}
 
         trainer.train()

From 26eba3ca878bb1a5d5cf378b3700baf7e476c2e6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Mon, 20 Jul 2026 17:01:27 +0000
Subject: [PATCH 23/44] remove unused import

---
 tests/experimental/test_server_distillation_trainer.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tests/experimental/test_server_distillation_trainer.py b/tests/experimental/test_server_distillation_trainer.py
index be9b60e5442..b85f7c689fd 100644
--- a/tests/experimental/test_server_distillation_trainer.py
+++ b/tests/experimental/test_server_distillation_trainer.py
@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import dataclasses
 import math
 from unittest.mock import MagicMock
 

From 5e05b9df46bf31622677bd1b1e57e6925aa8b10d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Mon, 20 Jul 2026 18:25:35 +0000
Subject: [PATCH 24/44] Rework the num_items_in_batch xfail: applied
 denominator, on-policy

---
 .../experimental/test_distillation_trainer.py | 45 ++++++++++---------
 1 file changed, 24 insertions(+), 21 deletions(-)

diff --git a/tests/experimental/test_distillation_trainer.py b/tests/experimental/test_distillation_trainer.py
index cfa46668b60..e64a94c9e54 100644
--- a/tests/experimental/test_distillation_trainer.py
+++ b/tests/experimental/test_distillation_trainer.py
@@ -631,41 +631,44 @@ def test_train_updates_params_on_and_off_policy(self, lmbda):
             assert not torch.equal(param, trainer.model.get_parameter(name)), f"Parameter {name} has not changed."
 
     @pytest.mark.xfail(
-        reason="num_items_in_batch is computed from the raw dataloader batches, before _prepare_inputs runs. "
-        "_RepeatBatchDataLoader yields the same generation batch once per accumulation step, so the count is "
-        "gradient_accumulation_steps times too large, and it is derived from the dataset labels even on steps whose "
-        "completions are replaced by generation. Fixed when the buffer moves to the GRPO-style _prepare_inputs."
+        strict=True,
+        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 "
+        "accumulation window, so the denominator the loss divides by does not equal the completion tokens actually "
+        "trained on. Un-xfail when the count moves to the GRPO-style _prepare_inputs (plan 5.6).",
     )
-    def test_num_items_in_batch_counts_the_tokens_trained_on(self):
-        """`num_items_in_batch` is the loss denominator, so it must count the tokens actually trained on.
+    def test_num_items_in_batch_counts_the_tokens_trained_on(self, monkeypatch):
+        """`num_items_in_batch` is the loss denominator, so it must count the completion tokens actually trained on.
 
-        The loss is reduced as `sum / num_items_in_batch`, so an inflated count silently scales the loss — and the
-        gradient — down. Documented here as a known failure; un-xfail it in the PR that fixes the denominator.
+        Capture the value where it is applied (`_reduce_divergence_loss`) rather than the argument transformers passes
+        to `compute_loss`: the GRPO-style fix computes the count during generation and the loss reads it from there, so
+        asserting on the applied denominator keeps the test valid — and able to turn green — across that move.
         """
-        recorded = []
+        recorded = []  # (denominator applied, completion tokens in this microbatch)
+        original = DistillationTrainer._reduce_divergence_loss
 
-        class _RecordingTrainer(DistillationTrainer):
-            def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
-                # Prompt positions are already -100, so this is the completion-token count either way.
-                recorded.append((num_items_in_batch, int((inputs["labels"] != -100).sum())))
-                return super().compute_loss(model, inputs, return_outputs, num_items_in_batch)
+        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)
+
+        monkeypatch.setattr(DistillationTrainer, "_reduce_divergence_loss", staticmethod(_recording))
 
         dataset = load_dataset("trl-internal-testing/zen", "conversational_language_modeling", split="train")
-        trainer = _RecordingTrainer(
+        trainer = DistillationTrainer(
             model=self.model_id,
             teacher_model=self.model_id,
-            args=self._make_args(gradient_accumulation_steps=2, max_steps=1),
+            args=self._make_args(lmbda=1.0, gradient_accumulation_steps=2, max_steps=1),
             train_dataset=dataset,
             processing_class=self.tokenizer,
         )
 
         trainer.train()
 
-        assert len(recorded) == 2, "expected one compute_loss call per accumulation step"
-        reported = recorded[0][0]
-        assert reported is not None, "transformers did not pass num_items_in_batch"
-        # The denominator should be the number of tokens summed over the accumulation window.
-        assert int(reported) == sum(trained_on for _, trained_on in recorded)
+        assert len(recorded) == 2, "expected one loss reduction per accumulation step"
+        denominator = recorded[0][0]
+        assert denominator is not None, "the loss was not reduced by a token count"
+        # The denominator must be the completion tokens summed over the whole accumulation window.
+        assert int(denominator) == sum(tokens for _, tokens in recorded)
 
     @pytest.mark.parametrize(
         "eval_dataset_type",

From 23784c3a05429db535407257be033c40cb410d2b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Mon, 20 Jul 2026 18:29:22 +0000
Subject: [PATCH 25/44] Use a non-strict xfail for the num_items_in_batch pin

Align with the rest of the suite: all 51 xfail markers in tests/ are non-strict.
strict=True was an outlier that would fail CI on an unexpected XPASS.
---
 tests/experimental/test_distillation_trainer.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tests/experimental/test_distillation_trainer.py b/tests/experimental/test_distillation_trainer.py
index e64a94c9e54..9b88ddf1ce1 100644
--- a/tests/experimental/test_distillation_trainer.py
+++ b/tests/experimental/test_distillation_trainer.py
@@ -631,7 +631,6 @@ def test_train_updates_params_on_and_off_policy(self, lmbda):
             assert not torch.equal(param, trainer.model.get_parameter(name)), f"Parameter {name} has not changed."
 
     @pytest.mark.xfail(
-        strict=True,
         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 "
         "accumulation window, so the denominator the loss divides by does not equal the completion tokens actually "

From 6a2be124dd390a4fc399e1d424e7b3d99c00c983 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Mon, 20 Jul 2026 19:55:24 +0000
Subject: [PATCH 26/44] Deprecate `messages`-format datasets

Now that the collator accepts a prompt-only `prompt` column, warn (once) when a
`messages`-format conversational language-modeling dataset is passed, pointing at
the prompt-only format. Deprecate-before-remove; the actual removal of `messages`
support lands later in Phase 3.
---
 tests/experimental/test_distillation_trainer.py  | 16 +++++++++++++++-
 .../distillation/distillation_trainer.py         | 10 ++++++++++
 2 files changed, 25 insertions(+), 1 deletion(-)

diff --git a/tests/experimental/test_distillation_trainer.py b/tests/experimental/test_distillation_trainer.py
index 63d6dd02d04..7b774028b06 100644
--- a/tests/experimental/test_distillation_trainer.py
+++ b/tests/experimental/test_distillation_trainer.py
@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import os
+import warnings
 
 import pytest
 import torch
@@ -21,7 +22,7 @@
 from transformers import AutoModelForCausalLM, AutoTokenizer
 
 from trl.experimental.distillation import DistillationConfig, DistillationTrainer
-from trl.experimental.distillation.distillation_trainer import _RepeatBatchDataLoader
+from trl.experimental.distillation.distillation_trainer import _DistillationCollator, _RepeatBatchDataLoader
 from trl.experimental.gkd.gkd_trainer import GKDTrainer
 
 from ..testing_utils import TrlTestCase, require_liger_kernel, require_torch_accelerator
@@ -187,6 +188,19 @@ 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,
diff --git a/trl/experimental/distillation/distillation_trainer.py b/trl/experimental/distillation/distillation_trainer.py
index c0e9e64e811..60c2f20747d 100644
--- a/trl/experimental/distillation/distillation_trainer.py
+++ b/trl/experimental/distillation/distillation_trainer.py
@@ -14,6 +14,7 @@
 
 import random
 import textwrap
+import warnings
 from collections import defaultdict
 from collections.abc import Callable
 from contextlib import nullcontext
@@ -172,6 +173,7 @@ def __init__(
         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.")
@@ -188,6 +190,14 @@ def __call__(self, examples: list[dict[str, Any]]) -> dict[str, torch.Tensor]:
                 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"

From 93f6f5552a0f6cd94791cd0b8b178864c7dbc0bd Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Mon, 20 Jul 2026 21:01:11 +0000
Subject: [PATCH 27/44] Drop the loss_top_k usage-tips bullet from the docs

loss_top_k was removed from DistillationConfig (it now lives on
ServerDistillationConfig), so the local usage guide no longer documents it.
The "three key parameters" count drops to two (lmbda, beta).
---
 docs/source/distillation_trainer.md | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/docs/source/distillation_trainer.md b/docs/source/distillation_trainer.md
index c23e291a93d..7cab1d23628 100644
--- a/docs/source/distillation_trainer.md
+++ b/docs/source/distillation_trainer.md
@@ -54,11 +54,10 @@ trainer.save_model()
 
 ## Usage tips
 
-The [`experimental.distillation.DistillationTrainer`] needs three key parameters set via [`experimental.distillation.DistillationConfig`]:
+The [`experimental.distillation.DistillationTrainer`] needs two key parameters set via [`experimental.distillation.DistillationConfig`]:
 
 * `lmbda`: controls the student data fraction, i.e., the proportion of on-policy student-generated outputs. When `lmbda=0.0`, training is fully off-policy (dataset completions only). When `lmbda=1.0`, training is fully on-policy (student generates all completions). For values in between, each gradient accumulation slice is randomly assigned as on- or off-policy based on `lmbda`.
 * `beta`: controls the interpolation in the Generalized Jensen-Shannon Divergence. When `beta=0.0` the loss approximates forward KL divergence, while `beta=1.0` approximates reverse KL divergence. Values in between interpolate.
-* `loss_top_k`: number of top tokens to use for the KL/JSD loss. Set to `0` for exact full-vocabulary computation (local teacher only), or `> 0` for a top-k approximation. See more about top-k with external teacher server below.
 
 ### On-policy vs. off-policy
 

From ddcb836a4d885cfcbee190f583eb55a190d06320 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Mon, 20 Jul 2026 21:42:02 +0000
Subject: [PATCH 28/44] Drop leftover lmbda=1.0 from the num_items xfail test
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

lmbda was removed from DistillationConfig in this PR, so the xfailed
num_items test was raising TypeError at construction — swallowed by the
xfail marker, so it never exercised the num_items path. The trainer is
always on-policy now, so the arg is unnecessary.
---
 tests/experimental/test_distillation_trainer.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/experimental/test_distillation_trainer.py b/tests/experimental/test_distillation_trainer.py
index de9b02a23e5..a98eeff3a7d 100644
--- a/tests/experimental/test_distillation_trainer.py
+++ b/tests/experimental/test_distillation_trainer.py
@@ -300,7 +300,7 @@ def _recording(jsd, labels=None, reduction="batchmean", num_items_in_batch=None)
         trainer = DistillationTrainer(
             model=self.model_id,
             teacher_model=self.model_id,
-            args=self._make_args(lmbda=1.0, gradient_accumulation_steps=2, max_steps=1),
+            args=self._make_args(gradient_accumulation_steps=2, max_steps=1),
             train_dataset=dataset,
             processing_class=self.tokenizer,
         )

From 46d1c267cc6436c689a2c92f195b7a03a8f97774 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Tue, 21 Jul 2026 00:36:00 +0000
Subject: [PATCH 29/44] Fix num_items_in_batch to count generated completion
 tokens

Changes loss values.

transformers derives num_items_in_batch from the raw dataloader labels, before
on-policy generation replaces the completions. That count is wrong for on-policy
training (it counts dataset completions, not the generated ones) and zero for
prompt-only datasets, where it made the loss sum/0 = NaN.

_fill_buffer already generates every accumulation-window slice upfront, so the
global generated-completion-token count is known before any loss is computed.
Compute it there (gathered across processes) and use it as the loss denominator
in compute_loss. Un-xfail the two num_items tests (items 03 and 12).
---
 tests/experimental/test_distillation_trainer.py | 14 +-------------
 .../distillation/distillation_trainer.py        | 17 +++++++++++++++++
 2 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/tests/experimental/test_distillation_trainer.py b/tests/experimental/test_distillation_trainer.py
index b68ae6236c3..8705acf9a63 100644
--- a/tests/experimental/test_distillation_trainer.py
+++ b/tests/experimental/test_distillation_trainer.py
@@ -288,14 +288,8 @@ 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."""
+        """The forward-looking prompt-only format trains end to end: the student generates, the teacher scores."""
         dataset = load_dataset("trl-internal-testing/zen", "conversational_prompt_only", split="train")
         trainer = DistillationTrainer(
             model=self.model_id,
@@ -309,12 +303,6 @@ def test_train_runs_with_prompt_only_dataset(self):
 
         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 "
-        "accumulation window, so the denominator the loss divides by does not equal the completion tokens actually "
-        "trained on. Un-xfail when the count moves to the GRPO-style _prepare_inputs (plan 5.6).",
-    )
     def test_num_items_in_batch_counts_the_tokens_trained_on(self, monkeypatch):
         """`num_items_in_batch` is the loss denominator, so it must count the completion tokens actually trained on.
 
diff --git a/trl/experimental/distillation/distillation_trainer.py b/trl/experimental/distillation/distillation_trainer.py
index 60c2f20747d..d77c9c2d257 100644
--- a/trl/experimental/distillation/distillation_trainer.py
+++ b/trl/experimental/distillation/distillation_trainer.py
@@ -520,6 +520,7 @@ def __init__(
         # ── Buffer state ──
         self._buffered_inputs = None
         self._buffered_text_logs = None
+        self._buffered_num_items = None
         self._buffer_step = 0
 
         # ── Generation config ──
@@ -718,6 +719,17 @@ def _fill_buffer(self, generation_batch: dict[str, torch.Tensor | Any], buffer_s
         # Generate student completions for every slice
         self._generate_student_completions(slices, list(range(buffer_steps)))
 
+        # Loss denominator (`num_items_in_batch`): the global number of completion tokens actually trained on this
+        # optimizer step. transformers derives its own count from the *raw dataloader* labels — before generation
+        # replaces the completions — which is wrong for on-policy training and zero for prompt-only datasets (dividing
+        # the loss by zero). Recompute it here from the generated labels, gathered across processes (issue #4719).
+        local_completion_tokens = sum(
+            int((s["labels"] != -100).sum()) for s in self._buffered_inputs if s is not None
+        )
+        self._buffered_num_items = self.accelerator.gather(
+            torch.tensor(local_completion_tokens, device=self.accelerator.device)
+        ).sum()
+
         # Gather text logs once per optimizer step (all processes must participate)
         if self.log_completions:
             prompts_all = []
@@ -1013,6 +1025,11 @@ def _get_teacher_logits(self, inputs: dict[str, torch.Tensor | Any]) -> torch.Te
             ).logits
 
     def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
+        # transformers computes `num_items_in_batch` from the raw dataloader labels, before on-policy generation
+        # replaces the completions; use the count over the generated completions instead (computed in `_fill_buffer`).
+        if self.model.training and self._buffered_num_items is not None:
+            num_items_in_batch = self._buffered_num_items
+
         if self.use_liger_loss:
             loss = self._compute_liger_loss(model, inputs, num_items_in_batch=num_items_in_batch)
             return (loss, None) if return_outputs else loss

From 010f25094dfc1152f822f1d6c7ce8fd589632fde Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Tue, 21 Jul 2026 00:42:53 +0000
Subject: [PATCH 30/44] Revert "Fix num_items_in_batch to count generated
 completion tokens"

This reverts commit 46d1c267cc6436c689a2c92f195b7a03a8f97774.
---
 tests/experimental/test_distillation_trainer.py | 14 +++++++++++++-
 .../distillation/distillation_trainer.py        | 17 -----------------
 2 files changed, 13 insertions(+), 18 deletions(-)

diff --git a/tests/experimental/test_distillation_trainer.py b/tests/experimental/test_distillation_trainer.py
index 8705acf9a63..b68ae6236c3 100644
--- a/tests/experimental/test_distillation_trainer.py
+++ b/tests/experimental/test_distillation_trainer.py
@@ -288,8 +288,14 @@ 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 trains end to end: the student generates, the teacher scores."""
+        """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,
@@ -303,6 +309,12 @@ def test_train_runs_with_prompt_only_dataset(self):
 
         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 "
+        "accumulation window, so the denominator the loss divides by does not equal the completion tokens actually "
+        "trained on. Un-xfail when the count moves to the GRPO-style _prepare_inputs (plan 5.6).",
+    )
     def test_num_items_in_batch_counts_the_tokens_trained_on(self, monkeypatch):
         """`num_items_in_batch` is the loss denominator, so it must count the completion tokens actually trained on.
 
diff --git a/trl/experimental/distillation/distillation_trainer.py b/trl/experimental/distillation/distillation_trainer.py
index d77c9c2d257..60c2f20747d 100644
--- a/trl/experimental/distillation/distillation_trainer.py
+++ b/trl/experimental/distillation/distillation_trainer.py
@@ -520,7 +520,6 @@ def __init__(
         # ── Buffer state ──
         self._buffered_inputs = None
         self._buffered_text_logs = None
-        self._buffered_num_items = None
         self._buffer_step = 0
 
         # ── Generation config ──
@@ -719,17 +718,6 @@ def _fill_buffer(self, generation_batch: dict[str, torch.Tensor | Any], buffer_s
         # Generate student completions for every slice
         self._generate_student_completions(slices, list(range(buffer_steps)))
 
-        # Loss denominator (`num_items_in_batch`): the global number of completion tokens actually trained on this
-        # optimizer step. transformers derives its own count from the *raw dataloader* labels — before generation
-        # replaces the completions — which is wrong for on-policy training and zero for prompt-only datasets (dividing
-        # the loss by zero). Recompute it here from the generated labels, gathered across processes (issue #4719).
-        local_completion_tokens = sum(
-            int((s["labels"] != -100).sum()) for s in self._buffered_inputs if s is not None
-        )
-        self._buffered_num_items = self.accelerator.gather(
-            torch.tensor(local_completion_tokens, device=self.accelerator.device)
-        ).sum()
-
         # Gather text logs once per optimizer step (all processes must participate)
         if self.log_completions:
             prompts_all = []
@@ -1025,11 +1013,6 @@ def _get_teacher_logits(self, inputs: dict[str, torch.Tensor | Any]) -> torch.Te
             ).logits
 
     def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
-        # transformers computes `num_items_in_batch` from the raw dataloader labels, before on-policy generation
-        # replaces the completions; use the count over the generated completions instead (computed in `_fill_buffer`).
-        if self.model.training and self._buffered_num_items is not None:
-            num_items_in_batch = self._buffered_num_items
-
         if self.use_liger_loss:
             loss = self._compute_liger_loss(model, inputs, num_items_in_batch=num_items_in_batch)
             return (loss, None) if return_outputs else loss

From 6c60211450802268e18d5df868be3ac9a695dfb3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Tue, 21 Jul 2026 00:36:00 +0000
Subject: [PATCH 31/44] Fix num_items_in_batch to count generated completion
 tokens

Changes loss values.

transformers derives num_items_in_batch from the raw dataloader labels, before
on-policy generation replaces the completions. That count is wrong for on-policy
training (it counts dataset completions, not the generated ones) and zero for
prompt-only datasets, where it made the loss sum/0 = NaN.

_fill_buffer already generates every accumulation-window slice upfront, so the
global generated-completion-token count is known before any loss is computed.
Compute it there (gathered across processes) and use it as the loss denominator
in compute_loss. Un-xfail the two num_items tests (items 03 and 12).
---
 tests/experimental/test_distillation_trainer.py | 14 +-------------
 .../distillation/distillation_trainer.py        | 17 +++++++++++++++++
 2 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/tests/experimental/test_distillation_trainer.py b/tests/experimental/test_distillation_trainer.py
index b68ae6236c3..8705acf9a63 100644
--- a/tests/experimental/test_distillation_trainer.py
+++ b/tests/experimental/test_distillation_trainer.py
@@ -288,14 +288,8 @@ 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."""
+        """The forward-looking prompt-only format trains end to end: the student generates, the teacher scores."""
         dataset = load_dataset("trl-internal-testing/zen", "conversational_prompt_only", split="train")
         trainer = DistillationTrainer(
             model=self.model_id,
@@ -309,12 +303,6 @@ def test_train_runs_with_prompt_only_dataset(self):
 
         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 "
-        "accumulation window, so the denominator the loss divides by does not equal the completion tokens actually "
-        "trained on. Un-xfail when the count moves to the GRPO-style _prepare_inputs (plan 5.6).",
-    )
     def test_num_items_in_batch_counts_the_tokens_trained_on(self, monkeypatch):
         """`num_items_in_batch` is the loss denominator, so it must count the completion tokens actually trained on.
 
diff --git a/trl/experimental/distillation/distillation_trainer.py b/trl/experimental/distillation/distillation_trainer.py
index 60c2f20747d..d77c9c2d257 100644
--- a/trl/experimental/distillation/distillation_trainer.py
+++ b/trl/experimental/distillation/distillation_trainer.py
@@ -520,6 +520,7 @@ def __init__(
         # ── Buffer state ──
         self._buffered_inputs = None
         self._buffered_text_logs = None
+        self._buffered_num_items = None
         self._buffer_step = 0
 
         # ── Generation config ──
@@ -718,6 +719,17 @@ def _fill_buffer(self, generation_batch: dict[str, torch.Tensor | Any], buffer_s
         # Generate student completions for every slice
         self._generate_student_completions(slices, list(range(buffer_steps)))
 
+        # Loss denominator (`num_items_in_batch`): the global number of completion tokens actually trained on this
+        # optimizer step. transformers derives its own count from the *raw dataloader* labels — before generation
+        # replaces the completions — which is wrong for on-policy training and zero for prompt-only datasets (dividing
+        # the loss by zero). Recompute it here from the generated labels, gathered across processes (issue #4719).
+        local_completion_tokens = sum(
+            int((s["labels"] != -100).sum()) for s in self._buffered_inputs if s is not None
+        )
+        self._buffered_num_items = self.accelerator.gather(
+            torch.tensor(local_completion_tokens, device=self.accelerator.device)
+        ).sum()
+
         # Gather text logs once per optimizer step (all processes must participate)
         if self.log_completions:
             prompts_all = []
@@ -1013,6 +1025,11 @@ def _get_teacher_logits(self, inputs: dict[str, torch.Tensor | Any]) -> torch.Te
             ).logits
 
     def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
+        # transformers computes `num_items_in_batch` from the raw dataloader labels, before on-policy generation
+        # replaces the completions; use the count over the generated completions instead (computed in `_fill_buffer`).
+        if self.model.training and self._buffered_num_items is not None:
+            num_items_in_batch = self._buffered_num_items
+
         if self.use_liger_loss:
             loss = self._compute_liger_loss(model, inputs, num_items_in_batch=num_items_in_batch)
             return (loss, None) if return_outputs else loss

From fbf4c3bc24926332fdc17a8a9e7a7c7fd32bfaf6 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Tue, 21 Jul 2026 00:51:20 +0000
Subject: [PATCH 32/44] style

---
 trl/experimental/distillation/distillation_trainer.py | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/trl/experimental/distillation/distillation_trainer.py b/trl/experimental/distillation/distillation_trainer.py
index d77c9c2d257..9c8d5145b3d 100644
--- a/trl/experimental/distillation/distillation_trainer.py
+++ b/trl/experimental/distillation/distillation_trainer.py
@@ -723,9 +723,7 @@ def _fill_buffer(self, generation_batch: dict[str, torch.Tensor | Any], buffer_s
         # optimizer step. transformers derives its own count from the *raw dataloader* labels — before generation
         # replaces the completions — which is wrong for on-policy training and zero for prompt-only datasets (dividing
         # the loss by zero). Recompute it here from the generated labels, gathered across processes (issue #4719).
-        local_completion_tokens = sum(
-            int((s["labels"] != -100).sum()) for s in self._buffered_inputs if s is not None
-        )
+        local_completion_tokens = sum(int((s["labels"] != -100).sum()) for s in self._buffered_inputs if s is not None)
         self._buffered_num_items = self.accelerator.gather(
             torch.tensor(local_completion_tokens, device=self.accelerator.device)
         ).sum()

From fa5e277a69aab620fe34332caf98294dc9df009c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Tue, 21 Jul 2026 01:05:02 +0000
Subject: [PATCH 33/44] Normalize the num_items loss correctly for DDP and
 grad-accumulation

The buffer-path num_items denominator is the global completion-token count over
the whole accumulation window, so two corrections are needed (both mirror GRPO):

- Divide by accelerator.num_processes in compute_loss, to compensate for DDP
  averaging the gradients across ranks.
- Set compute_loss_func to a non-None sentinel in super().__init__, so Trainer
  does not additionally scale the loss by gradient_accumulation_steps (it fires
  that scaling for models whose forward does not advertise loss kwargs).

Single-process + loss-kwargs models (the CPU test model) are unaffected.
---
 trl/experimental/distillation/distillation_trainer.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/trl/experimental/distillation/distillation_trainer.py b/trl/experimental/distillation/distillation_trainer.py
index 9c8d5145b3d..a4bbac23d06 100644
--- a/trl/experimental/distillation/distillation_trainer.py
+++ b/trl/experimental/distillation/distillation_trainer.py
@@ -486,6 +486,12 @@ def __init__(
             callbacks=callbacks,
             optimizers=optimizers,
             preprocess_logits_for_metrics=preprocess_logits_for_metrics,
+            # In Trainer, `training_step` scales the loss by `gradient_accumulation_steps` only if `compute_loss_func`
+            # is None. Here, loss scaling instead depends on the total number of completion tokens across the global
+            # accumulated batch. To control scaling ourselves, we must disable Trainer's built-in scaling. The simplest
+            # (though a bit hacky) way is to set `compute_loss_func` to any non-None value, which bypasses that behavior
+            # without rewriting `training_step`.
+            compute_loss_func="non-None value to disable scaling",
         )
 
         # ── Prepare teacher model (after super().__init__ so accelerator is ready) ──
@@ -1025,8 +1031,9 @@ def _get_teacher_logits(self, inputs: dict[str, torch.Tensor | Any]) -> torch.Te
     def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
         # transformers computes `num_items_in_batch` from the raw dataloader labels, before on-policy generation
         # replaces the completions; use the count over the generated completions instead (computed in `_fill_buffer`).
+        # Divide by the process count so the per-process loss compensates for DDP gradient averaging (as GRPO does).
         if self.model.training and self._buffered_num_items is not None:
-            num_items_in_batch = self._buffered_num_items
+            num_items_in_batch = self._buffered_num_items.clamp(min=1.0) / self.accelerator.num_processes
 
         if self.use_liger_loss:
             loss = self._compute_liger_loss(model, inputs, num_items_in_batch=num_items_in_batch)

From 8351cb0c09c731af0b49b721b3ad57db731ad3d0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Tue, 21 Jul 2026 01:34:12 +0000
Subject: [PATCH 34/44] Normalize the num_items loss correctly for DDP and
 grad-accumulation

The buffer-path num_items denominator is the global completion-token count over
the whole accumulation window, so two corrections are needed (both mirror GRPO):

- Divide by accelerator.num_processes in compute_loss, to compensate for DDP
  averaging the gradients across ranks.
- Set compute_loss_func to a non-None sentinel in super().__init__, so Trainer
  does not additionally scale the loss by gradient_accumulation_steps (it fires
  that scaling for models whose forward does not advertise loss kwargs).

Single-process + loss-kwargs models (the CPU test model) are unaffected.
---
 trl/experimental/distillation/distillation_trainer.py       | 5 +++++
 .../server_distillation/server_distillation_trainer.py      | 6 ++++++
 2 files changed, 11 insertions(+)

diff --git a/trl/experimental/distillation/distillation_trainer.py b/trl/experimental/distillation/distillation_trainer.py
index a4bbac23d06..a89d055a770 100644
--- a/trl/experimental/distillation/distillation_trainer.py
+++ b/trl/experimental/distillation/distillation_trainer.py
@@ -494,6 +494,11 @@ def __init__(
             compute_loss_func="non-None value to disable scaling",
         )
 
+        # Gradient accumulation requires scaled loss. Normally, loss scaling in the parent class depends on whether the
+        # model accepts loss-related kwargs. Since we compute our own loss, this check is irrelevant. We set
+        # self.model_accepts_loss_kwargs to False to enable scaling.
+        self.model_accepts_loss_kwargs = False
+
         # ── Prepare teacher model (after super().__init__ so accelerator is ready) ──
         if teacher_model is not None:
             # The divergence compares the full next-token distribution of the student against the teacher's, so both
diff --git a/trl/experimental/server_distillation/server_distillation_trainer.py b/trl/experimental/server_distillation/server_distillation_trainer.py
index 19f4e708ec5..a760840ec10 100644
--- a/trl/experimental/server_distillation/server_distillation_trainer.py
+++ b/trl/experimental/server_distillation/server_distillation_trainer.py
@@ -232,6 +232,12 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
                 labels=trimmed_labels,
             )
 
+        # The base trainer disables Trainer's built-in grad-accum loss scaling (via `compute_loss_func`) because it
+        # normalizes by the global completion-token count. The server path normalizes locally with `batchmean` and does
+        # not consume `num_items_in_batch`, so it must re-apply that scaling itself.
+        if self.model.training:
+            loss = loss / self.current_gradient_accumulation_steps
+
         return (loss, student_outputs) if return_outputs else loss
 
     def _get_teacher_token_logprobs_from_server(

From 289c40c90ed969309ade25d5f6d113141b2e6ac9 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Tue, 21 Jul 2026 02:00:25 +0000
Subject: [PATCH 35/44] Switch tests, docs, and example to prompt-only datasets

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
---
 .../15-switch-tests-docs-prompt-only.md       |  8 ++++
 docs/source/distillation_trainer.md           | 44 +++++++------------
 examples/scripts/distillation.py              | 23 ++--------
 .../experimental/test_distillation_trainer.py | 16 ++++---
 4 files changed, 37 insertions(+), 54 deletions(-)
 create mode 100644 distillation_prs/15-switch-tests-docs-prompt-only.md

diff --git a/distillation_prs/15-switch-tests-docs-prompt-only.md b/distillation_prs/15-switch-tests-docs-prompt-only.md
new file mode 100644
index 00000000000..35192666046
--- /dev/null
+++ b/distillation_prs/15-switch-tests-docs-prompt-only.md
@@ -0,0 +1,8 @@
+# Switch tests, docs, and example to prompt-only
+
+On-policy distillation only needs prompts (the student generates completions), so the training surface should use prompt-only datasets. Now unblocked by the `num_items` fix (#6478), which stops prompt-only batches from NaN-ing.
+
+- **Tests:** `conversational_language_modeling` → `conversational_prompt_only` in the end-to-end training/init tests. Kept on messages: the deprecation test and `_make_local_trainer` (feeds `test_loss_normalizes`, which reads completion tokens straight from the collator — switches to prompt-only once generation replaces the collator).
+- **Docs + example:** quick-start and example script now use the `prompt` column (`trl-lib/ultrafeedback-prompt`, a conversational prompt-only dataset). Also removed the dead `lmbda` / off-policy references (the param was already removed from the config).
+
+Verified: `pytest tests/experimental/test_distillation_trainer.py` — 40 passed.
diff --git a/docs/source/distillation_trainer.md b/docs/source/distillation_trainer.md
index 7cab1d23628..7d134085e59 100644
--- a/docs/source/distillation_trainer.md
+++ b/docs/source/distillation_trainer.md
@@ -21,10 +21,10 @@ The `DistillationTrainer` is designed for distilling teacher models of all sizes
 from datasets import load_dataset
 from trl.experimental.distillation import DistillationConfig, DistillationTrainer
 
-# 1. Load dataset and format as prompt-only chat messages
+# 1. Load dataset and format as a prompt-only column
 dataset = load_dataset("openai/gsm8k", "main", split="train")
 dataset = dataset.map(
-    lambda x: {"messages": [{"role": "user", "content": x["question"]}]},
+    lambda x: {"prompt": [{"role": "user", "content": x["question"]}]},
     remove_columns=dataset.column_names,
 )
 
@@ -35,7 +35,6 @@ config = DistillationConfig(
     bf16=True,
     save_strategy="no",
     # Distillation
-    lmbda=1.0,                      # fully on-policy (student generates)
     beta=1.0,                       # reverse KL
     # Teacher
     teacher_model_init_kwargs={"dtype": "bfloat16"},
@@ -54,14 +53,13 @@ trainer.save_model()
 
 ## Usage tips
 
-The [`experimental.distillation.DistillationTrainer`] needs two key parameters set via [`experimental.distillation.DistillationConfig`]:
+The [`experimental.distillation.DistillationTrainer`] trains the student fully on-policy: the student generates its own completions and learns to match the teacher's next-token distribution on them. The key parameter is set via [`experimental.distillation.DistillationConfig`]:
 
-* `lmbda`: controls the student data fraction, i.e., the proportion of on-policy student-generated outputs. When `lmbda=0.0`, training is fully off-policy (dataset completions only). When `lmbda=1.0`, training is fully on-policy (student generates all completions). For values in between, each gradient accumulation slice is randomly assigned as on- or off-policy based on `lmbda`.
 * `beta`: controls the interpolation in the Generalized Jensen-Shannon Divergence. When `beta=0.0` the loss approximates forward KL divergence, while `beta=1.0` approximates reverse KL divergence. Values in between interpolate.
 
-### On-policy vs. off-policy
+### On-policy generation
 
-Setting `lmbda=1.0` (fully on-policy) generally outperforms off-policy distillation because the student learns from its own mistakes rather than imitating trajectories it may never produce. The generation buffer ensures on-policy training stays efficient: prompts across gradient accumulation steps are batched into a single vLLM call.
+Fully on-policy training generally outperforms off-policy distillation because the student learns from its own mistakes rather than imitating trajectories it may never produce. The generation buffer keeps this efficient: prompts across gradient accumulation steps are batched into a single vLLM call.
 
 ### Using an external teacher server
 
@@ -74,7 +72,6 @@ config = DistillationConfig(
     teacher_model_server_url="http://teacher-host:8000",
     loss_top_k=1,       # required with teacher server when beta > 0
     beta=1.0,
-    lmbda=1.0,
 )
 
 trainer = DistillationTrainer(
@@ -93,50 +90,43 @@ When using the teacher server:
 
 ### Expected dataset type
 
-The dataset should be formatted as a [conversational](dataset_formats#conversational) [language modeling](dataset_formats#language-modeling) dataset:
+The dataset should be formatted as a [prompt-only](dataset_formats#prompt-only) dataset. The student generates its own completions on-policy, so only the prompt is needed:
 
 ```python
-{"messages": [{"role": "user", "content": "What color is the sky?"},
-              {"role": "assistant", "content": "It is blue."}]}
-```
-
-When using fully on-policy distillation (`lmbda=1.0`), the assistant turn can be omitted since the student will generate its own completions:
-
-```python
-{"messages": [{"role": "user", "content": "What color is the sky?"}]}
+{"prompt": [{"role": "user", "content": "What color is the sky?"}]}
 ```
 
 ## Example script
 
-Use [`examples/scripts/distillation.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/distillation.py) to launch distillation training from the command line. The script supports full training, mixed on/off-policy, and LoRA via the standard `ModelConfig` flags.
+Use [`examples/scripts/distillation.py`](https://github.com/huggingface/trl/blob/main/examples/scripts/distillation.py) to launch distillation training from the command line. The script supports full training and LoRA via the standard `ModelConfig` flags.
 
 ```bash
-# Full training (off-policy only, lmbda=0):
+# Full training:
 python examples/scripts/distillation.py \
     --model_name_or_path Qwen/Qwen2.5-0.5B-Instruct \
     --teacher_model_name_or_path Qwen/Qwen2.5-1.5B-Instruct \
-    --dataset_name trl-lib/chatbot_arena_completions \
+    --dataset_name trl-lib/ultrafeedback-prompt \
     --learning_rate 2e-5 \
     --per_device_train_batch_size 4 \
     --gradient_accumulation_steps 8 \
-    --lmbda 0.0 \
     --output_dir distilled-model \
     --num_train_epochs 1
 ```
 
 ```bash
-# Mixed on/off-policy (lmbda=0.5):
+# LoRA:
 python examples/scripts/distillation.py \
     --model_name_or_path Qwen/Qwen2.5-0.5B-Instruct \
     --teacher_model_name_or_path Qwen/Qwen2.5-1.5B-Instruct \
-    --dataset_name trl-lib/chatbot_arena_completions \
-    --learning_rate 2e-5 \
+    --dataset_name trl-lib/ultrafeedback-prompt \
+    --learning_rate 2e-4 \
     --per_device_train_batch_size 4 \
     --gradient_accumulation_steps 8 \
-    --lmbda 0.5 \
-    --beta 0.5 \
     --output_dir distilled-model \
-    --num_train_epochs 1
+    --num_train_epochs 1 \
+    --use_peft \
+    --lora_r 64 \
+    --lora_alpha 16
 ```
 
 ## DistillationTrainer
diff --git a/examples/scripts/distillation.py b/examples/scripts/distillation.py
index 807537bf27c..631b0e46858 100644
--- a/examples/scripts/distillation.py
+++ b/examples/scripts/distillation.py
@@ -23,31 +23,15 @@
 
 # docstyle-ignore
 """
-# Full training (off-policy only, lmbda=0):
+# Full training:
 ```
 python examples/scripts/distillation.py \
     --model_name_or_path Qwen/Qwen2.5-0.5B-Instruct \
     --teacher_model_name_or_path Qwen/Qwen2.5-1.5B-Instruct \
-    --dataset_name trl-lib/chatbot_arena_completions \
+    --dataset_name trl-lib/ultrafeedback-prompt \
     --learning_rate 2e-5 \
     --per_device_train_batch_size 4 \
     --gradient_accumulation_steps 8 \
-    --lmbda 0.0 \
-    --output_dir distilled-model \
-    --num_train_epochs 1
-```
-
-# Mixed on/off-policy (lmbda=0.5):
-```
-python examples/scripts/distillation.py \
-    --model_name_or_path Qwen/Qwen2.5-0.5B-Instruct \
-    --teacher_model_name_or_path Qwen/Qwen2.5-1.5B-Instruct \
-    --dataset_name trl-lib/chatbot_arena_completions \
-    --learning_rate 2e-5 \
-    --per_device_train_batch_size 4 \
-    --gradient_accumulation_steps 8 \
-    --lmbda 0.5 \
-    --beta 0.5 \
     --output_dir distilled-model \
     --num_train_epochs 1
 ```
@@ -57,11 +41,10 @@
 python examples/scripts/distillation.py \
     --model_name_or_path Qwen/Qwen2.5-0.5B-Instruct \
     --teacher_model_name_or_path Qwen/Qwen2.5-1.5B-Instruct \
-    --dataset_name trl-lib/chatbot_arena_completions \
+    --dataset_name trl-lib/ultrafeedback-prompt \
     --learning_rate 2e-4 \
     --per_device_train_batch_size 4 \
     --gradient_accumulation_steps 8 \
-    --lmbda 0.0 \
     --output_dir distilled-model \
     --num_train_epochs 1 \
     --use_peft \
diff --git a/tests/experimental/test_distillation_trainer.py b/tests/experimental/test_distillation_trainer.py
index 8705acf9a63..e8ce88aad12 100644
--- a/tests/experimental/test_distillation_trainer.py
+++ b/tests/experimental/test_distillation_trainer.py
@@ -221,6 +221,8 @@ 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")
         return DistillationTrainer(
             model=self.model_id,
@@ -248,7 +250,7 @@ def test_distillation_trainer_train_runs_with_local_teacher(self):
             save_steps=2,
             per_device_eval_batch_size=2,
         )
-        dataset = load_dataset("trl-internal-testing/zen", "conversational_language_modeling")
+        dataset = load_dataset("trl-internal-testing/zen", "conversational_prompt_only")
         trainer = DistillationTrainer(
             model=self.model_id,
             teacher_model=self.model_id,
@@ -319,7 +321,7 @@ def _recording(jsd, labels=None, reduction="batchmean", num_items_in_batch=None)
 
         monkeypatch.setattr(DistillationTrainer, "_reduce_divergence_loss", staticmethod(_recording))
 
-        dataset = load_dataset("trl-internal-testing/zen", "conversational_language_modeling", split="train")
+        dataset = load_dataset("trl-internal-testing/zen", "conversational_prompt_only", split="train")
         trainer = DistillationTrainer(
             model=self.model_id,
             teacher_model=self.model_id,
@@ -349,14 +351,14 @@ def _recording(jsd, labels=None, reduction="batchmean", num_items_in_batch=None)
         ],
     )
     def test_init_with_eval_dataset(self, eval_dataset_type):
-        train_dataset = load_dataset("trl-internal-testing/zen", "conversational_language_modeling", split="train")
+        train_dataset = load_dataset("trl-internal-testing/zen", "conversational_prompt_only", split="train")
 
         if eval_dataset_type == "none":
             eval_dataset = None
         else:
             streaming = "iterable" in eval_dataset_type
             eval_split = load_dataset(
-                "trl-internal-testing/zen", "conversational_language_modeling", split="test", streaming=streaming
+                "trl-internal-testing/zen", "conversational_prompt_only", split="test", streaming=streaming
             )
             if eval_dataset_type in ("dataset", "iterable_dataset"):
                 eval_dataset = eval_split
@@ -419,7 +421,7 @@ def test_distillation_trainer_with_liger(self):
         import importlib
 
         training_args = self._make_args(use_liger_kernel=True, use_cpu=False)
-        dataset = load_dataset("trl-internal-testing/zen", "conversational_language_modeling", split="train")
+        dataset = load_dataset("trl-internal-testing/zen", "conversational_prompt_only", split="train")
 
         trainer = DistillationTrainer(
             model=self.model_id,
@@ -439,7 +441,7 @@ def test_distillation_trainer_with_liger(self):
     def test_teacher_vocab_size_mismatch_raises(self):
         # The local-teacher loss compares full next-token distributions, so student and teacher must share a
         # vocabulary. A teacher with a different vocab_size is rejected (use GOLD for cross-tokenizer distillation).
-        dataset = load_dataset("trl-internal-testing/zen", "conversational_language_modeling", split="train")
+        dataset = load_dataset("trl-internal-testing/zen", "conversational_prompt_only", split="train")
         with pytest.raises(ValueError, match="vocab_size"):
             DistillationTrainer(
                 model=self.model_id,
@@ -452,7 +454,7 @@ def test_teacher_vocab_size_mismatch_raises(self):
     def test_teacher_model_init_kwargs_with_instantiated_teacher_raises(self):
         # `teacher_model_init_kwargs` only applies when the teacher is a model id; passing it alongside an already
         # instantiated teacher is a mistake worth surfacing.
-        dataset = load_dataset("trl-internal-testing/zen", "conversational_language_modeling", split="train")
+        dataset = load_dataset("trl-internal-testing/zen", "conversational_prompt_only", split="train")
         with pytest.raises(ValueError, match="teacher_model_init_kwargs"):
             DistillationTrainer(
                 model=self.model_id,

From b02741cab4e00c892959e96940ee3fd99472f39c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Tue, 21 Jul 2026 02:07:14 +0000
Subject: [PATCH 36/44] remove pr desc

---
 distillation_prs/15-switch-tests-docs-prompt-only.md | 8 --------
 1 file changed, 8 deletions(-)
 delete mode 100644 distillation_prs/15-switch-tests-docs-prompt-only.md

diff --git a/distillation_prs/15-switch-tests-docs-prompt-only.md b/distillation_prs/15-switch-tests-docs-prompt-only.md
deleted file mode 100644
index 35192666046..00000000000
--- a/distillation_prs/15-switch-tests-docs-prompt-only.md
+++ /dev/null
@@ -1,8 +0,0 @@
-# Switch tests, docs, and example to prompt-only
-
-On-policy distillation only needs prompts (the student generates completions), so the training surface should use prompt-only datasets. Now unblocked by the `num_items` fix (#6478), which stops prompt-only batches from NaN-ing.
-
-- **Tests:** `conversational_language_modeling` → `conversational_prompt_only` in the end-to-end training/init tests. Kept on messages: the deprecation test and `_make_local_trainer` (feeds `test_loss_normalizes`, which reads completion tokens straight from the collator — switches to prompt-only once generation replaces the collator).
-- **Docs + example:** quick-start and example script now use the `prompt` column (`trl-lib/ultrafeedback-prompt`, a conversational prompt-only dataset). Also removed the dead `lmbda` / off-policy references (the param was already removed from the config).
-
-Verified: `pytest tests/experimental/test_distillation_trainer.py` — 40 passed.

From 4a9273ffb8de9d782bc1ec7e163aae36ae4eea4f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Tue, 21 Jul 2026 02:09:43 +0000
Subject: [PATCH 37/44] Qualify expected dataset as conversational prompt-only

---
 docs/source/distillation_trainer.md | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/source/distillation_trainer.md b/docs/source/distillation_trainer.md
index 7d134085e59..8f68ac20cf8 100644
--- a/docs/source/distillation_trainer.md
+++ b/docs/source/distillation_trainer.md
@@ -90,7 +90,7 @@ When using the teacher server:
 
 ### Expected dataset type
 
-The dataset should be formatted as a [prompt-only](dataset_formats#prompt-only) dataset. The student generates its own completions on-policy, so only the prompt is needed:
+The dataset should be formatted as a [conversational](dataset_formats#conversational) [prompt-only](dataset_formats#prompt-only) dataset. The student generates its own completions on-policy, so only the prompt is needed:
 
 ```python
 {"prompt": [{"role": "user", "content": "What color is the sky?"}]}

From 2b5b73723bf5d301bbb6d51d889a034a1d745e91 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Tue, 21 Jul 2026 02:32:31 +0000
Subject: [PATCH 38/44] Remove messages-format support and prompt-length config

---
 .../experimental/test_distillation_trainer.py | 48 ++++------
 .../test_server_distillation_trainer.py       | 14 +--
 .../distillation/distillation_config.py       | 27 ------
 .../distillation/distillation_trainer.py      | 92 ++-----------------
 4 files changed, 29 insertions(+), 152 deletions(-)

diff --git a/tests/experimental/test_distillation_trainer.py b/tests/experimental/test_distillation_trainer.py
index e8ce88aad12..683f7c6888b 100644
--- a/tests/experimental/test_distillation_trainer.py
+++ b/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():
diff --git a/tests/experimental/test_server_distillation_trainer.py b/tests/experimental/test_server_distillation_trainer.py
index 87639f550ef..0406d1f22cc 100644
--- a/tests/experimental/test_server_distillation_trainer.py
+++ b/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,
diff --git a/trl/experimental/distillation/distillation_config.py b/trl/experimental/distillation/distillation_config.py
index 6babe3d43f4..044d2c3ac14 100644
--- a/trl/experimental/distillation/distillation_config.py
+++ b/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}.")
 
diff --git a/trl/experimental/distillation/distillation_trainer.py b/trl/experimental/distillation/distillation_trainer.py
index a89d055a770..53e55b19d2e 100644
--- a/trl/experimental/distillation/distillation_trainer.py
+++ b/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

From 4f91a905e77f853855fa61e1e3431b2722578cfc Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Tue, 21 Jul 2026 02:45:21 +0000
Subject: [PATCH 39/44] Pin signature columns to GRPO's [prompt, image, images]

---
 .../distillation/distillation_trainer.py             | 12 +++++-------
 1 file changed, 5 insertions(+), 7 deletions(-)

diff --git a/trl/experimental/distillation/distillation_trainer.py b/trl/experimental/distillation/distillation_trainer.py
index 53e55b19d2e..5c35a97e47d 100644
--- a/trl/experimental/distillation/distillation_trainer.py
+++ b/trl/experimental/distillation/distillation_trainer.py
@@ -570,14 +570,12 @@ 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 = ["prompt", "prompts", "prompt_attention_mask", "messages", "chat_template_kwargs", "tools"]
+        # If `self.args.remove_unused_columns` is True, non-signature columns are removed.
+        # By default, this method sets `self._signature_columns` to the model's expected inputs (usually, "input_ids"
+        # and "attention_mask"). In DistillationTrainer, we preprocess data, so using the model's signature columns
+        # doesn't work. Instead, we set them to the columns expected by the `training_step` method, hence the override.
         if self._signature_columns is None:
-            self._signature_columns = extra_columns
-        else:
-            for col in extra_columns:
-                if col not in self._signature_columns:
-                    self._signature_columns.append(col)
+            self._signature_columns = ["prompt", "image", "images"]
 
     def _get_train_sampler(self, dataset=None):
         if dataset is None:

From 3c4282730015d38a846e8de76e057ae839473abe Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Tue, 21 Jul 2026 02:57:58 +0000
Subject: [PATCH 40/44] Emit completion_mask alongside labels in the generated
 batch

---
 tests/experimental/test_distillation_trainer.py | 17 +++++++++++++++++
 .../distillation/distillation_trainer.py        | 12 +++++++-----
 2 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/tests/experimental/test_distillation_trainer.py b/tests/experimental/test_distillation_trainer.py
index 683f7c6888b..9bec167e98e 100644
--- a/tests/experimental/test_distillation_trainer.py
+++ b/tests/experimental/test_distillation_trainer.py
@@ -403,6 +403,23 @@ def test_loss_normalizes_by_num_items_in_batch(self):
         # Doubling the global count exactly halves the loss (sum / num_items is linear in 1/num_items).
         torch.testing.assert_close(loss_double, loss_mean / 2, rtol=1e-4, atol=1e-6)
 
+    def test_generated_batch_emits_completion_mask(self, monkeypatch):
+        """The generated batch emits `completion_mask`, equal to `labels != -100` (added ahead of the loss switch)."""
+        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"]
+        assert "completion_mask" in inputs
+        assert torch.equal(inputs["completion_mask"].bool(), inputs["labels"] != -100)
+
     @require_liger_kernel
     @require_torch_accelerator
     def test_distillation_trainer_with_liger(self):
diff --git a/trl/experimental/distillation/distillation_trainer.py b/trl/experimental/distillation/distillation_trainer.py
index 5c35a97e47d..780ae31bb41 100644
--- a/trl/experimental/distillation/distillation_trainer.py
+++ b/trl/experimental/distillation/distillation_trainer.py
@@ -742,7 +742,7 @@ def _generate_with_model(self, slices: list[dict[str, torch.Tensor | Any]], on_p
                 else:
                     prompt_token_lengths = torch.full((batch_size,), prompt_width, dtype=torch.long, device=device)
                 completion_lengths = self._get_completion_lengths(generated_tokens, prompt_width)
-                new_attention_mask, new_labels = self._build_sequence_batch(
+                new_attention_mask, new_labels, new_completion_mask = self._build_sequence_batch(
                     generated_tokens, prompt_width, prompt_token_lengths, completion_lengths
                 )
 
@@ -771,6 +771,7 @@ def _generate_with_model(self, slices: list[dict[str, torch.Tensor | Any]], on_p
                 updated["input_ids"] = generated_tokens
                 updated["attention_mask"] = new_attention_mask
                 updated["labels"] = new_labels
+                updated["completion_mask"] = new_completion_mask
 
                 self._buffered_inputs[slice_idx] = updated
                 self._buffered_text_logs[slice_idx] = (prompt_texts, completion_texts)
@@ -832,7 +833,7 @@ def _store_completions_in_buffer(
             completion_ids_padded = torch.stack(completion_tensors)
             new_input_ids = torch.cat([prompt_ids, completion_ids_padded], dim=1)
             completion_lengths = torch.tensor(completion_lengths, device=device, dtype=torch.long)
-            new_attention_mask, new_labels = self._build_sequence_batch(
+            new_attention_mask, new_labels, new_completion_mask = self._build_sequence_batch(
                 new_input_ids, prompt_width, prompt_token_lengths, completion_lengths
             )
 
@@ -848,6 +849,7 @@ def _store_completions_in_buffer(
             updated["input_ids"] = new_input_ids
             updated["attention_mask"] = new_attention_mask
             updated["labels"] = new_labels
+            updated["completion_mask"] = new_completion_mask
             # Update prompts to match the new padding width so prompt_length is consistent
             updated["prompts"] = prompt_ids
             updated["prompt_attention_mask"] = prompt_attention_mask
@@ -861,8 +863,8 @@ def _build_sequence_batch(
         prompt_width: int,
         prompt_token_lengths: torch.Tensor,
         completion_lengths: torch.Tensor,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        """Build attention mask and labels from prompt/completion lengths."""
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """Build attention mask, labels, and completion mask from prompt/completion lengths."""
         prompt_token_lengths = prompt_token_lengths.to(device=new_input_ids.device, dtype=torch.long)
         completion_lengths = completion_lengths.to(device=new_input_ids.device, dtype=torch.long)
         positions = torch.arange(new_input_ids.shape[1], device=new_input_ids.device).unsqueeze(0)
@@ -873,7 +875,7 @@ def _build_sequence_batch(
         new_labels = torch.full_like(new_input_ids, -100)
         new_labels[completion_mask] = new_input_ids[completion_mask]
 
-        return new_attention_mask, new_labels
+        return new_attention_mask, new_labels, completion_mask.int()
 
     # ──────────────────────────────────────────────────────────────────────
     #  Loss computation

From 6497bc5bc80fe94705d5157128a25e16a3aaee09 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Tue, 21 Jul 2026 03:15:22 +0000
Subject: [PATCH 41/44] Switch the loss to consume completion_mask instead of
 labels

---
 .../experimental/test_distillation_trainer.py | 12 +++++--
 .../distillation/distillation_trainer.py      | 32 ++++++++++++-------
 .../server_distillation_trainer.py            |  4 +--
 3 files changed, 31 insertions(+), 17 deletions(-)

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..7aa47ec0a86 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,21 @@ 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 +1042,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")

From b38280e33a31b549f5e613d337f150b32d2db8e3 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?= <gallouedec.quentin@gmail.com>
Date: Tue, 21 Jul 2026 03:32:34 +0000
Subject: [PATCH 42/44] Emit prompt_ids/prompt_mask/completion_ids in the batch

---
 .../experimental/test_distillation_trainer.py | 21 +++++++++++++++++++
 .../distillation/distillation_trainer.py      |  9 ++++++++
 2 files changed, 30 insertions(+)

diff --git a/tests/experimental/test_distillation_trainer.py b/tests/experimental/test_distillation_trainer.py
index 0d0fcf1bd09..fdad338f999 100644
--- a/tests/experimental/test_distillation_trainer.py
+++ b/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):
diff --git a/trl/experimental/distillation/distillation_trainer.py b/trl/experimental/distillation/distillation_trainer.py
index 7aa47ec0a86..61f36e573f9 100644
--- a/trl/experimental/distillation/distillation_trainer.py
+++ b/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

From f1d48968b9cd9603a3709c56a50279d06724a2e0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?=
 <quentin.gallouedec@huggingface.co>
Date: Tue, 21 Jul 2026 16:56:45 +0000
Subject: [PATCH 43/44] Remove stale teacher-server section from the
 distillation doc

---
 docs/source/distillation_trainer.md | 33 +----------------------------
 1 file changed, 1 insertion(+), 32 deletions(-)

diff --git a/docs/source/distillation_trainer.md b/docs/source/distillation_trainer.md
index 8f68ac20cf8..ebdff711051 100644
--- a/docs/source/distillation_trainer.md
+++ b/docs/source/distillation_trainer.md
@@ -6,11 +6,7 @@ The Distillation Trainer implements on-policy knowledge distillation as describe
 
 > Knowledge distillation (KD) is widely used for compressing a teacher model to reduce its inference cost and memory footprint, by training a smaller student model. However, current KD methods for auto-regressive sequence models suffer from distribution mismatch between output sequences seen during training and those generated by the student during inference. To address this issue, we introduce Generalized Knowledge Distillation (GKD). Instead of solely relying on a fixed set of output sequences, GKD trains the student on its self-generated output sequences by leveraging feedback from the teacher on such sequences. Unlike supervised KD approaches, GKD also offers the flexibility to employ alternative loss functions between the student and teacher, which can be useful when the student lacks the expressivity to mimic the teacher's distribution.
 
-The `DistillationTrainer` is designed for distilling teacher models of all sizes into smaller students efficiently. It extends the ideas from the `GKDTrainer` with three key optimizations:
-
-1. **Generation buffer** – decouples the training microbatch size from the generation batch size, letting vLLM batch many prompts in a single call across gradient accumulation steps. This alone can speed up training by up to 40x.
-2. **Teacher server support** – moves the teacher to an external vLLM server so it does not need to fit on the same GPUs as the student.
-3. **Binary-encoded logprob payloads** – packs log-probabilities into base64-encoded NumPy arrays instead of nested JSON lists, shrinking transfer payloads by ~5x.
+The `DistillationTrainer` trains a smaller student model to match a teacher's next-token distribution on the student's own on-policy generations, extending the ideas from the `GKDTrainer`. A generation buffer decouples the training microbatch size from the generation batch size, letting vLLM batch many prompts in a single call across gradient accumulation steps.
 
 > [!NOTE]
 > The Distillation Trainer is currently part of the `trl.experimental` namespace. APIs may change without notice while the feature is iterated on.
@@ -61,33 +57,6 @@ The [`experimental.distillation.DistillationTrainer`] trains the student fully o
 
 Fully on-policy training generally outperforms off-policy distillation because the student learns from its own mistakes rather than imitating trajectories it may never produce. The generation buffer keeps this efficient: prompts across gradient accumulation steps are batched into a single vLLM call.
 
-### Using an external teacher server
-
-For teachers that do not fit on training GPUs (e.g., 100B+ parameters), host the teacher on a separate vLLM server and set `use_teacher_server=True` with `teacher_model_server_url`:
-
-```python
-config = DistillationConfig(
-    output_dir="distilled-model",
-    use_teacher_server=True,
-    teacher_model_server_url="http://teacher-host:8000",
-    loss_top_k=1,       # required with teacher server when beta > 0
-    beta=1.0,
-)
-
-trainer = DistillationTrainer(
-    model="Qwen/Qwen3-4B",
-    args=config,
-    train_dataset=dataset,
-)
-trainer.train()
-```
-
-When using the teacher server:
-- `loss_top_k` must be `> 0` when `beta=0.0` (forward KL)
-- `loss_top_k` must be exactly `1` when `beta > 0` (reverse KL or JSD)
-- `reverse_kl_top_1_mode="argmax"` is not supported
-- Liger kernel is not supported
-
 ### Expected dataset type
 
 The dataset should be formatted as a [conversational](dataset_formats#conversational) [prompt-only](dataset_formats#prompt-only) dataset. The student generates its own completions on-policy, so only the prompt is needed:

From ac1f71f8402ec470357f307655ff70acdb5665f0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Quentin=20Gallou=C3=A9dec?=
 <quentin.gallouedec@huggingface.co>
Date: Wed, 22 Jul 2026 16:40:10 +0000
Subject: [PATCH 44/44] Drop stale local-teacher loss comment

---
 trl/experimental/distillation/distillation_trainer.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/trl/experimental/distillation/distillation_trainer.py b/trl/experimental/distillation/distillation_trainer.py
index 7aa47ec0a86..e7ed831db42 100644
--- a/trl/experimental/distillation/distillation_trainer.py
+++ b/trl/experimental/distillation/distillation_trainer.py
@@ -982,7 +982,6 @@ def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=N
         prompt_length = self._compute_prompt_length(inputs)
         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, :]