diff --git a/docs/source/async_grpo_trainer.md b/docs/source/async_grpo_trainer.md index deb22caaa47..0da72f4d733 100644 --- a/docs/source/async_grpo_trainer.md +++ b/docs/source/async_grpo_trainer.md @@ -363,6 +363,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 | +| `frac_reward_zero_std`, `reward_kept` | fraction of groups dropped because every scored reward was equal, or fewer than two rollouts were scored (GRPO only counts the first case), 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 fewer than two 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 batch is unaffected, it still fills its `samples_per_step` from the kept groups, so a step takes more generation when many groups are dropped. A dropped group still counts toward an epoch, which stays one pass over the prompts. ### Tools diff --git a/tests/experimental/test_async_grpo_trainer.py b/tests/experimental/test_async_grpo_trainer.py index 61acadbb359..0290c4b7646 100644 --- a/tests/experimental/test_async_grpo_trainer.py +++ b/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 @@ -98,6 +99,7 @@ def __init__( samples_per_weight_sync: int = 10, fork_k: int = 1, candidate_sets: bool = False, + drop_every: int = 0, ): self.rollout_buffer = queue.Queue() # What a real worker forwards to its child; the trainer writes the adapter name here under adapter sync @@ -107,6 +109,7 @@ def __init__( self._samples_per_weight_sync = samples_per_weight_sync self._model_version = 0 self._fork_k = fork_k + self._drop_every = drop_every self._model = None self._sample_iter = self._make_sample_iter(tokenizer, dataset, num_generations) @@ -139,6 +142,9 @@ def _sample(self, tokenizer, ids, n_prompt): def _make_sample_iter(self, tokenizer, dataset, num_generations): for group_id, row in enumerate(itertools.cycle(dataset)): + if self._drop_every and group_id % self._drop_every == 0: + yield group_id # a zero-advantage group: the worker sends only its id + continue completions = [ [{"role": "assistant", "content": f"{row['completion'][0]['content']} {idx}"}] for idx in range(num_generations) @@ -766,6 +772,16 @@ 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_dropped_groups(self): + trainer = self._trainer_at_step_5({0, 1, 3, 4}, dataset_start_index=10) + trainer._dropped_groups = {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_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 @@ -2228,8 +2244,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). @@ -2238,19 +2254,70 @@ 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, [{}] * 3, "c0") seq_b = TrainingSequence([1, 2, 20], [0, 0, 1], [0, 0, -0.2], [[]] * 3, [{}] * 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, [{}] * 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) + + +def all_none_reward(completions, **kwargs): + return [None] * len(completions) + + +def one_scored_reward(completions, **kwargs): + return [None, 1.0] + + +class TestDropZeroAdvantageGroups(TrlTestCase): + @pytest.fixture(autouse=True) + def _init_accelerate_state(self): + PartialState() + + def _score(self, loop, reward_funcs): + for group_id, reward_func in enumerate(reward_funcs): + seq_a = TrainingSequence([1, 2, 10], [0, 0, 1], [0, 0, -0.1], [[]] * 3, [{}] * 3, "c0") + seq_b = TrainingSequence([1, 2, 20], [0, 0, 1], [0, 0, -0.2], [[]] * 3, [{}] * 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())) + items, metrics = [], defaultdict(list) + with contextlib.suppress(queue.Empty): + while True: + items.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 items, metrics + + @pytest.mark.parametrize( + "reward_func, reward", + [(constant_reward, (2.0, 2)), (all_none_reward, (0.0, 0)), (one_scored_reward, (1.0, 1))], + ) + def test_dropped_group_reaches_the_buffer_as_its_group_id(self, reward_func, reward): + items, metrics = self._score(_rollout_loop(), [reward_func, two_reward]) + assert items[0] == 0 + assert [s.group_id for s in items[1:]] == [1, 1] + assert metrics["frac_reward_zero_std"] == [(1.0, 1), (0.0, 1)] + assert metrics["reward"] == [reward, (4.0, 2)] # the dropped group still counts in the policy's reward @pytest.mark.skipif( @@ -2266,7 +2333,7 @@ class TestEpochStop(TrlTestCase): what actually regresses if the wiring breaks. """ - def _train(self, fork_k, num_train_epochs=2): + def _train(self, fork_k, num_train_epochs=2, drop_every=0): model_id = "trl-internal-testing/small-Qwen2ForCausalLM-2.5" dataset = load_dataset("trl-internal-testing/zen", "conversational_prompt_completion", split="train") args = AsyncGRPOConfig( @@ -2279,7 +2346,9 @@ def _train(self, fork_k, num_train_epochs=2): vllm_server_timeout=5.0, report_to="none", ) - worker = _StubRolloutWorker(AutoTokenizer.from_pretrained(model_id), dataset, num_generations=3, fork_k=fork_k) + worker = _StubRolloutWorker( + AutoTokenizer.from_pretrained(model_id), dataset, num_generations=3, fork_k=fork_k, drop_every=drop_every + ) trainer = AsyncGRPOTrainer( model=model_id, reward_funcs=dummy_reward_func, @@ -2307,6 +2376,13 @@ def test_epoch_stop_is_fork_independent(self): # stop in FEWER prompt-passes (the pre-fix bug). assert forked.state.global_step > no_fork.state.global_step + def test_epoch_stop_counts_dropped_groups(self): + trainer, num_prompts = self._train(fork_k=1, drop_every=2) + + consumed = trainer._trained_groups | trainer._dropped_groups + assert 2 * num_prompts <= len(consumed) < 3 * num_prompts + assert len(trainer._trained_groups) < 2 * num_prompts + def test_epoch_stop_saves_the_final_step(self): trainer, _ = self._train(fork_k=1) diff --git a/trl/experimental/async_grpo/async_grpo_config.py b/trl/experimental/async_grpo/async_grpo_config.py index 7de6aba8a8c..2913d43f070 100644 --- a/trl/experimental/async_grpo/async_grpo_config.py +++ b/trl/experimental/async_grpo/async_grpo_config.py @@ -450,6 +450,12 @@ def __post_init__(self): "sampled token with `top_k` (e.g. `top_k=512`)." ) + if self.num_generations < 2: + raise ValueError( + "GRPO requires at least 2 generations per prompt to calculate the advantages. You provided " + f"{self.num_generations}, which is less than the minimum required." + ) + if self.parallelism_config is not None and ( self.parallelism_config.cp_enabled or self.parallelism_config.sp_enabled ): diff --git a/trl/experimental/async_grpo/async_grpo_trainer.py b/trl/experimental/async_grpo/async_grpo_trainer.py index 66a5d5956eb..c7c2f9787e8 100644 --- a/trl/experimental/async_grpo/async_grpo_trainer.py +++ b/trl/experimental/async_grpo/async_grpo_trainer.py @@ -124,10 +124,12 @@ class RolloutWorkerProtocol(Protocol): Attributes: rollout_buffer (`queue.Queue` or `multiprocessing.queues.Queue`): - Queue the trainer drains; the worker pushes scored `RolloutSample`s onto it. The two queue types are - structurally identical (`get` / `put_nowait` / `qsize`) but nominally unrelated, so both are allowed: the - default [`AsyncRolloutWorker`] runs its loop in a spawned process and uses `multiprocessing.Queue`, while - an in-process worker uses `queue.Queue`. + Queue the trainer drains; the worker pushes scored `RolloutSample`s onto it, or the bare `int` group id of + a group it scored but has nothing to train on (every advantage zero), so the prompt still counts as + consumed for epochs and checkpoint resume. The two queue types are structurally identical (`get` / + `put_nowait` / `qsize`) but nominally unrelated, so both are allowed: the default [`AsyncRolloutWorker`] + runs its loop in a spawned process and uses `multiprocessing.Queue`, while an in-process worker uses + `queue.Queue`. metrics_queue (`queue.Queue` or `multiprocessing.queues.Queue`): Queue the trainer drains in `log()` for metrics the worker measured itself. Each item is one dict shaped like the trainer's metric sink — `{key: float}` for a gauge or a counter, `{key: (numerator, denominator)}` @@ -239,11 +241,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): @@ -252,8 +254,9 @@ 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) + trainer = self._trainer + consumed = trainer._groups_before_resume + len(trainer._trained_groups | trainer._dropped_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. @@ -310,6 +313,7 @@ def __init__( check_health_fn, stale_after_s, metrics, + dropped_groups, max_staleness=3, poll_interval_s=5.0, report_to=None, @@ -322,6 +326,7 @@ def __init__( # `dispatch_batches=True`), which is also the only process where the queue wait is real, so its metrics need no # communication at all — they are appended straight into the sink the trainer reduces in `log()`. self.metrics = metrics + self.dropped_groups = dropped_groups # Blocking-get time, accumulated here and flushed by `training_step` at the optimizer-step boundary — the only # place that knows when a step's worth of waiting is done. self.wait_s = 0.0 @@ -366,6 +371,10 @@ def __iter__(self): if wait_s > 1.0: logger.info(f"waited {wait_s:.1f}s for sample (qsize={self.queue.qsize()})") + if isinstance(sample, int): + self.dropped_groups.add(sample) + continue + version = self.model_version_fn() staleness = version - sample.model_version if staleness > self.max_staleness: @@ -1272,6 +1281,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._dropped_groups: set[int] = set() # Tracks restart to match `num_train_epochs` self._groups_before_resume = 0 self._epoch_stop_groups: int | None = None @@ -1461,6 +1471,7 @@ def get_train_dataloader(self) -> DataLoader: check_health_fn=self.rollout_worker.check_health, stale_after_s=self.args.heartbeat_stale_after_s, metrics=self._metrics["train"], + dropped_groups=self._dropped_groups, max_staleness=self.args.max_staleness, report_to=self.args.report_to, ) @@ -1941,9 +1952,9 @@ def _save_checkpoint(self, model, trial): 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._trained_groups | self._dropped_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} diff --git a/trl/experimental/async_grpo/async_rollout_worker.py b/trl/experimental/async_grpo/async_rollout_worker.py index d5a6288cc3e..b90674229c5 100644 --- a/trl/experimental/async_grpo/async_rollout_worker.py +++ b/trl/experimental/async_grpo/async_rollout_worker.py @@ -926,19 +926,20 @@ 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, ) self._total_groups_scored += 1 - for sample in samples: + for item in samples or [group.group_id]: t_blocked = None while True: try: - sample.enqueued_at = time.time() - self.rollout_buffer.put_nowait(sample) + if isinstance(item, RolloutSample): + item.enqueued_at = time.time() + self.rollout_buffer.put_nowait(item) break except queue.Full: if stop_event.is_set(): @@ -1253,10 +1254,10 @@ async def _score_group(self, group: RolloutGroup) -> list[RolloutSample]: rewards[all_nan_mask] = np.nan scored_mask = ~np.isnan(rewards) + scored = rewards[scored_mask] # NOTE: for NaN reward we set advantage to 0 ! advantages = np.zeros_like(rewards) if scored_mask.any(): - scored = rewards[scored_mask] advantages[scored_mask] = (scored - scored.mean()) / (scored.std() + 1e-8) reward_mean = float(scored.mean()) reward_std = float(scored.std()) @@ -1264,6 +1265,14 @@ 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}") + zero_std = bool(np.all(scored == scored[:1])) + self._rates["reward"][0] += float(scored.sum()) + self._rates["reward"][1] += len(scored) + self._rates["frac_reward_zero_std"][0] += zero_std + self._rates["frac_reward_zero_std"][1] += 1 + if zero_std: + return [] + total_calls = sum(group.tool_call_counts) tool_metrics = ( [ @@ -1294,7 +1303,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)