diff --git a/docs/guides/single-controller.md b/docs/guides/single-controller.md index 401b7dc78d8..d5ba74d4187 100644 --- a/docs/guides/single-controller.md +++ b/docs/guides/single-controller.md @@ -95,6 +95,7 @@ Field definitions: - `max_buffered_rollouts` — hard cap on unconsumed rollout groups buffered in the data plane. Validated at setup against the gated sampler's required capacity; a value too small deadlocks the rollout pump, so setup raises instead of silently blocking. Sized from the widest window the run ever uses, so `warmup_lookahead_versions` rather than `max_lookahead_versions` when it is set. - `min_groups_for_streaming_train` — minimum ready groups the trainer waits for before dispatching a batch. Set to `num_prompts_per_step` for sync/legacy semantics; lower for streaming. (PPO) Must equal `num_prompts_per_step` — the critic has no split train API, so one `train_from_meta` call is one optimizer step, and streaming a step across chunks would step the critic once per chunk. - `sampler.warmup_lookahead_versions` (PPO) — lookahead used while `ppo.policy_training_start_step` critic warmup is in progress, shrinking back to `max_lookahead_versions` afterwards. The SC equivalent of `ppo.async_ppo.warmup_generation_lead_steps`. +- `drop_incomplete_targets_on_restore` — needs a sampler that stamps the target step: `in_order`, or a `custom` sampler that stamps; setup raises under the other built-ins. On resume a target step holding fewer than `num_prompts_per_step` groups is gap-filled by default: the rollout pump dispatches only the missing prompts and drops the rest of that dataloader batch. `true` discards the restored groups and dispatches the step whole instead. Neither regenerates the original prompts. The SC equivalent of `ppo.async_ppo.drop_incomplete_targets_on_restore`. ## Implementation Structure @@ -170,6 +171,7 @@ SC reads its async knobs from `async_rl:` and **requires `grpo.async_grpo: null` | *(no legacy equivalent — matches legacy full-batch train semantics)* | `min_groups_for_streaming_train: ${grpo.num_prompts_per_step}`, or `${ppo.num_prompts_per_step}` on a PPO run | | *(no legacy equivalent — matches legacy `max_trajectory_age + 1` batches in flight)* | `max_inflight_prompts: num_prompts_per_step × (max_lookahead_versions + 1)` | | *(no legacy equivalent — legacy sizes its buffer to `num_prompts_per_step × max_trajectory_age_steps × 2`)* | `max_buffered_rollouts: num_prompts_per_step × (max_lookahead_versions + 1)` (tight; see the [Config → behavior map](#config--behavior-map) for per-sampler values) | +| `drop_incomplete_targets_on_restore` | `drop_incomplete_targets_on_restore` (same; needs a target-stamping sampler such as `in_order`) | ## Known Missing Features diff --git a/examples/configs/grpo_math_1B_megatron_single_controller.yaml b/examples/configs/grpo_math_1B_megatron_single_controller.yaml index ce161cf3ac9..2df11464136 100644 --- a/examples/configs/grpo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/grpo_math_1B_megatron_single_controller.yaml @@ -21,6 +21,9 @@ async_rl: max_inflight_prompts: ${grpo.num_prompts_per_step} # Cap on unconsumed rollout groups buffered in the DataPlane (backpressure). max_buffered_rollouts: 64 + # in_order only. On resume, drop a target step's restored groups when it is + # short of a full batch instead of gap-filling it from subsequent prompts. + drop_incomplete_targets_on_restore: false # Enable per-rollout diagnostic prints (prompt content / completion previews). diagnostics: false diff --git a/examples/configs/ppo_math_1B_megatron_single_controller.yaml b/examples/configs/ppo_math_1B_megatron_single_controller.yaml index f94bcc7a39e..fd52dd6cda6 100644 --- a/examples/configs/ppo_math_1B_megatron_single_controller.yaml +++ b/examples/configs/ppo_math_1B_megatron_single_controller.yaml @@ -33,6 +33,9 @@ async_rl: max_inflight_prompts: ${ppo.num_prompts_per_step} # Cap on unconsumed rollout groups buffered in the DataPlane (backpressure). max_buffered_rollouts: 64 + # in_order only. On resume, drop a target step's restored groups when it is + # short of a full batch instead of gap-filling it from subsequent prompts. + drop_incomplete_targets_on_restore: false # Enable per-rollout diagnostic prints (prompt content / completion previews). diagnostics: false diff --git a/nemo_rl/algorithms/async_utils/replay_buffer.py b/nemo_rl/algorithms/async_utils/replay_buffer.py index f5061aa75ef..d83399d22a3 100644 --- a/nemo_rl/algorithms/async_utils/replay_buffer.py +++ b/nemo_rl/algorithms/async_utils/replay_buffer.py @@ -1036,13 +1036,17 @@ async def load_state_dict( max_groups: int, expected_partition_id: str, expected_group_size: int, + groups_per_step: int, + drop_incomplete_targets_on_restore: bool, ) -> int: """Validate and re-put checkpointed groups into the buffer. The preflight runs entirely before any DataPlane write (legacy precedent: validate, then truncate): 1. Validate the envelope and raise ValueError on malformed state. - 2. Truncate to ``max_groups``, keeping the freshest groups, so the + 2. Under ``drop_incomplete_targets_on_restore``, drop every group + stamped for a target step that is short of ``groups_per_step``. + 3. Truncate to ``max_groups``, keeping the freshest groups, so the restored count can never exceed the buffer's capacity. Groups carrying a ``target_step`` are never truncated — an over-capacity in-order checkpoint raises instead (see Raises). @@ -1061,6 +1065,13 @@ async def load_state_dict( expected_group_size: num_generations_per_prompt; every group must hold exactly this many rows (a changed group size silently breaks the group-relative baseline). + groups_per_step: num_prompts_per_step; the group count that makes a + target step whole. Read only when dropping incomplete targets. + drop_incomplete_targets_on_restore: Drop the restored groups of a + target step that is short of a full batch, so the rollout pump + dispatches that step from subsequent dataloader prompts instead + of gap-filling the hole. The original prompts are not + regenerated. Groups carrying no target_step are never dropped. Returns: Number of groups restored into the buffer. @@ -1120,6 +1131,34 @@ async def load_state_dict( ) seen_sample_ids.add(sid) + num_dropped_incomplete = 0 + if drop_incomplete_targets_on_restore: + target_counts = Counter( + group["target_step"] + for group in groups + if group["target_step"] is not None + ) + incomplete = { + target + for target, count in target_counts.items() + if count < groups_per_step + } + if incomplete: + kept = [ + group for group in groups if group["target_step"] not in incomplete + ] + num_dropped_incomplete = len(groups) - len(kept) + groups = kept + print( + "Dropping partially restored target step(s) " + + ", ".join( + f"{target}={target_counts[target]}/{groups_per_step}" + for target in sorted(incomplete) + ) + + "; the rollout pump refills them from subsequent prompts", + flush=True, + ) + if state["saved_capacity"] != max_groups: print( "TQReplayBuffer capacity changed: " @@ -1165,6 +1204,11 @@ async def load_state_dict( self._group_ids.append(group["group_id"]) summary = f"📦 Restored {len(groups)} replay group(s) from checkpoint" + if num_dropped_incomplete: + summary += ( + f"; dropped {num_dropped_incomplete} group(s) from incomplete " + "target step(s)" + ) if num_truncated: summary += f"; truncated {num_truncated} group(s) over capacity" print(summary, flush=True) diff --git a/nemo_rl/algorithms/single_controller.py b/nemo_rl/algorithms/single_controller.py index 88f6183b7d1..be028658e54 100644 --- a/nemo_rl/algorithms/single_controller.py +++ b/nemo_rl/algorithms/single_controller.py @@ -410,6 +410,9 @@ async def _maybe_restore_replay_buffer(self) -> None: Skipped with a warning when the checkpoint was written under a different sampler: restored groups carry the saving sampler's weight/target-step stamps, which another policy may never select. + + A target step restored short of a full batch is gap-filled by the rollout + pump, unless async_rl.drop_incomplete_targets_on_restore drops it instead. """ if self._last_checkpoint_path is None: return @@ -442,6 +445,10 @@ async def _maybe_restore_replay_buffer(self) -> None: max_groups=self._async_cfg.max_buffered_rollouts, expected_partition_id=self._partition_id, expected_group_size=self._algo_cfg.num_generations_per_prompt, + groups_per_step=self._algo_cfg.num_prompts_per_step, + drop_incomplete_targets_on_restore=( + self._async_cfg.drop_incomplete_targets_on_restore + ), ) # Each buffered group holds one _buffer_capacity permit; the load # truncation guarantees restored <= capacity, so this never blocks. diff --git a/nemo_rl/algorithms/single_controller_utils/config.py b/nemo_rl/algorithms/single_controller_utils/config.py index ce28051b297..c5365ee35da 100644 --- a/nemo_rl/algorithms/single_controller_utils/config.py +++ b/nemo_rl/algorithms/single_controller_utils/config.py @@ -420,6 +420,9 @@ class AsyncRLConfig(BaseModel, extra="allow"): max_inflight_prompts: int = 32 # Cap on unconsumed rollout groups buffered in the DataPlane (backpressure). max_buffered_rollouts: int = 64 + # in_order only (setup raises otherwise): on resume, drop the restored groups + # of a target step short of a full batch instead of gap-filling it. + drop_incomplete_targets_on_restore: bool = False # Enable per-rollout diagnostic prints (prompt content / completion previews). diagnostics: bool = False @@ -931,6 +934,18 @@ def validate_single_controller_config(master_config: MasterConfig) -> None: "loss_fn.force_on_policy_ratio=false so prev_logprobs are used" ) + # Only a stamped target step can be short of a batch. "custom" is left alone + # as in _validate_failure_settings: only its author knows whether it stamps. + if ( + async_config.drop_incomplete_targets_on_restore + and async_config.sampler.name not in ("in_order", "custom") + ): + raise NotImplementedError( + f"async_rl.sampler.name={async_config.sampler.name!r} stamps no target " + "step, so async_rl.drop_incomplete_targets_on_restore would change " + "nothing about the resume. Remove it, or switch to the in_order sampler." + ) + # Top-k retention keys off checkpointing.metric_name, but SC has no # validation loop yet (see _save_checkpoint), so a "val:" metric would # never be collected and top-k would silently degrade to a no-op. diff --git a/tests/unit/single_controller/test_checkpointing.py b/tests/unit/single_controller/test_checkpointing.py index 481df1ee695..d71c83a1ab6 100644 --- a/tests/unit/single_controller/test_checkpointing.py +++ b/tests/unit/single_controller/test_checkpointing.py @@ -50,7 +50,11 @@ import yaml from torchdata.stateful_dataloader import StatefulDataLoader -from nemo_rl.algorithms.async_utils.staleness_sampler import WindowedSamplerConfig +from nemo_rl.algorithms.async_utils.staleness_sampler import ( + InOrderSamplerConfig, + SamplerConfig, + WindowedSamplerConfig, +) from nemo_rl.algorithms.grpo import ( GRPOConfig, GRPOSaveState, @@ -272,6 +276,8 @@ async def load_state_dict( max_groups: int, expected_partition_id: str, expected_group_size: int, + groups_per_step: int, + drop_incomplete_targets_on_restore: bool, ) -> int: self.load_calls.append( { @@ -279,6 +285,10 @@ async def load_state_dict( "max_groups": max_groups, "expected_partition_id": expected_partition_id, "expected_group_size": expected_group_size, + "groups_per_step": groups_per_step, + "drop_incomplete_targets_on_restore": ( + drop_incomplete_targets_on_restore + ), } ) return self.load_return @@ -321,18 +331,27 @@ def _actor_master_config( checkpoint_must_save_by: Optional[str] = None, ft_save_period: Optional[int] = None, num_prompts_per_step: int = 2, + num_generations_per_prompt: int = 2, max_num_epochs: int = 1, + sampler: Optional[SamplerConfig] = None, + drop_incomplete_targets_on_restore: bool = False, ) -> MasterConfig: """MasterConfig for in-process SingleControllerActor tests. All fields are populated (init_tmp_checkpoint dumps the whole config to config.yaml); values satisfy validate_single_controller_config. """ - sampler_cfg = WindowedSamplerConfig(max_staleness_versions=1) + sampler_cfg = ( + sampler + if sampler is not None + else WindowedSamplerConfig(max_staleness_versions=1) + ) return MasterConfig.model_construct( policy={ # One optimizer.step per RL step: prompts * generations == gbs. - "train_global_batch_size": num_prompts_per_step * 2, + "train_global_batch_size": ( + num_prompts_per_step * num_generations_per_prompt + ), "generation": {"colocated": {"enabled": False}}, }, loss_fn=ClippedPGLossConfig(), @@ -342,7 +361,7 @@ def _actor_master_config( max_num_steps=max_num_steps, max_num_epochs=max_num_epochs, num_prompts_per_step=num_prompts_per_step, - num_generations_per_prompt=2, + num_generations_per_prompt=num_generations_per_prompt, seed=42, ), logger={ @@ -371,6 +390,7 @@ def _actor_master_config( min_groups_for_streaming_train=1, max_inflight_prompts=4, max_buffered_rollouts=4, + drop_incomplete_targets_on_restore=drop_incomplete_targets_on_restore, ), ) @@ -469,6 +489,17 @@ async def _main(): return asyncio.run(_main()) +def _run_buffer_restore(mc: MasterConfig, actor_args: SingleControllerActorArgs): + """Construct the actor and await only the replay-buffer restore.""" + + async def _main(): + actor = _ACTOR_CLS(mc, actor_args, SetupTimingMetrics()) + await actor._maybe_restore_replay_buffer() + return actor + + return asyncio.run(_main()) + + def _run_restore_then_train_pump( mc: MasterConfig, actor_args: SingleControllerActorArgs ): @@ -1301,6 +1332,8 @@ def test_run_restores_replay_buffer_and_permits(self, tmp_path): "max_groups": 4, "expected_partition_id": _PARTITION_ID, "expected_group_size": 2, + "groups_per_step": 2, + "drop_incomplete_targets_on_restore": False, } ] # Each restored group holds one _buffer_capacity permit. @@ -1309,6 +1342,38 @@ def test_run_restores_replay_buffer_and_permits(self, tmp_path): # run()'s finally must tear the synchronizer down exactly once. assert actor._weight_synchronizer.shutdown_count == 1 + def test_the_drop_flag_is_forwarded_to_the_buffer(self, tmp_path): + """The knob is in_order-only, so the restore must carry it, not default it.""" + ckpt_dir = tmp_path / "resume_ckpt" + ckpt_dir.mkdir() + envelope = {"groups": ["g0"]} + torch.save(envelope, ckpt_dir / "replay_buffer.pt") + # Generations != prompts so a swap of the two group counts is visible. + mc = _actor_master_config( + tmp_path, + max_num_steps=0, + num_generations_per_prompt=4, + sampler=InOrderSamplerConfig(max_lookahead_versions=1), + drop_incomplete_targets_on_restore=True, + ) + save_state = _initial_grpo_save_state() + save_state.sampler_name = "in_order" + buffer = _FakeTQBuffer(load_return=1) + + _run_buffer_restore( + mc, + _make_actor_args( + tq_buffer=buffer, + last_checkpoint_path=str(ckpt_dir), + save_state=save_state, + ), + ) + + assert len(buffer.load_calls) == 1 + assert buffer.load_calls[0]["drop_incomplete_targets_on_restore"] is True + assert buffer.load_calls[0]["expected_group_size"] == 4 + assert buffer.load_calls[0]["groups_per_step"] == 2 + def test_restored_permits_are_released_by_a_live_pump(self, tmp_path): # The restore takes one capacity permit per group; a running pump must # give them all back. Every other restore test uses max_num_steps=0, diff --git a/tests/unit/single_controller/test_resiliency_config.py b/tests/unit/single_controller/test_resiliency_config.py index e3408612ffa..45a2e68a1ef 100644 --- a/tests/unit/single_controller/test_resiliency_config.py +++ b/tests/unit/single_controller/test_resiliency_config.py @@ -712,3 +712,39 @@ def test_a_custom_sampler_is_not_second_guessed(self): rollout_failure={"max_consecutive_dropped_prompts": 2}, ) validate_single_controller_config(cfg) + + +class TestDropIncompleteTargetsOnRestore: + """The knob needs a stamping sampler: nothing else has a target step to drop.""" + + def test_default_is_off(self): + assert AsyncRLConfig().drop_incomplete_targets_on_restore is False + + @pytest.mark.parametrize( + "sampler", + [ + {"name": "windowed"}, + {"name": "weight_fifo"}, + {"name": "ready_first"}, + ], + ids=lambda sampler: sampler["name"], + ) + def test_rejected_under_every_built_in_but_in_order(self, sampler): + cfg = _master_config( + sampler=sampler, + drop_incomplete_targets_on_restore=True, + ) + with pytest.raises(NotImplementedError, match="stamps no target step"): + validate_single_controller_config(cfg) + + def test_accepted_under_in_order(self): + cfg = _master_config(drop_incomplete_targets_on_restore=True) + validate_single_controller_config(cfg) + + def test_a_custom_sampler_is_not_second_guessed(self): + """Whether it stamps is its author's to know, as with the drop budget.""" + cfg = _master_config( + sampler={"name": "custom", "target": "some_module:SomeSampler"}, + drop_incomplete_targets_on_restore=True, + ) + validate_single_controller_config(cfg) diff --git a/tests/unit/single_controller/test_tq_replay_buffer.py b/tests/unit/single_controller/test_tq_replay_buffer.py index 80081cbdd8c..1ad62a0cf36 100644 --- a/tests/unit/single_controller/test_tq_replay_buffer.py +++ b/tests/unit/single_controller/test_tq_replay_buffer.py @@ -35,6 +35,9 @@ # Each record yields _N_GENS training rows. _N_GENS = 2 +# Groups that make one target step whole (production: num_prompts_per_step). +# Deliberately != _N_GENS so the two thresholds cannot be confused for each other. +_GROUPS_PER_STEP = 3 def _stub_record_to_train_batch( @@ -618,6 +621,8 @@ def _load( max_groups: int = 8, expected_partition_id: str = "rollout_data", expected_group_size: int = _N_GENS, + groups_per_step: int = _GROUPS_PER_STEP, + drop_incomplete_targets_on_restore: bool = False, ) -> int: return _run( buf.load_state_dict( @@ -625,6 +630,8 @@ def _load( max_groups=max_groups, expected_partition_id=expected_partition_id, expected_group_size=expected_group_size, + groups_per_step=groups_per_step, + drop_incomplete_targets_on_restore=drop_incomplete_targets_on_restore, ) ) @@ -877,3 +884,71 @@ def test_target_stamped_groups_within_capacity_load_fine(self): assert _load(buf, state, max_groups=2) == 2 assert buf.target_step_list == [1, 2] + + +class TestTQReplayBufferLoadDropIncompleteTargets: + """``drop_incomplete_targets_on_restore`` discards partially restored steps.""" + + @staticmethod + def _partial_envelope() -> dict[str, Any]: + # Target 1 is whole (_GROUPS_PER_STEP groups); target 2 lost one to the + # in-flight cutoff at save time, so two of its three groups are here. + # Two rather than one so that a threshold mutated to expected_group_size + # (== _N_GENS == 2) stops dropping it. + return _make_envelope( + [ + _make_group_entry("g0", weight=1, target_step=1), + _make_group_entry("g1", weight=1, target_step=1), + _make_group_entry("g2", weight=1, target_step=1), + _make_group_entry("g3", weight=2, target_step=2), + _make_group_entry("g4", weight=2, target_step=2), + ], + saved_capacity=8, + ) + + def test_incomplete_target_dropped_and_complete_target_kept(self): + dp = FakeDataPlaneClient() + buf = _make_buffer(dp) + + restored = _load( + buf, self._partial_envelope(), drop_incomplete_targets_on_restore=True + ) + + assert restored == 3 + assert buf.target_step_list == [1, 1, 1] + put_sample_ids = [sid for c in dp.put_calls for sid in c["sample_ids"]] + assert "g3_g0" not in put_sample_ids + assert "g4_g0" not in put_sample_ids + + def test_incomplete_target_retained_by_default(self): + buf = _make_buffer(FakeDataPlaneClient()) + + assert _load(buf, self._partial_envelope()) == 5 + assert buf.target_step_list == [1, 1, 1, 2, 2] + + def test_unstamped_groups_are_never_dropped(self): + # Only a stamped step can be short of a batch; under a sampler that + # stamps nothing the knob must leave the whole checkpoint alone. + state = _make_envelope( + [_make_group_entry(f"g{w}", weight=w) for w in (1, 2, 3)], + saved_capacity=8, + ) + buf = _make_buffer(FakeDataPlaneClient()) + + assert _load(buf, state, drop_incomplete_targets_on_restore=True) == 3 + assert buf.target_step_list == [None, None, None] + + def test_dropping_runs_before_truncation(self): + # The drop is what brings this checkpoint back inside capacity, so it + # must happen first -- truncating a target-stamped envelope raises. + buf = _make_buffer(FakeDataPlaneClient()) + + restored = _load( + buf, + self._partial_envelope(), + max_groups=3, + drop_incomplete_targets_on_restore=True, + ) + + assert restored == 3 + assert buf.target_step_list == [1, 1, 1]