Filter zero-advantage groups in AsyncGRPO

docs/source/async_grpo_trainer.md

@@ -354,6 +354,9 @@ A **rollout** is **one full** conversation: a prompt generated to completion, in
 | `rollout/vllm_running_requests`                                    | requests the server's engines are decoding (vLLM's `num_requests_running`, summed over engines). The gap to `rollout/inflight` is what sits in the API server process, building responses                                                   |
 | `rollout/vllm_queued_requests`                                     | requests the server has accepted but not yet scheduled (vLLM's `num_requests_waiting`, summed over engines). The cap does not grow while this is non-zero                                                                                    |
 | `rollout/vllm_preemptions_total`                                   | requests the server preempted during the logging window (vLLM's `num_preemptions_total`, summed over engines). Reported only: a weight sync resets the server's caches, which preempts every running request, so it is not a pressure signal                                                                     |
+| `rollout/zero_advantage_groups_total`, `reward_kept`               | groups dropped because every rollout got the same reward, and the mean reward of the groups that were trained on. `reward` is the policy's reward over every scored group, dropped ones included, so it is the one to track                      |
+
+A group whose rollouts all received the same reward (or none of which could be scored) has zero advantage everywhere, so the worker drops it instead of enqueuing it: it would cost a forward and backward and move nothing. The loop does not resample to replace a dropped group, so `batch/groups_per_step` goes down while the batch keeps its `samples_per_step` samples. A dropped group still counts toward an epoch, which stays one pass over the prompts.
 
 ### Tools
 

tests/experimental/test_async_grpo_trainer.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 import asyncio
+import contextlib
 import itertools
 import json
 import math
@@ -33,7 +34,7 @@
 from accelerate import PartialState
 from datasets import Dataset, load_dataset
 from requests.adapters import BaseAdapter
-from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedModel
+from transformers import AutoModelForCausalLM, AutoTokenizer, PreTrainedModel, TrainerControl
 from transformers.testing_utils import torch_device
 
 
@@ -45,6 +46,7 @@
     FixedCountBatcher,
     RolloutWorkerProtocol,
     TokenBudgetBatcher,
+    _EpochStopCallback,
     _iter_vllm_named_params,
     _reduce_metric,
     round_lora_rank,
@@ -633,6 +635,7 @@ def _checkpoint_trainer(output_dir, dataset_start_index=0):
     # The checkpoint paths only run for an `AsyncRolloutWorker`, which requires vLLM
     trainer.rollout_worker = MagicMock(spec=AsyncRolloutWorker)
     trainer.rollout_worker._loop_kwargs = {"dataset_start_index": dataset_start_index}
+    trainer.rollout_worker.filtered_groups = queue.Queue()
     return trainer
 
 
@@ -668,6 +671,25 @@ def test_save_checkpoint_writes_rollout_state(self):
             data = json.load(f)
         assert data["prompt_index"] == 15  # dataset_start_index(10) + first_untrained(5)
 
+    def test_save_checkpoint_skips_filtered_groups(self):
+        trainer = self._trainer_at_step_5({0, 1, 3, 4}, dataset_start_index=10)
+        trainer.rollout_worker.filtered_groups.put(2)
+
+        with patch.object(_BaseTrainer, "_save_checkpoint"):
+            trainer._save_checkpoint(trainer.model, None)
+
+        with open(os.path.join(self.tmp_dir, "checkpoint-5", "rollout_state.json")) as f:
+            assert json.load(f)["prompt_index"] == 15
+
+    def test_epoch_stop_counts_filtered_groups(self):
+        trainer = self._trainer_at_step_5({0, 1, 3, 4}, dataset_start_index=0)
+        trainer.rollout_worker.filtered_groups.put(2)
+        control = TrainerControl()
+
+        _EpochStopCallback(trainer, target_groups=5).on_step_end(trainer.args, trainer.state, control)
+
+        assert control.should_training_stop
+
     def test_save_checkpoint_writes_rollout_state_before_the_hub_push(self):
         # `super()._save_checkpoint` is what uploads the checkpoint folder under `hub_strategy="checkpoint"`, so the
         # file has to exist by the time it runs. Writing it from an `on_save` callback would not: `Trainer` fires
@@ -749,6 +771,7 @@ def _rollout_loop(**kwargs):
         "failed_event": mp.Event(),
         "exception_info_queue": mp.Queue(),
         "metrics_queue": mp.Queue(),
+        "filtered_groups": mp.Queue(),
         **kwargs,
     }
     return _AsyncRolloutLoop(**kwargs)
@@ -1972,8 +1995,8 @@ def test_metrics_are_per_conversation_and_independent(self):
 
         samples = asyncio.run(_rollout_loop(reward_funcs=[two_reward])._score_group(group))
 
-        assert samples[0].metrics["reward"] == 1.0
-        assert samples[1].metrics["reward"] == 3.0 and samples[2].metrics["reward"] == 3.0
+        assert samples[0].metrics["reward_kept"] == 1.0
+        assert samples[1].metrics["reward_kept"] == 3.0 and samples[2].metrics["reward_kept"] == 3.0
         assert samples[0].metrics["reward_std"] == pytest.approx(1.0)
         assert samples[1].metrics["rewards/two_reward"] == 3.0
         # The fork's two rows must not share a metrics dict (the score loop mutates it per sample).
@@ -1982,19 +2005,59 @@ def test_metrics_are_per_conversation_and_independent(self):
     def test_all_none_reward_conversation_is_unscorable(self):
         # A conversation for which every reward func returns None gets advantage 0 and NaN reward.
         def maybe_none(completions, **kwargs):
-            return [None, 2.0]
+            return [None, 2.0, 4.0]
 
         seq_a = TrainingSequence([1, 2, 10], [0, 0, 1], [0, 0, -0.1], [[]] * 3, "c0")
         seq_b = TrainingSequence([1, 2, 20], [0, 0, 1], [0, 0, -0.2], [[]] * 3, "c1")
-        group = _group([[seq_a], [seq_b]], completions_ids=[[10], [20]])
+        seq_c = TrainingSequence([1, 2, 30], [0, 0, 1], [0, 0, -0.3], [[]] * 3, "c2")
+        group = _group([[seq_a], [seq_b], [seq_c]], completions_ids=[[10], [20], [30]])
 
         samples = asyncio.run(_rollout_loop(reward_funcs=[maybe_none])._score_group(group))
 
-        assert len(samples) == 2
+        assert len(samples) == 3
         assert samples[0].advantage == 0.0  # unscorable -> advantage 0
-        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
+        assert math.isnan(samples[0].metrics["reward_kept"])
+        assert samples[1].advantage < 0 < samples[2].advantage  # advantage over the scorable rows only
+        assert samples[1].metrics["reward_kept"] == 2.0
+
+
+def constant_reward(completions, **kwargs):
+    return [1.0] * len(completions)
+
+
+class TestFilterZeroAdvantageGroups(TrlTestCase):
+    @pytest.fixture(autouse=True)
+    def _init_accelerate_state(self):
+        PartialState()
+
+    def _score(self, loop):
+        for group_id, reward_func in enumerate([constant_reward, two_reward]):
+            seq_a = TrainingSequence([1, 2, 10], [0, 0, 1], [0, 0, -0.1], [[]] * 3, "c0")
+            seq_b = TrainingSequence([1, 2, 20], [0, 0, 1], [0, 0, -0.2], [[]] * 3, "c1")
+            group = _group([[seq_a], [seq_b]], completions_ids=[[10], [20]])
+            group.group_id = group_id
+            group.queued_at = time.monotonic()
+            loop.reward_funcs = [reward_func]
+            loop._groups_to_score.put_nowait(group)
+            loop._groups_to_score.put_nowait(None)
+            loop._loop.run_until_complete(loop._score_loop(asyncio.Event()))
+        samples, metrics = [], defaultdict(list)
+        with contextlib.suppress(queue.Empty):
+            while True:
+                samples.append(loop.rollout_buffer.get(timeout=1))
+        with contextlib.suppress(queue.Empty):
+            while True:
+                for key, value in loop._metrics_queue.get(timeout=1).items():
+                    metrics[key].append(value)
+        return samples, metrics
+
+    def test_constant_reward_group_never_reaches_the_buffer(self):
+        loop = _rollout_loop()
+        samples, metrics = self._score(loop)
+        assert [s.group_id for s in samples] == [1, 1]
+        assert sum(metrics["rollout/zero_advantage_groups_total"]) == 1
+        assert loop._filtered_groups.get(timeout=1) == 0
+        assert metrics["reward"] == [(2.0, 2), (4.0, 2)]  # both groups, the filtered one included
 
 
 @pytest.mark.skipif(

trl/experimental/async_grpo/async_grpo_trainer.py

@@ -13,6 +13,7 @@
 # limitations under the License.
 
 
+import contextlib
 import contextvars
 import itertools
 import json
@@ -237,11 +238,11 @@ def on_train_begin(self, _args, _state, _control, **_kwargs):
 class _EpochStopCallback(TrainerCallback):
     """Stop after `num_train_epochs` full passes over the prompt dataset.
 
-    An epoch is counted in distinct prompt-groups actually trained (accumulated in the collator, which runs on the main
-    process just before the model forward). This is fork-independent: all generations of a prompt and all forked rows
-    of a conversation share one `group_id`, so a conversation forking into many rows still counts once. Only the main
-    process collates (`dispatch_batches=True`), so the stop decision is reduced across ranks to keep data-parallel
-    workers in lockstep.
+    An epoch is counted in distinct prompt-groups consumed: trained (accumulated in the collator, which runs on the
+    main process just before the model forward) or dropped for zero advantage. This is fork-independent: all
+    generations of a prompt and all forked rows of a conversation share one `group_id`, so a conversation forking into
+    many rows still counts once. Only the main process collates (`dispatch_batches=True`), so the stop decision is
+    reduced across ranks to keep data-parallel workers in lockstep.
     """
 
     def __init__(self, trainer: "AsyncGRPOTrainer", target_groups: int):
@@ -250,8 +251,8 @@ def __init__(self, trainer: "AsyncGRPOTrainer", target_groups: int):
 
     def on_step_end(self, args, _state, control, **_kwargs):
         acc = self._trainer.accelerator
-        trained = self._trainer._groups_before_resume + len(self._trainer._trained_groups)
-        reached = torch.tensor(int(trained >= self._target), device=acc.device)
+        consumed = self._trainer._groups_before_resume + len(self._trainer._consumed_groups())
+        reached = torch.tensor(int(consumed >= self._target), device=acc.device)
         if int(acc.reduce(reached, reduction="sum").item()) >= 1:
             control.should_training_stop = True
             # `max_steps` is only a safety ceiling here, so `DefaultFlowCallback` never reaches its end-of-training save.
@@ -1259,6 +1260,7 @@ def __init__(
         # Epoch handling: stop after num_train_epochs full passes over the PROMPT dataset, counted as distinct
         # prompt-groups trained (fork-independent).
         self._trained_groups: set[int] = set()
+        self._filtered_groups: set[int] = set()
         # Tracks restart to match `num_train_epochs`
         self._groups_before_resume = 0
         self._epoch_stop_groups: int | None = None
@@ -1880,14 +1882,21 @@ def _sync_weight_lora(self, t0: float) -> tuple[float, float, float]:
                 self.rollout_worker.update_model_version(self.model_version)
         return t_pause, t_barrier, t_transfer
 
+    def _consumed_groups(self) -> set[int]:
+        if isinstance(self.rollout_worker, AsyncRolloutWorker):
+            with contextlib.suppress(queue.Empty):
+                while True:
+                    self._filtered_groups.add(self.rollout_worker.filtered_groups.get_nowait())
+        return self._trained_groups | self._filtered_groups
+
     def _save_checkpoint(self, model, trial):
         if self.accelerator.is_main_process and isinstance(self.rollout_worker, AsyncRolloutWorker):
             checkpoint_folder = f"{PREFIX_CHECKPOINT_DIR}-{self.state.global_step}"
             checkpoint_dir = os.path.join(self._get_output_dir(trial=trial), checkpoint_folder)
             os.makedirs(checkpoint_dir, exist_ok=True)
-            trained = self._trained_groups
-            first_untrained = next(g for g in itertools.count() if g not in trained)
-            prompt_index = self.rollout_worker._loop_kwargs["dataset_start_index"] + first_untrained
+            consumed = self._consumed_groups()
+            first_unconsumed = next(g for g in itertools.count() if g not in consumed)
+            prompt_index = self.rollout_worker._loop_kwargs["dataset_start_index"] + first_unconsumed
             # `model_version` rides along so adapter names keep counting across a resume: restarting at v1 would
             # republish a different adapter under a name a still-running server already holds.
             rollout_state = {"prompt_index": prompt_index, "model_version": self.model_version}

trl/experimental/async_grpo/async_rollout_worker.py

@@ -301,6 +301,7 @@ def _child_main(
     failed_event: MPEvent,
     exception_info_queue: MPQueue,
     metrics_queue: MPQueue,
+    filtered_groups: MPQueue,
 ) -> None:
     _scrub_child_env()
     # `accelerate.logging.get_logger` requires `PartialState()` to have been called.
@@ -316,6 +317,7 @@ def _child_main(
         failed_event=failed_event,
         exception_info_queue=exception_info_queue,
         metrics_queue=metrics_queue,
+        filtered_groups=filtered_groups,
     )
     child_ready_event.set()
     _spawn_stop_watcher(rollout_loop, stop_event)
@@ -351,6 +353,7 @@ def __init__(
         failed_event: MPEvent,
         exception_info_queue: MPQueue,
         metrics_queue: MPQueue,
+        filtered_groups: MPQueue,
         tools: list[Callable] | None = None,
         environment_factory: Callable[[], object] | dict[str, Callable[[], object]] | None = None,
         num_generations: int = 8,
@@ -404,6 +407,7 @@ def __init__(
         self._failed_event = failed_event  # shared mp.Event
         self._exception_info_queue = exception_info_queue  # shared mp.Queue(maxsize=1)
         self._metrics_queue = metrics_queue  # shared mp.Queue; drained by the trainer in `log()`
+        self._filtered_groups = filtered_groups  # shared mp.Queue; drained by the trainer in `_save_checkpoint()`
         # Metric accumulators
         self._counters: dict[str, float] = defaultdict(float)
         self._rates: dict[str, list[float]] = defaultdict(lambda: [0.0, 0.0])
@@ -903,7 +907,7 @@ async def _score_loop(self, stop_event: asyncio.Event) -> None:
                 print_prompt_completions_sample(
                     prompts=[s.prompt for s in samples],
                     completions=[s.completion for s in samples],
-                    rewards={"reward": [s.metrics["reward"] for s in samples]},
+                    rewards={"reward": [s.metrics["reward_kept"] for s in samples]},
                     advantages=[s.advantage for s in samples],
                     step=self._total_groups_scored,
                     num_samples=self.num_completions_to_print,
@@ -1215,6 +1219,13 @@ async def _score_group(self, group: RolloutGroup) -> list[RolloutSample]:
             reward_mean = reward_std = float("nan")
         logger.info(f"Rollout metrics: reward_mean={reward_mean:.4f}, reward_std={reward_std:.4f}")
 
+        self._rates["reward"][0] += float(np.nansum(rewards))
+        self._rates["reward"][1] += int(scored_mask.sum())
+        if not scored_mask.any() or reward_std == 0.0:
+            self._counters["rollout/zero_advantage_groups_total"] += 1
+            self._filtered_groups.put_nowait(group.group_id)
+            return []
+
         total_calls = sum(group.tool_call_counts)
         tool_metrics = (
             [
@@ -1245,7 +1256,7 @@ async def _score_group(self, group: RolloutGroup) -> list[RolloutSample]:
             )
         ):
             metrics = {
-                "reward": float(reward),
+                "reward_kept": float(reward),
                 "reward_std": reward_std,
                 **{
                     f"rewards/{name}": float(func_reward)
@@ -1325,6 +1336,7 @@ def __init__(
         # Metrics the child measures, drained by the trainer in `log()`. Bounded and dropped-on-full in the child, so a
         # trainer that stops draining can never block generation.
         self.metrics_queue = ctx.Queue(maxsize=4096)
+        self.filtered_groups = ctx.Queue()
         # Forwarded verbatim to _AsyncRolloutLoop in the child. queue_maxsize is also
         # forwarded — the child reads it for "rollout buffer full" log lines.
         loop_kwargs["queue_maxsize"] = queue_maxsize
@@ -1375,6 +1387,7 @@ def start(self) -> None:
                 self._failed_event,
                 self._exception_info_queue,
                 self.metrics_queue,
+                self.filtered_groups,
             ),
             name="grpo-rollout-worker-child",
             daemon=True,