From 692f7749cdf8b0516324756ea713f8b023c846ad Mon Sep 17 00:00:00 2001
From: aminediro <amine.dirhoussi@huggingface.co>
Date: Thu, 10 Sep 2026 14:55:17 +0000
Subject: [PATCH 1/5] AsyncGRPO: cancel stale in-flight rollouts and survive
 generation failures

---
 docs/source/async_grpo_trainer.md             |   5 +-
 tests/experimental/test_async_grpo_trainer.py | 117 ++++++++++++++++-
 .../async_grpo/async_grpo_trainer.py          |   1 +
 .../async_grpo/async_rollout_worker.py        | 120 ++++++++++++------
 4 files changed, 200 insertions(+), 43 deletions(-)

diff --git a/docs/source/async_grpo_trainer.md b/docs/source/async_grpo_trainer.md
index a915a318136..c2d4e922f63 100644
--- a/docs/source/async_grpo_trainer.md
+++ b/docs/source/async_grpo_trainer.md
@@ -36,7 +36,7 @@ The rollout worker runs in a separate process spawned from the trainer, so rewar
 
 After every `weight_sync_steps` training steps, the updated weights are transferred to the vLLM server via NCCL so that subsequent generations reflect the latest policy.
 
-Because generation and training run concurrently, the training samples may have been generated by a slightly older version of the model. The `max_staleness` parameter controls how many weight updates a sample can lag behind before being discarded.
+Because generation and training run concurrently, the training samples may have been generated by a slightly older version of the model. The `max_staleness` parameter controls how many weight updates a sample can lag behind before being discarded. The worker applies the same limit to its in-flight rollouts: when the policy advances, it cancels the generations of any group that started too many versions ago, so vLLM does not finish work the trainer would drop.
 
 The number of concurrent requests sent to the vLLM server is controlled by `max_inflight_tasks`. By default it is set automatically to `max_staleness × per_device_train_batch_size × gradient_accumulation_steps × num_processes` — the maximum number of samples the trainer can consume before they become stale. Generating more than this is wasteful since the excess samples will be discarded.
 
@@ -220,6 +220,9 @@ A **rollout** is **one full** conversation: a prompt generated to completion, in
 | `rollout/score_s`, `rollout/score_wait_s`, `rollout/score_block_s` | scoring: time to score a group, group wait time to be scored, and how long generation was blocked because the scoring queue was full                                                                                                         |
 | `rollout/vllm_retry_total`                                         | retried vLLM requests. A degraded server otherwise looks like unexplained slowness. It sits here rather than in `completions/` because it counts requests to the server, not generated text: a retried request produced no completion at all |
 | `rollout/backpressure_s`                                           | how long generation was blocked because the rollout queue was full. See [the rollout queue](#the-rollout-queue)                                                                                                                              |
+| `rollout/failed_total`                                             | rollouts that raised (a request that failed every retry, a completion that could not be parsed). The rollout is dropped and its group scored with the rest; only a group where every rollout failed takes the worker down                    |
+| `rollout/dropped_groups_total`                                     | groups dropped because a single rollout survived: a group-relative advantage needs at least two                                                                                                                                               |
+| `rollout/stale_groups_total`                                       | in-flight groups cancelled because the policy moved more than `max_staleness` versions past the one they started at. The trainer would drop their samples anyway (see `sample/dropped_stale_total`), so finishing them only burns vLLM compute |
 
 ### Tools
 
diff --git a/tests/experimental/test_async_grpo_trainer.py b/tests/experimental/test_async_grpo_trainer.py
index 8ad9a4d1c37..51fadc2cb68 100644
--- a/tests/experimental/test_async_grpo_trainer.py
+++ b/tests/experimental/test_async_grpo_trainer.py
@@ -487,7 +487,7 @@ def test_rollout_loop_skips_to_start_index(self):
         dataset = Dataset.from_dict({"prompt": [f"row_{i}" for i in range(10)]})
         loop = self._make_rollout_loop(dataset, dataset_start_index=3)
         it = loop._repeat_iterator()
-        _group_id, row = next(it)
+        _group_id, _index, row = next(it)
         assert row["prompt"] == "row_3"
 
     def test_inner_training_loop_sets_dataset_start_index_from_file(self):
@@ -1188,3 +1188,118 @@ def test_epoch_stop_is_fork_independent(self):
         # 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
+
+
+def _strict_reward(completions, answer, **kwargs):
+    return [1.0 for _ in zip(completions, answer, strict=True)]
+
+
+_ROLLOUT = (
+    [{"role": "assistant", "content": "c"}],
+    [1, 2],
+    [TrainingSequence([1, 2], [0, 1], [0.0, -0.1], "r")],
+    0,
+    0,
+    None,
+)
+
+
+class TestGenerateLoop(TrlTestCase):
+    def _loop(self, num_generations, max_inflight_tasks, max_staleness=4):
+        PartialState()
+        with patch("trl.experimental.async_grpo.async_rollout_worker.add_response_schema", side_effect=lambda x: x):
+            return _AsyncRolloutLoop(
+                model_name="test",
+                dataset=Dataset.from_dict({"prompt": [f"q{i}" for i in range(8)]}),
+                reward_funcs=[dummy_reward_func],
+                processing_class=MagicMock(),
+                rollout_buffer=queue.Queue(),
+                metrics_queue=queue.Queue(),
+                model_version_value=mp.Value("i", 0),
+                heartbeat_value=mp.Value("d", 0.0),
+                failed_event=mp.Event(),
+                exception_info_queue=queue.Queue(),
+                num_generations=num_generations,
+                max_inflight_tasks=max_inflight_tasks,
+                max_staleness=max_staleness,
+            )
+
+    async def _groups(self, loop, generate_one, n, bump_version_to=None):
+        """Run the generate loop until `n` groups reach the score queue."""
+        loop._generate_one = generate_one
+        stop = asyncio.Event()
+        task = asyncio.create_task(loop._generate_loop(stop))
+        if bump_version_to is not None:
+            await asyncio.sleep(0.1)
+            loop._model_version_value.value = bump_version_to
+        groups = [await asyncio.wait_for(loop._groups_to_score.get(), 5) for _ in range(n)]
+        stop.set()
+        await task
+        return groups
+
+    def test_failed_rollout_is_dropped_and_the_group_scored_with_the_rest(self):
+        loop = self._loop(num_generations=4, max_inflight_tasks=4)
+        calls = itertools.count()
+
+        async def generate_one(prompt, tool_dict, tools, group_id):
+            if next(calls) == 1:
+                raise RuntimeError("boom")
+            return _ROLLOUT
+
+        (group,) = asyncio.run(self._groups(loop, generate_one, 1))
+        assert group.group_id == 0
+        assert len(group.completions) == 3
+
+    def test_group_with_a_single_surviving_rollout_is_dropped(self):
+        loop = self._loop(num_generations=3, max_inflight_tasks=3)
+        calls = itertools.count()
+
+        async def generate_one(prompt, tool_dict, tools, group_id):
+            if next(calls) < 2:
+                raise RuntimeError("boom")
+            return _ROLLOUT
+
+        (group,) = asyncio.run(self._groups(loop, generate_one, 1))
+        assert group.group_id == 1
+
+    def test_group_where_every_rollout_fails_raises(self):
+        loop = self._loop(num_generations=2, max_inflight_tasks=2)
+
+        async def generate_one(prompt, tool_dict, tools, group_id):
+            raise RuntimeError("boom")
+
+        loop._generate_one = generate_one
+        with pytest.raises(RuntimeError, match="boom"):
+            asyncio.run(asyncio.wait_for(loop._generate_loop(asyncio.Event()), 5))
+
+    def test_stale_in_flight_groups_are_cancelled_when_the_policy_advances(self):
+        loop = self._loop(num_generations=2, max_inflight_tasks=4, max_staleness=1)
+
+        async def generate_one(prompt, tool_dict, tools, group_id):
+            if loop.model_version == 0:
+                await asyncio.Event().wait()
+            return _ROLLOUT
+
+        groups = asyncio.run(self._groups(loop, generate_one, 2, bump_version_to=2))
+        assert [g.group_id for g in groups] == [2, 3]
+        assert [g.model_version for g in groups] == [2, 2]
+
+    def test_partially_dispatched_stale_group_is_regenerated_as_a_smaller_group(self):
+        loop = self._loop(num_generations=4, max_inflight_tasks=2, max_staleness=0)
+
+        async def generate_one(prompt, tool_dict, tools, group_id):
+            if loop.model_version == 0:
+                await asyncio.Event().wait()
+            return _ROLLOUT
+
+        (group,) = asyncio.run(self._groups(loop, generate_one, 1, bump_version_to=1))
+        assert group.group_id == 0
+        assert len(group.completions) == 2
+        assert group.model_version == 1
+
+    def test_reward_kwargs_are_trimmed_to_the_surviving_rollouts(self):
+        PartialState()
+        group = _group([[_ROLLOUT[2][0]]] * 3, [[1, 2]] * 3)
+        group.reward_kwargs = {"answer": ["a"] * 4}
+        samples = asyncio.run(_bare_loop([_strict_reward])._score_group(group))
+        assert len(samples) == 3
diff --git a/trl/experimental/async_grpo/async_grpo_trainer.py b/trl/experimental/async_grpo/async_grpo_trainer.py
index b47e5cfd867..b95579365f2 100644
--- a/trl/experimental/async_grpo/async_grpo_trainer.py
+++ b/trl/experimental/async_grpo/async_grpo_trainer.py
@@ -1014,6 +1014,7 @@ def __init__(
                     queue_maxsize=self.args.queue_maxsize,
                     vllm_server_url=self.args.vllm_server_base_url,
                     max_tokens=self.args.max_completion_length,
+                    max_staleness=self.args.max_staleness,
                     temperature=self.args.temperature,
                     top_p=self.args.top_p,
                     top_k=self.args.top_k,
diff --git a/trl/experimental/async_grpo/async_rollout_worker.py b/trl/experimental/async_grpo/async_rollout_worker.py
index d936ed9c8de..f6b680cb664 100644
--- a/trl/experimental/async_grpo/async_rollout_worker.py
+++ b/trl/experimental/async_grpo/async_rollout_worker.py
@@ -314,6 +314,7 @@ def __init__(
         score_queue_maxsize: int = 16,
         vllm_server_url: str = "http://localhost:8000",
         max_tokens: int = 32,
+        max_staleness: int = 4,
         temperature: float = 1.0,
         top_p: float = 1.0,
         top_k: int = 0,
@@ -367,6 +368,7 @@ def __init__(
         self.max_inflight_tasks = max_inflight_tasks
         self.queue_maxsize = queue_maxsize
         self.max_tokens = max_tokens
+        self.max_staleness = max_staleness
         self.temperature = temperature
         self.top_p = top_p
         self.top_k = top_k
@@ -504,14 +506,38 @@ async def _generate_loop(self, stop_event: asyncio.Event) -> None:
         inflight_tasks: dict[asyncio.Task, tuple[int, int, Any, object, Messages]] = {}
         free_slots = set(range(self.max_inflight_tasks))
         work_iter = self._repeat_iterator()
+        last_version = self.model_version
 
         self._generation_start_time = time.monotonic()
         try:
             while True:
                 # Wall-clock for cross-process comparison; parent uses time.time() in check_health.
                 self._heartbeat_value.value = time.time()
+
+                version = self.model_version
+                if version != last_version:
+                    last_version = version
+                    stale = {
+                        group_id
+                        for group_id, group in pending_groups.items()
+                        if version - group.model_version > self.max_staleness
+                    }
+                    for task, (group_id, slot, name, environment, _prompt) in list(inflight_tasks.items()):
+                        if group_id in stale:
+                            task.cancel()
+                            del inflight_tasks[task]
+                            free_slots.add(slot)
+                            if environment is not None:
+                                self._environment_pool[name].append(environment)
+                    for group_id in stale:
+                        del pending_groups[group_id]
+                        del pending_completed[group_id]
+                    if stale:
+                        self._counters["rollout/stale_groups_total"] += len(stale)
+                        logger.info(f"[generate] cancelled {len(stale)} stale group(s) at version {version}")
+
                 while free_slots and not stop_event.is_set():
-                    group_id, row = next(work_iter)
+                    group_id, index, row = next(work_iter)
                     slot = free_slots.pop()
                     # The environment is selected per example via its `environment` field (multi-env); only its tools
                     # are exposed in the example's prompt. When there are no environments, every example shares the
@@ -587,7 +613,7 @@ async def _generate_loop(self, stop_event: asyncio.Event) -> None:
                             env_rewards=[],
                             rollout_rewards=[],
                         )
-                        pending_completed[group_id] = 0
+                        pending_completed[group_id] = index
 
                     task = asyncio.create_task(
                         self._generate_one(prompt, tool_dict=tool_dict, tools=tools, group_id=group_id)
@@ -608,45 +634,58 @@ async def _generate_loop(self, stop_event: asyncio.Event) -> None:
                 for task in done:
                     group_id, slot, name, environment, prompt = inflight_tasks.pop(task)
                     free_slots.add(slot)
-                    if task.exception() is not None:
-                        raise task.exception()
-
-                    (
-                        completion,
-                        completion_ids,
-                        sequences,
-                        tool_call_count,
-                        tool_failure_count,
-                        rollout_reward,
-                    ) = task.result()
                     group = pending_groups[group_id]
-                    group.prompts.append(prompt)
-                    group.completions.append(completion)
-                    group.completions_ids.append(completion_ids)
-                    group.completions_sequences.append(sequences)
-                    group.tool_call_counts.append(tool_call_count)
-                    group.tool_failure_counts.append(tool_failure_count)
-                    group.rollout_rewards.append(rollout_reward)
-                    # The environment owns the reward: score it now, while this rollout's environment still holds its
-                    # final state and before returning it to the pool. `get_reward` may be async awaiting yields to
-                    # inflight requests instead of halting them. The env is returned to the pool only after scoring, so
-                    # a concurrent rollout can't draw and reset it during the await. Record `(env class, reward)` so
-                    # `_score_group` can place it in the matching env's reward column; rollouts whose env owns no reward
-                    # record `None` (turned into NaN and ignored) to stay aligned with the group's other per-rollout lists.
-                    if self._env_reward_types:
-                        env_type = type(environment)
-                        if env_type in self._env_reward_types:
-                            get_reward = environment.get_reward
-                            reward = await get_reward() if inspect.iscoroutinefunction(get_reward) else get_reward()
-                            group.env_rewards.append((env_type, reward))
-                        else:
-                            group.env_rewards.append(None)
+                    error = task.exception()
+                    if error is not None:
+                        logger.warning(f"[generate] rollout failed for group {group_id}, dropping it", exc_info=error)
+                        self._counters["rollout/failed_total"] += 1
+                    else:
+                        (
+                            completion,
+                            completion_ids,
+                            sequences,
+                            tool_call_count,
+                            tool_failure_count,
+                            rollout_reward,
+                        ) = task.result()
+                        group.prompts.append(prompt)
+                        group.completions.append(completion)
+                        group.completions_ids.append(completion_ids)
+                        group.completions_sequences.append(sequences)
+                        group.tool_call_counts.append(tool_call_count)
+                        group.tool_failure_counts.append(tool_failure_count)
+                        group.rollout_rewards.append(rollout_reward)
+                        # The environment owns the reward: score it now, while this rollout's environment still holds
+                        # its final state and before returning it to the pool. `get_reward` may be async awaiting
+                        # yields to inflight requests instead of halting them. The env is returned to the pool only
+                        # after scoring, so a concurrent rollout can't draw and reset it during the await. Record
+                        # `(env class, reward)` so `_score_group` can place it in the matching env's reward column;
+                        # rollouts whose env owns no reward record `None` (turned into NaN and ignored) to stay aligned
+                        # with the group's other per-rollout lists.
+                        if self._env_reward_types:
+                            env_type = type(environment)
+                            if env_type in self._env_reward_types:
+                                get_reward = environment.get_reward
+                                reward = (
+                                    await get_reward() if inspect.iscoroutinefunction(get_reward) else get_reward()
+                                )
+                                group.env_rewards.append((env_type, reward))
+                            else:
+                                group.env_rewards.append(None)
+                        self._total_completion_tokens += len(completion_ids)
                     if environment is not None:
                         self._environment_pool[name].append(environment)
-                    self._total_completion_tokens += len(completion_ids)
                     pending_completed[group_id] += 1
 
                     if pending_completed[group_id] == self.num_generations:
+                        del pending_groups[group_id]
+                        del pending_completed[group_id]
+                        if not group.completions:
+                            raise error
+                        if len(group.completions) < 2:
+                            logger.warning(f"[generate] dropping group {group_id}: a single rollout succeeded")
+                            self._counters["rollout/dropped_groups_total"] += 1
+                            continue
                         group.queued_at = time.monotonic()
                         t_blocked = None
                         while True:
@@ -662,8 +701,6 @@ async def _generate_loop(self, stop_event: asyncio.Event) -> None:
                         if t_blocked is not None:
                             # Generation held back by scoring
                             self._push_metrics({"rollout/score_block_s": time.monotonic() - t_blocked})
-                        del pending_groups[group_id]
-                        del pending_completed[group_id]
         finally:
             for task in inflight_tasks:
                 task.cancel()
@@ -828,7 +865,7 @@ def _push_rollout_metrics(
             }
         )
 
-    def _repeat_iterator(self) -> Iterator[tuple[int, dict[str, Any]]]:
+    def _repeat_iterator(self) -> Iterator[tuple[int, int, dict[str, Any]]]:
         group_id = 0
         while True:
             try:
@@ -836,8 +873,8 @@ def _repeat_iterator(self) -> Iterator[tuple[int, dict[str, Any]]]:
             except StopIteration:
                 self._dataset_iter = iter(self.dataset)
                 row = next(self._dataset_iter)
-            for _ in range(self.num_generations):
-                yield group_id, row
+            for index in range(self.num_generations):
+                yield group_id, index, row
             group_id += 1
 
     async def _generate_one(
@@ -971,11 +1008,12 @@ async def _generate_one_turn(self, prompt_ids: list[int]) -> tuple[list[int], li
         return choice["token_ids"], choice["logprobs"]["token_logprobs"]
 
     async def _score_group(self, group: RolloutGroup) -> list[RolloutSample]:
+        n = len(group.completions)
         kwargs = dict(
             completions=group.completions,
             prompts=group.prompts,
             completion_ids=group.completions_ids,
-            **group.reward_kwargs,
+            **{key: values[:n] for key, values in group.reward_kwargs.items()},
         )
         all_rewards = await asyncio.gather(
             *[

From 9f150a451ce67be895a282ed6760e1aa33b1881d Mon Sep 17 00:00:00 2001
From: Kashif Rasul <kashif.rasul@gmail.com>
Date: Mon, 14 Sep 2026 10:05:17 +0200
Subject: [PATCH 2/5] Fix stale rollout cancellation edge cases

---
 docs/source/async_distillation_trainer.md     |  4 +-
 .../test_async_distillation_trainer.py        | 63 +++++++++++++++----
 tests/experimental/test_async_grpo_trainer.py | 24 ++++++-
 .../async_distillation_config.py              |  4 +-
 .../async_distillation_trainer.py             |  1 +
 .../async_rollout_worker.py                   | 30 ++++++++-
 .../async_grpo/async_rollout_worker.py        | 21 ++++---
 .../async_grpo/openenv_harness.py             | 55 +++++++++++++---
 8 files changed, 168 insertions(+), 34 deletions(-)

diff --git a/docs/source/async_distillation_trainer.md b/docs/source/async_distillation_trainer.md
index 4e32137e2e9..a8fa0a13ea0 100644
--- a/docs/source/async_distillation_trainer.md
+++ b/docs/source/async_distillation_trainer.md
@@ -74,7 +74,8 @@ logprob for at each `beta` regime.
 
 After every `weight_sync_steps` training steps, the updated student weights are transferred to its vLLM server via
 NCCL. As with [`experimental.async_grpo.AsyncGRPOTrainer`], generation runs ahead of training, so samples may reflect a slightly stale
-policy; `max_staleness` controls how many weight updates a sample can lag behind before being discarded.
+policy; `max_staleness` controls how many weight updates a sample can lag behind before in-flight work is cancelled or
+a queued sample is discarded.
 
 ## Quick start
 
@@ -229,6 +230,7 @@ A **rollout** is one prompt taken all the way through: generated by the student,
 | `rollout/inflight`         | rollouts in flight, i.e. generating or being scored                                                                                                                                                                                          |
 | `rollout/vllm_retry_total` | retried vLLM requests, to either the student's server or a teacher's. A degraded server otherwise looks like unexplained slowness. It sits here rather than in `completions/` because it counts requests to a server, not generated text     |
 | `rollout/backpressure_s`   | how long generation was blocked because the rollout queue was full. See [the rollout queue](#the-rollout-queue)                                                                                                                              |
+| `rollout/stale_samples_total` | in-flight samples cancelled because the policy moved more than `max_staleness` versions past the one they started at                                                                                                                        |
 
 ### Samples arriving from the queue
 
diff --git a/tests/experimental/test_async_distillation_trainer.py b/tests/experimental/test_async_distillation_trainer.py
index 9a4ee90af24..8866a88bf92 100644
--- a/tests/experimental/test_async_distillation_trainer.py
+++ b/tests/experimental/test_async_distillation_trainer.py
@@ -356,6 +356,23 @@ def _bare_loop(tokenizer, teacher_server_urls):
 TWO_TEACHERS = {"math": "http://math:8002", "code": "http://code:8003"}
 
 
+def _rollout_loop(dataset, **kwargs):
+    ctx = mp.get_context("spawn")
+    loop_kwargs = dict(
+        model_name="test",
+        dataset=dataset,
+        processing_class=MagicMock(),
+        rollout_buffer=ctx.Queue(),
+        metrics_queue=ctx.Queue(),
+        model_version_value=ctx.Value("i", 0),
+        heartbeat_value=ctx.Value("d", 0.0),
+        failed_event=ctx.Event(),
+        exception_info_queue=ctx.Queue(),
+    )
+    loop_kwargs.update(kwargs)
+    return _AsyncRolloutLoop(**loop_kwargs)
+
+
 class TestWorkerMetrics:
     """The worker's payload has the same shape as the trainer's sink, so draining it in `log()` is an append."""
 
@@ -922,6 +939,39 @@ def test_stops_once_the_prompt_target_is_reached(self, trained, before_resume, s
         assert control.should_training_stop is should_stop
 
 
+class TestGenerateLoop(TrlTestCase):
+    def test_stale_in_flight_samples_are_cancelled_when_the_policy_advances(self):
+        loop = _rollout_loop(
+            Dataset.from_dict({"prompt": [f"q{i}" for i in range(8)]}), max_inflight_tasks=2, max_staleness=0
+        )
+        cancelled = []
+
+        async def run():
+            stop = asyncio.Event()
+
+            async def generate_and_score_one(prompt_id, row):
+                version = loop.model_version
+                if version == 0:
+                    try:
+                        await asyncio.Event().wait()
+                    except asyncio.CancelledError:
+                        cancelled.append(prompt_id)
+                        raise
+                stop.set()
+                return types.SimpleNamespace(completion_mask=[1], model_version=version, enqueued_at=None)
+
+            loop._generate_and_score_one = generate_and_score_one
+            task = asyncio.create_task(loop._generate_loop(stop))
+            await asyncio.sleep(0.1)
+            loop._model_version_value.value = 1
+            await asyncio.wait_for(task, 5)
+
+        asyncio.run(run())
+        assert len(cancelled) == 2
+        assert loop.rollout_buffer.get(timeout=5).model_version == 1
+        assert loop._metrics_queue.get(timeout=5)["rollout/stale_samples_total"] == 2
+
+
 class TestRolloutStateCheckpoint(TrlTestCase):
     """Prompt-index checkpoint/resume logic — no GPU or vLLM required."""
 
@@ -974,17 +1024,8 @@ def record(*_args, **_kwargs):
         assert written_before_super == [True]
 
     def test_rollout_loop_skips_to_start_index(self):
-        ctx = mp.get_context("spawn")
-        loop = _AsyncRolloutLoop(
-            model_name="test",
-            dataset=Dataset.from_dict({"prompt": [f"row_{i}" for i in range(10)]}),
-            processing_class=MagicMock(),
-            rollout_buffer=ctx.Queue(),
-            model_version_value=ctx.Value("i", 0),
-            heartbeat_value=ctx.Value("d", 0.0),
-            failed_event=ctx.Event(),
-            exception_info_queue=ctx.Queue(),
-            metrics_queue=ctx.Queue(),
+        loop = _rollout_loop(
+            Dataset.from_dict({"prompt": [f"row_{i}" for i in range(10)]}),
             dataset_start_index=3,
         )
         _prompt_id, row = next(loop._repeat_iterator())
diff --git a/tests/experimental/test_async_grpo_trainer.py b/tests/experimental/test_async_grpo_trainer.py
index 2bef75bc7f0..636756995f6 100644
--- a/tests/experimental/test_async_grpo_trainer.py
+++ b/tests/experimental/test_async_grpo_trainer.py
@@ -1457,13 +1457,21 @@ async def generate_one(prompt, tool_dict, tools, group_id):
 
     def test_stale_in_flight_groups_are_cancelled_when_the_policy_advances(self):
         loop = self._loop(num_generations=2, max_inflight_tasks=4, max_staleness=1)
+        cancelled = []
 
         async def generate_one(prompt, tool_dict, tools, group_id):
             if loop.model_version == 0:
-                await asyncio.Event().wait()
+                try:
+                    await asyncio.Event().wait()
+                except asyncio.CancelledError:
+                    await asyncio.sleep(0)
+                    cancelled.append(group_id)
+                    raise
+            assert len(cancelled) == 4
             return _ROLLOUT
 
         groups = asyncio.run(self._groups(loop, generate_one, 2, bump_version_to=2))
+        assert len(cancelled) == 4
         assert [g.group_id for g in groups] == [2, 3]
         assert [g.model_version for g in groups] == [2, 2]
 
@@ -1480,6 +1488,20 @@ async def generate_one(prompt, tool_dict, tools, group_id):
         assert len(group.completions) == 2
         assert group.model_version == 1
 
+    def test_stale_group_with_one_undispatched_rollout_is_skipped(self):
+        loop = self._loop(num_generations=4, max_inflight_tasks=3, max_staleness=0)
+
+        async def generate_one(prompt, tool_dict, tools, group_id):
+            if loop.model_version == 0:
+                await asyncio.Event().wait()
+            if group_id == 0:
+                raise RuntimeError("the stale tail must not be dispatched")
+            return _ROLLOUT
+
+        (group,) = asyncio.run(self._groups(loop, generate_one, 1, bump_version_to=1))
+        assert group.group_id == 1
+        assert group.model_version == 1
+
     def test_reward_kwargs_are_trimmed_to_the_surviving_rollouts(self):
         PartialState()
         group = _group([[_ROLLOUT[2][0]]] * 3, [[1, 2]] * 3)
diff --git a/trl/experimental/async_distillation/async_distillation_config.py b/trl/experimental/async_distillation/async_distillation_config.py
index 9a24e9a2917..9f7598f9f9b 100644
--- a/trl/experimental/async_distillation/async_distillation_config.py
+++ b/trl/experimental/async_distillation/async_distillation_config.py
@@ -150,7 +150,7 @@ class AsyncDistillationConfig(_BaseConfig):
             num_processes`.
         max_staleness (`int`, *optional*, defaults to `4`):
             Maximum number of weight update steps a rollout sample can lag behind the current model version before
-            being discarded.
+            an in-flight sample is cancelled or a queued sample is discarded.
         queue_maxsize (`int`, *optional*, defaults to `1024`):
             Maximum number of rollout samples to buffer in the rollout queue.
         weight_sync_steps (`int`, *optional*, defaults to `1`):
@@ -358,7 +358,7 @@ class AsyncDistillationConfig(_BaseConfig):
         default=4,
         metadata={
             "help": "Maximum number of weight update steps a rollout sample can lag behind the current model "
-            "version before being discarded."
+            "version before an in-flight sample is cancelled or a queued sample is discarded."
         },
     )
     queue_maxsize: int = field(
diff --git a/trl/experimental/async_distillation/async_distillation_trainer.py b/trl/experimental/async_distillation/async_distillation_trainer.py
index 127b62799f1..0bcd5c48230 100644
--- a/trl/experimental/async_distillation/async_distillation_trainer.py
+++ b/trl/experimental/async_distillation/async_distillation_trainer.py
@@ -1133,6 +1133,7 @@ def __init__(
                     teacher_top_k=self.args.teacher_top_k,
                     teacher_temperature=self.args.teacher_temperature,
                     max_tokens=self.args.max_completion_length,
+                    max_staleness=self.args.max_staleness,
                     temperature=self.args.temperature,
                     top_p=self.args.top_p,
                     top_k=self.args.top_k,
diff --git a/trl/experimental/async_distillation/async_rollout_worker.py b/trl/experimental/async_distillation/async_rollout_worker.py
index 756c22a94ce..8354a6ab35e 100644
--- a/trl/experimental/async_distillation/async_rollout_worker.py
+++ b/trl/experimental/async_distillation/async_rollout_worker.py
@@ -208,6 +208,7 @@ def __init__(
         teacher_top_k: int = 8,
         teacher_temperature: float = 1.0,
         max_tokens: int = 32,
+        max_staleness: int = 4,
         temperature: float = 1.0,
         top_p: float = 1.0,
         top_k: int = 0,
@@ -243,6 +244,7 @@ def __init__(
         self.max_inflight_tasks = max_inflight_tasks
         self.queue_maxsize = queue_maxsize
         self.max_tokens = max_tokens
+        self.max_staleness = max_staleness
         self.temperature = temperature
         self.top_p = top_p
         self.top_k = top_k
@@ -321,19 +323,41 @@ async def _resolve_teacher_model_names(self) -> None:
             logger.info(f"teacher {teacher_id!r} at {url} serves {self.teacher_model_names[teacher_id]}")
 
     async def _generate_loop(self, stop_event: asyncio.Event) -> None:
-        inflight_tasks: dict[asyncio.Task, int] = {}
+        # Keep the dispatch version beside the slot: a sample does not expose its model version until its task returns.
+        inflight_tasks: dict[asyncio.Task, tuple[int, int]] = {}
         free_slots = set(range(self.max_inflight_tasks))
         work_iter = self._repeat_iterator()
+        last_version = self.model_version
 
         self._generation_start_time = time.monotonic()
         try:
             while True:
                 self._heartbeat_value.value = time.time()
+
+                version = self.model_version
+                if version != last_version:
+                    last_version = version
+                    stale_tasks = [
+                        task
+                        for task, (_slot, task_version) in inflight_tasks.items()
+                        if version - task_version > self.max_staleness
+                    ]
+                    for task in stale_tasks:
+                        task.cancel()
+                    if stale_tasks:
+                        await asyncio.gather(*stale_tasks, return_exceptions=True)
+                    for task in stale_tasks:
+                        slot, _task_version = inflight_tasks.pop(task)
+                        free_slots.add(slot)
+                    if stale_tasks:
+                        self._counters["rollout/stale_samples_total"] += len(stale_tasks)
+                        logger.info(f"cancelled {len(stale_tasks)} stale rollout(s) at version {version}")
+
                 while free_slots and not stop_event.is_set():
                     prompt_id, row = next(work_iter)
                     slot = free_slots.pop()
                     task = asyncio.create_task(self._generate_and_score_one(prompt_id, row))
-                    inflight_tasks[task] = slot
+                    inflight_tasks[task] = (slot, self.model_version)
 
                 if not inflight_tasks:
                     if stop_event.is_set():
@@ -346,7 +370,7 @@ async def _generate_loop(self, stop_event: asyncio.Event) -> None:
                     continue
 
                 for task in done:
-                    slot = inflight_tasks.pop(task)
+                    slot, _task_version = inflight_tasks.pop(task)
                     free_slots.add(slot)
                     if task.exception() is not None:
                         raise task.exception()
diff --git a/trl/experimental/async_grpo/async_rollout_worker.py b/trl/experimental/async_grpo/async_rollout_worker.py
index 6fc173da737..a9f5edf0c48 100644
--- a/trl/experimental/async_grpo/async_rollout_worker.py
+++ b/trl/experimental/async_grpo/async_rollout_worker.py
@@ -533,13 +533,16 @@ async def _generate_loop(self, stop_event: asyncio.Event) -> None:
                         for group_id, group in pending_groups.items()
                         if version - group.model_version > self.max_staleness
                     }
-                    for task, (group_id, slot, name, environment, _prompt) in list(inflight_tasks.items()):
-                        if group_id in stale:
-                            task.cancel()
-                            del inflight_tasks[task]
-                            free_slots.add(slot)
-                            if environment is not None:
-                                self._environment_pool[name].append(environment)
+                    stale_tasks = [task for task, values in inflight_tasks.items() if values[0] in stale]
+                    for task in stale_tasks:
+                        task.cancel()
+                    if stale_tasks:
+                        await asyncio.gather(*stale_tasks, return_exceptions=True)
+                    for task in stale_tasks:
+                        _group_id, slot, name, environment, _prompt = inflight_tasks.pop(task)
+                        free_slots.add(slot)
+                        if environment is not None:
+                            self._environment_pool[name].append(environment)
                     for group_id in stale:
                         del pending_groups[group_id]
                         del pending_completed[group_id]
@@ -549,6 +552,10 @@ async def _generate_loop(self, stop_event: asyncio.Event) -> None:
 
                 while free_slots and not stop_event.is_set():
                     group_id, index, row = next(work_iter)
+                    # Missing state at a nonzero index means stale cancellation discarded the group's earlier work.
+                    # Do not dispatch a lone remainder: it cannot produce a group-relative advantage.
+                    if group_id not in pending_groups and self.num_generations - index < 2:
+                        continue
                     slot = free_slots.pop()
                     # The environment is selected per example via its `environment` field (multi-env); only its tools
                     # are exposed in the example's prompt. When there are no environments, every example shares the
diff --git a/trl/experimental/async_grpo/openenv_harness.py b/trl/experimental/async_grpo/openenv_harness.py
index 16cd00164c7..bf693f4c857 100644
--- a/trl/experimental/async_grpo/openenv_harness.py
+++ b/trl/experimental/async_grpo/openenv_harness.py
@@ -17,11 +17,12 @@
 import asyncio
 import functools
 import json
+import threading
 import time
 import uuid
 from collections.abc import Callable
-from concurrent.futures import ThreadPoolExecutor
-from dataclasses import dataclass
+from concurrent.futures import Future, ThreadPoolExecutor
+from dataclasses import dataclass, field
 from typing import Any, Protocol, TypedDict, cast
 
 from accelerate.logging import get_logger
@@ -78,6 +79,13 @@ class HarnessTurn:
     """One agent turn from the trace, passed to `train_turn_fn` to decide whether it is trained."""
 
     messages: list[Message]  # the conversation sent to the model this turn (the prompt)
+
+
+@dataclass
+class _CancellationToken:
+    event: threading.Event = field(default_factory=threading.Event)
+    session: object | None = None
+    sampling_future: Future | None = None
     tools: list[dict] | None  # tools available to the model this turn
     content: str  # the assistant's text content this turn
     tool_calls: list[dict]  # the tool calls the assistant emitted (empty for a pure-text turn)
@@ -147,14 +155,30 @@ def __init__(
         self._session_pool = ThreadPoolExecutor(
             max_workers=max(1, self.max_inflight_tasks), thread_name_prefix="harness-session"
         )
-        # In-flight sessions, so `_run_loops` can close them on stop (see there). set ops are atomic under the GIL.
+        # In-flight sessions, so `_run_loops` can close them on stop (see there). Set ops are atomic under the GIL.
         self._live_sessions: set = set()
 
     async def _generate_one(self, prompt, tool_dict, tools, group_id=0):
         # TODO(@openenv): provide an async version for performance
         #  OpenEnv's harness layer is synchronous, so run the whole session on the pool.
         loop = asyncio.get_running_loop()
-        result, metrics = await loop.run_in_executor(self._session_pool, self._run_session, prompt, group_id)
+        token = _CancellationToken()
+        future = loop.run_in_executor(self._session_pool, self._run_session, prompt, group_id, token)
+        try:
+            result, metrics = await asyncio.shield(future)
+        except asyncio.CancelledError:
+            # Executor cancellation does not stop its thread, and white-box sampling runs in a separate event-loop
+            # future. Signal both paths, then keep the worker slot occupied until its thread has unwound.
+            token.event.set()
+            if token.sampling_future is not None:
+                token.sampling_future.cancel()
+            if token.session is not None:
+                try:
+                    await asyncio.to_thread(token.session.close)
+                except Exception:
+                    logger.warning("closing cancelled harness session failed", exc_info=True)
+            await asyncio.gather(future, return_exceptions=True)
+            raise
         # Pushed here and not in `_run_session`: the accumulators are plain dicts, and the pool runs many sessions at
         # once, so a push off the event loop would race the score loop's.
         if metrics is not None:
@@ -179,7 +203,7 @@ def _close(session):
         finally:
             self._session_pool.shutdown(wait=True)
 
-    def _run_session(self, prompt, group_id=0):
+    def _run_session(self, prompt, group_id=0, token: _CancellationToken | None = None):
         """Drive one OpenEnv session to completion, on a pool thread.
 
         Returns `(the _generate_one tuple, rollout metrics or None)`. The metrics are handed back rather than pushed
@@ -202,16 +226,20 @@ def _run_session(self, prompt, group_id=0):
         except Exception:
             logger.warning("harness session create failed; scoring rollout as unscorable", exc_info=True)
             return self._EMPTY_ROLLOUT, None
-        self._live_sessions.add(session)  # tracked so a stop can close it (unblocks wait_for_completion below)
+        token = token or _CancellationToken()
+        token.session = session
+        self._live_sessions.add(session)
         timed_out = False
         trace: list[TraceEntry] = []
         tool_calls_by_name: dict[str, int] = {}
         try:
+            if token.event.is_set():
+                return self._EMPTY_ROLLOUT, None
             if self._adapter is not None:
                 # white-box: the adapter runs the tool loop, calling `_sample_turn` each turn.
                 turns: list[TurnRecord] = []
                 result = self._adapter.run_white_box(
-                    functools.partial(self._sample_turn, turns), session, self._limits
+                    functools.partial(self._sample_turn, turns, token), session, self._limits
                 )
                 completion = result.messages
                 tool_call_count = int(result.metrics.get("tool_calls", len(result.tool_trace)))
@@ -267,7 +295,9 @@ def _run_session(self, prompt, group_id=0):
             except Exception:
                 logger.warning("harness session close failed", exc_info=True)
 
-    def _sample_turn(self, turns: list[TurnRecord], messages, tools, sampling) -> ModelStepResult:
+    def _sample_turn(
+        self, turns: list[TurnRecord], token: _CancellationToken, messages, tools, sampling
+    ) -> ModelStepResult:
         """OpenEnv `ModelStep`: sample one assistant turn against vLLM and record a `TurnRecord` into `turns`."""
         prompt_ids = self.tokenizer.apply_chat_template(
             messages,
@@ -279,7 +309,14 @@ def _sample_turn(self, turns: list[TurnRecord], messages, tools, sampling) -> Mo
             **self.chat_template_kwargs,
         )
         # ModelStep is sync on a pool thread; bridge the async vLLM POST onto the loop's event loop.
-        turn_ids, logprobs = asyncio.run_coroutine_threadsafe(self._generate_one_turn(prompt_ids), self._loop).result()
+        future = asyncio.run_coroutine_threadsafe(self._generate_one_turn(prompt_ids), self._loop)
+        token.sampling_future = future
+        if token.event.is_set():
+            future.cancel()
+        try:
+            turn_ids, logprobs = future.result()
+        finally:
+            token.sampling_future = None
         turns.append(TurnRecord(prompt_ids, turn_ids, logprobs))
         message = parse_response(self.tokenizer, turn_ids, prefix=prompt_ids)
         return ModelStepResult(

From 44c45d9878613b14293d9d61d479a9473c1169a2 Mon Sep 17 00:00:00 2001
From: Kashif Rasul <kashif.rasul@gmail.com>
Date: Mon, 14 Sep 2026 10:07:20 +0200
Subject: [PATCH 3/5] Fix async distillation docstring formatting

---
 .../async_distillation/async_distillation_config.py           | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/trl/experimental/async_distillation/async_distillation_config.py b/trl/experimental/async_distillation/async_distillation_config.py
index 9f7598f9f9b..00d5aea51d5 100644
--- a/trl/experimental/async_distillation/async_distillation_config.py
+++ b/trl/experimental/async_distillation/async_distillation_config.py
@@ -149,8 +149,8 @@ class AsyncDistillationConfig(_BaseConfig):
             `-1` (auto), which sets it to `max_staleness * per_device_train_batch_size * gradient_accumulation_steps *
             num_processes`.
         max_staleness (`int`, *optional*, defaults to `4`):
-            Maximum number of weight update steps a rollout sample can lag behind the current model version before
-            an in-flight sample is cancelled or a queued sample is discarded.
+            Maximum number of weight update steps a rollout sample can lag behind the current model version before an
+            in-flight sample is cancelled or a queued sample is discarded.
         queue_maxsize (`int`, *optional*, defaults to `1024`):
             Maximum number of rollout samples to buffer in the rollout queue.
         weight_sync_steps (`int`, *optional*, defaults to `1`):

From 87e897e23b90731c52a926ceaea9fb4cd2b20e47 Mon Sep 17 00:00:00 2001
From: Kashif Rasul <kashif.rasul@gmail.com>
Date: Mon, 14 Sep 2026 10:14:20 +0200
Subject: [PATCH 4/5] Fix harness cancellation token fields

---
 trl/experimental/async_grpo/openenv_harness.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/trl/experimental/async_grpo/openenv_harness.py b/trl/experimental/async_grpo/openenv_harness.py
index bf693f4c857..f79258dbe24 100644
--- a/trl/experimental/async_grpo/openenv_harness.py
+++ b/trl/experimental/async_grpo/openenv_harness.py
@@ -79,6 +79,9 @@ class HarnessTurn:
     """One agent turn from the trace, passed to `train_turn_fn` to decide whether it is trained."""
 
     messages: list[Message]  # the conversation sent to the model this turn (the prompt)
+    tools: list[dict] | None  # tools available to the model this turn
+    content: str  # the assistant's text content this turn
+    tool_calls: list[dict]  # the tool calls the assistant emitted (empty for a pure-text turn)
 
 
 @dataclass
@@ -86,9 +89,6 @@ class _CancellationToken:
     event: threading.Event = field(default_factory=threading.Event)
     session: object | None = None
     sampling_future: Future | None = None
-    tools: list[dict] | None  # tools available to the model this turn
-    content: str  # the assistant's text content this turn
-    tool_calls: list[dict]  # the tool calls the assistant emitted (empty for a pure-text turn)
 
 
 def _tools_to_schema(tools: list) -> list[dict] | None:

From a15bbe0bad4784fde183a4ea2fea0de7d1581212 Mon Sep 17 00:00:00 2001
From: Kashif Rasul <kashif.rasul@gmail.com>
Date: Mon, 14 Sep 2026 10:22:26 +0200
Subject: [PATCH 5/5] Initialize accelerate state in rollout loop tests

---
 tests/experimental/test_async_distillation_trainer.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/experimental/test_async_distillation_trainer.py b/tests/experimental/test_async_distillation_trainer.py
index 8866a88bf92..dce8192661c 100644
--- a/tests/experimental/test_async_distillation_trainer.py
+++ b/tests/experimental/test_async_distillation_trainer.py
@@ -25,6 +25,7 @@
 
 import pytest
 import torch
+from accelerate import PartialState
 from datasets import Dataset, load_dataset
 from transformers import AutoTokenizer
 from transformers.testing_utils import torch_device
@@ -357,6 +358,7 @@ def _bare_loop(tokenizer, teacher_server_urls):
 
 
 def _rollout_loop(dataset, **kwargs):
+    PartialState()
     ctx = mp.get_context("spawn")
     loop_kwargs = dict(
         model_name="test",