Merge remote-tracking branch 'upstream/main' into fix-torch-jit-script-method-py314-warning

Author: albertvillanova

docs/source/distillation_trainer.md

@@ -54,34 +54,15 @@ 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.
-* `distillation_objective`: selects the training objective. Use `"jsd"` for the generalized JSD/KL objective, or `"iw_opd"` for Importance-Weighted On-Policy Distillation, which reweights sampled-token reverse-KL policy-gradient updates by prefix teacher-student agreement. IW-OPD requires `lmbda=1.0`.
-* `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
 
 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.
 
-### Importance-Weighted On-Policy Distillation
-
-Set `distillation_objective="iw_opd"` to use Importance-Weighted On-Policy Distillation from [On the Position Bias of On-Policy Distillation](https://huggingface.co/papers/2606.22600). IW-OPD computes a sampled-token OPD advantage from teacher and student log-probabilities, then upweights earlier tokens and downweights later tokens according to accumulated teacher-student drift.
-
-```python
-config = DistillationConfig(
-    output_dir="distilled-model",
-    distillation_objective="iw_opd",
-    iw_opd_gamma=0.5,
-    lmbda=1.0,
-)
-```
-
-IW-OPD is only available for fully on-policy training. It is incompatible with `use_liger_kernel=True`; when
-`use_vllm=True`, set `vllm_sync_frequency=1`. The vLLM path requests sampled-token rollout log-probabilities and
-uses them as the detached IW-OPD rollout policy log-probs.
-
 ### 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`:

docs/source/paper_index.md

@@ -684,37 +684,6 @@ training_args = GRPOConfig(
 )
 ```
 
-
-### Rethinking the Trust Region in LLM Reinforcement Learning
-
-**📜 Paper**: https://huggingface.co/papers/2602.04879
-
-DPPO replaces PPO/GRPO's heuristic ratio-clipping with a principled trust region based on direct policy divergence estimates. PPO-style clipping masks tokens based on the probability ratio π/μ, which over-penalizes low-probability tokens and under-penalizes high-probability ones. DPPO instead masks based on direct approximations of policy divergence (TV or KL), ensuring updates stay within a theoretically grounded trust region. Four divergence approximations are supported: `binary_tv`, `binary_kl`, `topk_tv`, and `topk_kl`.
-
-```python
-from trl.experimental.dppo import DPPOConfig, DPPOTrainer
-
-training_args = DPPOConfig(
-    divergence_type="binary_tv",  # divergence approximation
-    divergence_topk=20,  # K for top-K divergence modes (Section 7 / Appendix G.2 of the paper)
-    epsilon=0.15,  # δ_low threshold (Appendix F of the paper)
-    epsilon_high=0.15,  # δ_high threshold (Appendix F of the paper)
-    clip_ratio_c=20.0,  # IS ratio upper bound C (Section 5.4 of the paper)
-    beta=0.0,  # KL regularization coefficient
-    use_vllm=True,
-)
-
-trainer = DPPOTrainer(
-    model="your-model",
-    reward_funcs=[...],
-    args=training_args,
-    train_dataset=dataset,
-)
-trainer.train()
-```
-
-The official code [sail-sg/Stable-RL](https://github.com/sail-sg/Stable-RL)
-
 ## Optimal Advantage Regression
 
 Papers relating to the [`experimental.a2po.A2POTrainer`].
@@ -1665,37 +1634,37 @@ Papers relating to training a student model with the help of a teacher model.
 
 **📜 Paper**: https://huggingface.co/papers/2306.13649
 
-Introduces Generalized Knowledge Distillation (GKD), which addresses distribution mismatch in KD for auto-regressive models by training the student on its own generated outputs with teacher feedback, instead of a fixed set of sequences. GKD supports flexible loss functions (e.g. beyond KL when the student cannot match the teacher) and integrates with RL fine-tuning (RLHF). The paper reports results on summarization, translation, arithmetic reasoning, and instruction-tuning. Used in TRL via [`experimental.distillation.DistillationTrainer`] and [`experimental.gkd.GKDTrainer`]. To reproduce the paper's setting, use this configuration:
+Introduces Generalized Knowledge Distillation (GKD), which addresses distribution mismatch in KD for auto-regressive models by training the student on its own generated outputs with teacher feedback, instead of a fixed set of sequences. GKD supports flexible loss functions (e.g. beyond KL when the student cannot match the teacher) and integrates with RL fine-tuning (RLHF). The paper reports results on summarization, translation, arithmetic reasoning, and instruction-tuning. Used in TRL via [`experimental.gkd.GKDTrainer`], which exposes the paper's on/off-policy mixing (`lmbda`). [`experimental.distillation.DistillationTrainer`] implements the same generalized-JSD objective for the always-on-policy case. To reproduce the paper's setting, use this configuration:
 
 ```python
-from trl.experimental.distillation import DistillationConfig
+from trl.experimental.gkd import GKDConfig
 
 # XSum summarization task (Table A.1 of the paper)
-training_args = DistillationConfig(
+training_args = GKDConfig(
     lmbda=0.5,  # λ student data fraction (Section 3 of the paper)
     beta=0.5,  # β Generalized JSD interpolation, 0=KL, 1=reverse KL (Section 3 of the paper)
     temperature=1.0,  # student training temperature (Appendix A of the paper)
     max_steps=40000,  # training steps (Table A.1 of the paper)
     learning_rate=3e-4,  # learning rate (Table A.1 of the paper)
     per_device_train_batch_size=32,  # batch size (Table A.1 of the paper)
     warmup_steps=2000,  # warm-up steps (Table A.1 of the paper)
-    max_completion_length=64,  # max output tokens (Table A.1 of the paper)
+    max_new_tokens=64,  # max output tokens (Table A.1 of the paper)
 )
 ```
 
 ### On the Position Bias of On-Policy Distillation
 
 **📜 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

docs/source/vllm_integration.md

@@ -3,7 +3,7 @@
 This document will guide you through the process of using vLLM with TRL for faster generation in online methods like GRPO and Online DPO. We first summarize a tl;dr on how to use vLLM with TRL, and then we will go into the details of how it works under the hood.
 
 > [!WARNING]
-> TRL currently only supports vLLM versions from `0.16.0` to `0.24.0`. Please ensure you have a version in this range installed to avoid compatibility issues.
+> TRL currently only supports vLLM versions from `0.17.0` to `0.25.1`. Please ensure you have a version in this range installed to avoid compatibility issues.
 
 > [!TIP]
 > The following trainers currently support generation with vLLM:

pyproject.toml

@@ -80,7 +80,7 @@ test = [
     "pytest"
 ]
 vllm = [
-    "vllm>=0.16.0,<=0.24.0",
+    "vllm>=0.17.0,<=0.25.1",
     "fastapi",
     "pydantic",
     "aiohttp>=3.13.3",

tests/distributed/test_distributed.py

@@ -148,6 +148,45 @@ def test_dpo(self, config, get_config_path):
         )
         # fmt: on
 
+    @pytest.mark.parametrize(
+        "config",
+        [
+            "ddp",
+            pytest.param(
+                "zero2",
+                marks=pytest.mark.xfail(
+                    Version(transformers.__version__) == Version("5.1.0"),
+                    reason="Upstream incompatibility: deepspeed and transformers==5.1.0 (see transformers#43780)",
+                ),
+            ),
+            pytest.param(
+                "zero3",
+                marks=pytest.mark.xfail(
+                    Version(transformers.__version__) == Version("5.1.0"),
+                    reason="Upstream incompatibility: deepspeed and transformers==5.1.0 (see transformers#43780)",
+                ),
+            ),
+        ],
+    )
+    def test_dpo_precompute_ref_log_probs(self, config, get_config_path):
+        # `--eval_strategy epoch` passes an eval dataset, so reference log-probs are precomputed for both the train and
+        # eval splits (two passes), which is what previously broke multi-GPU precompute (fingerprint cache mismatch, and
+        # a corrupted ZeRO-3 parameter coordinator from re-initializing DeepSpeed on the policy model per pass).
+        # fmt: off
+        run_command(
+            [
+                "accelerate", "launch", "--config_file", get_config_path(config), "trl/scripts/dpo.py",
+                "--output_dir", self.tmp_dir,
+                "--model_name_or_path", "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5",
+                "--dataset_name", "trl-internal-testing/zen",
+                "--dataset_config", "standard_preference",
+                "--precompute_ref_log_probs",
+                "--eval_strategy", "epoch",
+            ],
+            os.environ.copy(),
+        )
+        # fmt: on
+
     @require_liger_kernel
     @pytest.mark.parametrize(
         "config",

tests/experimental/test_async_grpo_trainer.py

@@ -56,14 +56,17 @@ def dummy_reward_func(completions, **kwargs):
 class _StubRolloutWorker:
     """Minimal rollout worker stub for testing the trainer in isolation."""
 
-    def __init__(self, tokenizer, dataset, num_generations: int = 8, samples_per_weight_sync: int = 10):
+    def __init__(
+        self, tokenizer, dataset, num_generations: int = 8, samples_per_weight_sync: int = 10, fork_k: int = 1
+    ):
         self.rollout_buffer = queue.Queue()
         self._samples_per_weight_sync = samples_per_weight_sync
         self._model_version = 0
+        self._fork_k = fork_k
         self._sample_iter = self._make_sample_iter(tokenizer, dataset, num_generations)
 
     def _make_sample_iter(self, tokenizer, dataset, num_generations):
-        for row in itertools.cycle(dataset):
+        for group_id, row in enumerate(itertools.cycle(dataset)):
             completions = [
                 [{"role": "assistant", "content": f"{row['completion'][0]['content']} {idx}"}]
                 for idx in range(num_generations)
@@ -79,16 +82,22 @@ def _make_sample_iter(self, tokenizer, dataset, num_generations):
             advantages = (rewards - rewards.mean()) / rewards.std()
             for idx in range(num_generations):
                 completion_ids = prompt_completion_ids[idx][len(prompt_ids) :]
-                yield RolloutSample(
+                sample = RolloutSample(
                     prompt=row["prompt"],
                     completion=completions[idx],
                     input_ids=prompt_ids + completion_ids,
                     completion_mask=[0] * len(prompt_ids) + [1] * len(completion_ids),
                     old_log_probs=[0.0] * len(prompt_ids) + [-0.5] * len(completion_ids),
                     advantage=float(advantages[idx]),
                     model_version=self._model_version,
+                    group_id=group_id,
                     metrics={"reward": float(rewards[idx]), "reward_std": float(rewards.std())},
                 )
+                # fork_k rows per generation, all sharing this group_id: mimics message-mode forking a
+                # conversation into several training rows. Only the shared group_id matters for the epoch
+                # count, so identical duplicates are enough here (row shapes are covered by the reconciler tests).
+                for _ in range(self._fork_k):
+                    yield sample
 
     def _fill_queue(self):
         for _ in range(self._samples_per_weight_sync):
@@ -326,6 +335,7 @@ def _rollout_sample(length: int, advantage: float = 0.0, reward: float = 0.0) ->
         "completion_mask": [0] + [1] * (length - 1),
         "old_log_probs": [0.0] * length,
         "advantage": advantage,
+        "group_id": 0,
         "metrics": {"reward": reward},
     }
 
@@ -653,6 +663,7 @@ def _group(completions_sequences, completions_ids):
         tool_call_counts=[0] * n,
         tool_failure_counts=[0] * n,
         model_version=7,
+        group_id=0,
         env_rewards=[None] * n,
     )
 
@@ -722,3 +733,53 @@ def maybe_none(completions, **kwargs):
         assert math.isnan(samples[0].metrics["reward"])
         assert samples[1].advantage == 0.0  # only one scorable row -> zero-centered
         assert samples[1].metrics["reward"] == 2.0
+
+
+@pytest.mark.skipif(
+    not is_ampere_or_newer() and torch_device != "xpu",
+    reason="Flash Attention 2 requires Ampere or newer GPU, or XPU",
+)
+class TestEpochStop(TrlTestCase):
+    """`num_train_epochs` stops after N full passes over the PROMPTS, counted as distinct group_ids.
+
+    The point of the fix is that this is fork-independent: a conversation that message-mode forks into several training
+    rows still counts as one prompt-group, so the number of epochs does not depend on the fork rate (only the number of
+    optimizer steps does). Driven end-to-end through the real trainer (stub worker + real forward/backward), which is
+    what actually regresses if the wiring breaks.
+    """
+
+    def _train(self, fork_k, num_train_epochs=2):
+        model_id = "trl-internal-testing/tiny-Qwen2ForCausalLM-2.5"
+        dataset = load_dataset("trl-internal-testing/zen", "conversational_prompt_completion", split="train")
+        args = AsyncGRPOConfig(
+            output_dir=self.tmp_dir,
+            num_train_epochs=num_train_epochs,  # epoch-driven: no explicit max_steps
+            per_device_train_batch_size=3,
+            num_generations=3,
+            max_completion_length=8,
+            token_budget=-1,  # FixedCountBatcher: deterministic samples/step, so groups accrue predictably
+            vllm_server_timeout=5.0,
+            report_to="none",
+        )
+        worker = _StubRolloutWorker(AutoTokenizer.from_pretrained(model_id), dataset, num_generations=3, fork_k=fork_k)
+        trainer = AsyncGRPOTrainer(
+            model=model_id, reward_funcs=dummy_reward_func, args=args, train_dataset=dataset, rollout_worker=worker
+        )
+        trainer.train()
+        return trainer, len(dataset)
+
+    def test_epoch_stop_is_fork_independent(self):
+        no_fork, num_prompts = self._train(fork_k=1)
+        forked, _ = self._train(fork_k=3)
+
+        # Both stop after num_train_epochs=2 full passes over the prompts, i.e. ~2 x num_prompts distinct
+        # groups, plus at most one micro-batch of overshoot — regardless of the fork rate. (HF's own
+        # state.epoch is meaningless here: it's global_step/max_steps over an infinite IterableDataset,
+        # so we judge epochs by distinct prompt-groups trained, which is what the callback targets.)
+        for trainer in (no_fork, forked):
+            assert 2 * num_prompts <= len(trainer._trained_groups) < 3 * num_prompts
+
+        # Forks add rows, hence fixed-size optimizer steps: 3x rows per conversation must take strictly more
+        # steps for the same 2 epochs. If forks leaked into the epoch count, the forked run would instead
+        # stop in FEWER prompt-passes (the pre-fix bug).
+        assert forked.state.global_step > no_fork.state.global_step