AsyncGRPO: cancel stale in-flight rollouts and survive generation fai…

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
 

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

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,

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,16 +865,16 @@ 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:
                 row = next(self._dataset_iter)
             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(
             *[