feat(sc): support drop_incomplete_targets_on_restore in in_order

[yuki-97]

What does this PR do ?

Ports drop_incomplete_targets_on_restore from the legacy async-PPO path to the SingleController path, so a resume can discard a partially restored target step instead of gap-filling it.

Legacy async PPO/GRPO has had this knob since #3410 (ppo.async_ppo.drop_incomplete_targets_on_restore, nemo_rl/algorithms/async_utils/replay_buffer.py:531). SC had no equivalent: a resume always gap-filled. This adds async_rl.drop_incomplete_targets_on_restore (default false, so existing runs are byte-identical).

What a "partially restored target step" is. In-flight rollouts are not written to the checkpoint, so a step whose batch was still filling comes back with fewer than num_prompts_per_step groups stamped for it.

• false (default) — gap-fill. _rollout_pump dispatches only the missing prompts for that step and drops the rest of that dataloader batch (nemo_rl/algorithms/single_controller.py:630-632).
• true — drop. The restored groups of that step are discarded at load time and the step is dispatched whole from subsequent prompts.

Neither setting regenerates the original prompts, and both consume exactly one dataloader batch per admitted step.

A complete restored target step is already skipped — no new branch is needed for it. The same clamp does both jobs: num_prompts = max(0, prompt_batch.size - buffered). When the step is full, num_prompts is 0, the dispatch loop body never runs, and that dataloader batch is consumed without generating anything. This PR only changes what buffered is for incomplete steps. Two things worth knowing about that skip, both pre-existing: the skipped batch's prompts are discarded (the dataloader is not rewound, matching legacy's last_target_weight_already_generated skip), and admit is still called and still advances the dispatch index, which it must — otherwise the stamp sequence drifts one step off the trainer version.

in_order only. validate_single_controller_config raises NotImplementedError under any other sampler, custom included. in_order is the only sampler that stamps a target_step; under windowed / weight_fifo / ready_first every group is fungible across steps, so no restored group belongs to a step that could be incomplete and the flag would silently change nothing. custom is rejected because setup cannot tell whether it stamps without importing it.

Ordering: the drop runs before the over-capacity truncation in load_state_dict. Truncating a target-stamped envelope raises, and the drop can bring an over-capacity checkpoint back under max_buffered_rollouts — so it must come first.

Issues

None closed. Part of the SingleController parity work tracked in #2625.

Usage

async_rl:
  sampler:
    name: in_order
    max_lookahead_versions: 1
  # Discard a target step restored short of a full batch and dispatch it whole,
  # instead of gap-filling it from subsequent prompts.
  drop_incomplete_targets_on_restore: true

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• Stacked on yukih/sc-ppo — base branch is that PR, not main.

• No new deadlock surface. Dropped groups never acquire a _buffer_capacity permit (load_state_dict returns the post-drop count), and per-step capacity pressure is identical either way: gap-fill holds k restored permits and dispatches batch - k; drop holds 0 and dispatches batch. Dataloader consumption is also identical, so _clamp_max_num_steps's step budget is unaffected. Dropping strictly reduces the restored group count, so it can only make the over-capacity raise fire less often.

• One caveat. _batch_shortfall is not checkpointed, so a GRPO step that was legitimately shortened by on_dropped_prompt="shrink" reads as incomplete after a resume; with the flag on it is re-dispatched in full. That is extra rollout work, not a stall. PPO is unaffected — setup rejects drop budgets there.

• Tests: tests/unit/single_controller/test_tq_replay_buffer.py (drop / unstamped groups untouched / drop-before-truncate), test_resiliency_config.py (accepted under in_order, raises under the other four), test_checkpointing.py (the flag reaches the buffer).

  uv run --group test pytest tests/unit/single_controller/test_tq_replay_buffer.py tests/unit/single_controller/test_resiliency_config.py tests/unit/single_controller/test_checkpointing.py

• Test results. Yellow is the legacy path with the same settings, as a baseline. Red is v2. Blue is v2 resumed from checkpoint 20. Green is v2 with this PR, resumed from checkpoint 20 with drop_incomplete_targets_on_restore: true.
  

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[yuki-97]

/ok to test 4d45be5

[yuki-97]

Team review (5 agents: RL expert, bug finder, test agent, design reviewer, devil's advocate).

No correctness findings. The deadlock/liveness audit came back clean: peak buffer occupancy is N*(L+1) <= C and is identical with the flag on or off (gap-fill holds k restored permits and dispatches N-k; drop holds 0 and dispatches N), which rests on drop_last=True on the dataloader so every batch is exactly num_prompts_per_step. Dropping only ever reduces held permits. The sibling-inheritance check against the legacy async-PPO implementation found three structural deltas, all inherited or outcome-identical. Drop-before-truncate is right and the ordering is enforced by an executable test, not just the docstring.

Four suggestions below, all tests and docs.

Two caveats on this review: nothing was executed (uv could not run in this environment), so the mutation-survival claims are hand-traced against the source rather than run; and the perf evidence attached to the PR is a PPO run, where the shrink caveat you disclosed cannot fire by construction -- not asking for more, just noting the GRPO x shrink path has no run behind it.

Generated by Claude Code

[yuki-97]

/ok to test 2fcc765

[terrykong]

note for those reading that we'll probably focus on getting #3480 in so we can restore 100% accurately instead of dropping