[RayenTian]

What does this PR do ?

Extends single-controller dynamic sampling (#4588, which supports only in_order) to the ready_first and windowed samplers, and makes weight_fifo fail fast at setup because it cannot work with dynamic sampling.

This PR is stacked on #4588 (base branch ruit/dapo_dynamic_sampling_sc); review that first.

Issues

#1365

Background

Dynamic sampling drops prompt groups whose generations all got the same reward, then resamples a replacement so each step keeps its batch size. In #4588 the replacement and shortfall accounting is keyed on the target_step a sampler stamps on each slot, and only InOrderSampler stamps one. Every other sampler was rejected by config validation.

What this PR supports

async_rl.sampler.name	Dynamic sampling	Spare pool	A group lost for good (no spare left / budget spent)
in_order	Supported (unchanged from #4588)	Yes	Shrinks the step the group was stamped for.
ready_first	New	Yes	Shrinks the step being assembled at that moment (or the next one if it already finished selecting). Batch boundaries can shift, but the total consumed is right.
windowed	New	No	Nothing: rollout keeps producing, so later groups make up for it. Data runs out sooner, so the run can end before len(dataloader) steps; the last step closes on the groups left, subject to min_step_batch_fraction. on_dropped_prompt: replace is not required.
weight_fifo	Rejected at setup (see below)	n/a	n/a
custom	Allowed. Without DropAccountingSampler it gets the gated defaults: spare pool on, no short close at exhaustion, a lost group charged to its own target_step (or the step being assembled if the sampler stamps none). A stamping sampler therefore behaves like in_order, a non-stamping one like ready_first. An ungated, continuously producing sampler should implement the protocol with uses_replacement_pool = False and closes_short_at_exhaustion = True.

Dynamic sampling stays GRPO-only (PPO is rejected, as in #4588).

Result

How it works

• New optional DropAccountingSampler protocol (@runtime_checkable, so existing custom samplers are unaffected). A sampler declares uses_replacement_pool, closes_short_at_exhaustion, and charge_step_for_lost_group(target_step, assembling_version) (which step a permanently lost group is charged to, or None). BaseSampler supplies defaults matching the gated, unstamped behaviour; InOrderSampler overrides them explicitly with its current behaviour; WindowedSampler overrides them for the no-pool / close-short case.
• The controller tracks the step currently being assembled (_assembling_version) and charges groups from unstamped samplers to it. State (shortfalls, spare pool, replacement counts) stays in the controller; no new checkpointed state, and the checkpoint field sampler_stamps_target_steps keeps its name for compatibility.
• windowed closes its last step on whatever groups are left once rollout is exhausted, instead of raising "rollout exhausted before a complete training step".
• Per-step reward metrics (unfiltered_reward, zero_variance_fraction, ...) now also work for unstamped samplers: RolloutManager takes a step provider (the assembling step) and uses it when a group has no target_step. For unstamped samplers this is approximate, since a group is tallied under the step being assembled when it finishes.
• Config validation: dropped the sampler restriction on dynamic sampling and the hard error for drop budgets under unstamped samplers. on_dropped_prompt: replace is still required, except for windowed.

Why weight_fifo is rejected

weight_fifo builds a step only from groups that share one start_weight: select picks the oldest in-window version and waits for that version alone to fill the step. Dynamic sampling breaks this:

1. A group is filtered after it finishes generating, so its replacement is dispatched later, under a newer weight version.
2. The replacement lands in the newer version's bucket, so the older bucket is permanently short by one group per replacement.
3. The gate is closed (the trainer cannot advance until this step trains) and nothing is in flight, so nothing can ever top the old bucket up. Training hangs.

Observed on a real run (weight_fifo, max_staleness_versions=1, 8 prompts/step), at the stall:

trainer_v=2 dispatch_index=3 gate_window=1 inflight=0 buffer_len=16
per_start_weight={1: '6 ready/6 slots', 2: '10 ready/10 slots'} spare_pool=73

16 groups were ready, enough for two steps, yet none could be selected because the oldest version had 6 of the 8 needed. The spare pool still had 73 prompts, so this is not budget exhaustion. The same buffer contents are selectable under ready_first (reproduced against the real sampler classes with a fake buffer).

Making weight_fifo work would mean either letting a step mix weight versions (which is what ready_first already is) or stamping replacements with the old version (which mislabels freshly generated data). Neither is worth it, so setup now fails with:

grpo.use_dynamic_sampling cannot be combined with async_rl.sampler.name='weight_fifo'. weight_fifo builds each step only from groups generated under the same weight version, and a replacement for a filtered group is generated later, under a newer version. The step can then never fill, and training hangs. Set async_rl.sampler.name to 'in_order' or 'ready_first', or turn off grpo.use_dynamic_sampling.

Known limit: weight_fifo with on_dropped_prompt: replace and drops from infrastructure failures (not dynamic sampling) can hit the same bucket mismatch. That is not changed here.

Usage

grpo:
  use_dynamic_sampling: true
  num_generations_per_prompt: 16   # must be >= 2
  num_prompts_per_step: 8

async_rl:
  sampler:
    name: ready_first              # or in_order / windowed; NOT weight_fifo
    max_staleness_versions: 1      # ready_first / windowed; in_order uses max_lookahead_versions
  rollout_failure:
    on_dropped_prompt: replace     # not required for windowed
    max_replacement_attempts: 16   # >= 1
    replacement_reserve_prompts: 16
    min_step_batch_fraction: 1.0

Testing

• New and updated unit tests under tests/unit/single_controller/: sampler hooks (test_sampler_interface.py), spare-pool and shortfall behaviour for unstamped samplers (test_rollout_pump.py, test_single_controller_actor.py), per-step reward metrics (test_step_reward_metrics.py), config validation including the weight_fifo rejection (test_ppo_setup.py, test_resiliency_config.py).
• Not run locally: the unit tests have only been syntax-checked and pre-commit (ruff, pyrefly) passed. Relying on CI.
• The weight_fifo stall above was reproduced on a cluster run of the V2 DAPO recipe; the PR does not add a nightly test.

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 feat(sc): enable single-controller dynamic sampling #4588; the diff against that base is 13 files, mostly single_controller.py, staleness_sampler.py, config.py and tests.
• Docs: the Dynamic sampling section of docs/guides/single-controller.md now has the per-sampler table and the weight_fifo note.

🤖 Generated with Claude Code

[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.

[RayenTian]

/ok to test 4f9bbf6

[RayenTian]

/ok to test 4fae850

[RayenTian]

/ok to test db5f343

[svcnvidia-nemo-ci]

The optional sampler protocol is a useful separation between gated and continuously producing policies. I inspected all 16 changed files and traced admission, shortfall accounting, training claims, checkpoint restoration, and reward tallies.

Critical (1)

One correctness issue remains: losses arriving during a nonfinal training chunk can over-credit the current step, leaving the following step short without accounting for it.

Suggestions (1)

Quantitative training evidence would help validate the newly supported sampling policies.

Validation was static only: no tests or linters were executed, and no delegated independent review was available.

Note: formal approval is not available for this request, so this review is posted as a comment only.

Review generated by codex. This was a full review. Request a full review with /review reset=true.

[svcnvidia-nemo-ci]

1 action item — please fix in this PR.

TL;DR — multiple losses during a nonfinal chunk can shrink its target below the number already selected, discarding accounting needed by the next step.

AI-1 — [BUG] Prevent shortfall credits from crossing selected work

PR-introduced; confidence: 97. In _credit_shortfall, every unstamped loss still goes to the same assembling step until the train pump advances it.

Minimal case: ready-first, four prompts per step, streaming minimum three, lookahead one, shrink mode, two tolerated drops, and batch fraction 0.5:

1. Select three groups for step 0. The selection-time check leaves the assembling version at 0 because the target is still four.
2. While that chunk trains asynchronously, two lookahead prompts are permanently dropped. Both credits land on step 0, reducing its target to two despite three groups already being selected.
3. The loop exits, then prunes both credits. With two admitted batches, only three groups remain, but step 1 still requires four and raises “rollout exhausted before a complete training step was assembled.” Expected: two legal three-group steps, with one loss charged to each.

Action: Coordinate this credit transition with selection progress in the train pump. Track groups already selected, including the currently training chunk; once a loss consumes the final unselected slot, advance unstamped accounting before another loss can be charged there. Also reconcile losses arriving while selection awaits its training claim. The invariant is that a closed step's selected groups plus charged losses equal its configured size—credits must never excuse already-selected work. This spans the selection/credit transition, so I cannot safely offer a one-line replacement.

Preserve: Keep stamped in-order losses on their stamped step, windowed losses uncharged until exhaustion, and custom samplers' explicit charge decisions intact. Retain next-step charging for a loss during an already-final chunk, as asserted by the existing regression.

Acceptance checks (proposed, not executed): Drive the two-batch case above with two drops during the three-group chunk: expect shortfalls {0: 1, 1: 1} before cleanup, two optimizer steps consuming three groups each, and normal exhaustion. Repeat with only one drop: step 0 closes at three and step 1 remains full. Retain the existing final-chunk assertion that the loss remains under step 1, plus the stamped, windowed, and custom charging assertions.

[svcnvidia-nemo-ci]

1 action item — suggestion.

AI-1 — [PERF-EVIDENCE] Add quantitative validation for the supported policies

PR-introduced validation gap; confidence: 100. The newly supported training policies change batch boundaries, the number of completed updates, and which generated rewards enter training. The supplied run evidence diagnoses the rejected weight-FIFO stall, but does not measure throughput or convergence for the supported paths; this is not a demonstrated performance regression.

Action: Attach reproducible small-run results comparing dynamic sampling under in-order, ready-first, and windowed: resolved configs/seeds, generated and trained group counts, dropped/replaced counts, groups or tokens per second, and reward curves against generated/trained samples.

Preservation and acceptance: Keep the existing in-order baseline and use comparable model/data and generation budgets; do not compare only equal step counts because windowed can finish fewer updates. The evidence should show that the supported paths complete as configured and make their sampling cost and reward trajectory visible. This is an experimental-validation request, not a proposed code patch; no runs were executed in this review.