feat(sc): support drop_incomplete_targets_on_restore in in_order

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
 

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
 

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
 

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)

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.

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.

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,13 +276,19 @@ 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(
             {
                 "state": state,
                 "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,

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)