feat(sc): enable single-controller dynamic sampling by RayenTian · Pull Request #4588 · NVIDIA-NeMo/RL

RayenTian · 2026-10-08T06:23:43Z

What does this PR do ?

Enable DAPO-style dynamic sampling (grpo.use_dynamic_sampling) on the single-controller (SC) path for in-order sampler : zero-variance prompt groups are dropped and replaced with a fresh prompt, so every step still trains on num_prompts_per_step groups.

Issues

#1365

Result

Non token capture

V1 + Sync + dynamic sampling V.S. V2 + Async + dynamic sampling

token capture

V2 + Gym + token capture V.S. V2 + Gym + non token capture

Testing

Unit tests cover:

zero-variance detection and filtering on both the regular and token-capture paths
replacement, shrinking and the validation errors above
the per-step reward metrics

No nightly test is added in this PR. A complete V2 (single-controller) DAPO recipe will be added together with its nightly test once reward shaping is supported on the single-controller path. Until then, dynamic sampling here is covered by unit tests only.

Scope and limitations

Only the in_order sampler is supported in this PR (a custom sampler works only if it stamps a target step itself). With ready_first, weight_fifo or windowed, setup fails with a ValueError when use_dynamic_sampling is enabled. The other samplers will be supported in #4784

Why: replacement and shortfall accounting are keyed by each group's target_step, and only InOrderSampler stamps one. The other built-in samplers leave _GatedSampler._stamp returning None (windowed does not even advance its dispatch index). Without a stamped step, a filtered group cannot be credited short or replaced, so the gated samplers would wait on a group nobody is generating.

Usage

You can potentially add a usage example below

grpo:
  use_dynamic_sampling: true
  num_generations_per_prompt: 16   # must be >= 2, otherwise no group has reward spread
  num_prompts_per_step: 8

async_rl:
  sampler:
    name: in_order                 # the only sampler supported by this PR
    max_lookahead_versions: 1      # on-policy; the repo's SC default allows 1
  rollout_failure:
    on_dropped_prompt: replace     # required: resample a substitute for each filtered group
    max_replacement_attempts: 16   # per-slot cap, must be >= 1
    replacement_reserve_prompts: 16  # low-water mark of the spare-prompt pool, must be >= 1
    min_step_batch_fraction: 1.0   # 1.0 = every step must train on exactly num_prompts_per_step groups

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

...

copy-pr-bot · 2026-10-08T06:23:47Z

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.

cmunley1 · 2026-10-12T21:06:53Z

Tried this with Lightning 3.5 SingleController, ready_first, token_capture.enabled=true, GBS 4096. The filter never fired. Committed groups were 32.9% zero-variance, same as the run without it (32.7%), and dropped_prompt_groups stayed 0.

With token capture the pump calls generate_for_finalization and hands off to the finalizer actors (single_controller.py, around L1596). It never calls generate_and_push, where the check lives. Should the check move before the finalizer commit?

Also, ready_first stamps no step, so on_dropped_prompt=replace does nothing. Does a filtered group still use up an admission slot with no refund? At ~33% I'd expect the window to starve after a few steps. Has this run past 3 steps? Logging zero_variance_groups as a metric would make a no-op like this visible.

smahdavi4 · 2026-10-14T21:41:00Z

Can we also have a metric to keep track of the reward before dynamic sampling? it would be very useful because in dynamic sampling train/reward might stay flat while the model is still learning.

RayenTian · 2026-10-18T04:53:49Z

/ok to test 8e98c7c

RayenTian · 2026-10-19T04:36:00Z

/ok to test e55ff23

RayenTian · 2026-10-19T08:44:32Z

/ok to test 0b888b4

RayenTian · 2026-10-19T09:42:07Z

/ok to test 4a29761

RayenTian · 2026-10-20T01:23:59Z

Self review with agent is down. @pthombre would you please take a look at the token capture changes? @saumishr would you please also review the reserved prompt group changes? @yuki-97 would you please help review on the single controller part? Thank you all a lot!

yuki-97

Merge blockers — please rebase on main and resolve conflicts.

Thanks for bringing dynamic sampling to the single-controller path, and for the clear write-up of its scope and limitations.

Great to see the V1-vs-V2 and token-capture-vs-not runs posted in the description; the effect of the filter can be read straight off them.

Scope: the structure of the single-controller dynamic sampling path, at 4a29761. A detailed review will follow once the two comments below are addressed.

Remaining comments are on:

Moving token-capture support to a follow-up PR that borrows from later steps like the non-capture path
A dynamic-sampling budget of its own, counted separately from rollout fault tolerance

Addresses both review comments on #4588. Scope it to the non-capture path. On the token-capture path a replacement cannot borrow a finished group from a later step the way the non-capture path does, so every drop waits out a full fresh rollout with the trainer idle. Dynamic sampling makes drops routine, so that cost would land on every step. Setup now rejects use_dynamic_sampling together with token_capture.enabled, and FinalizedGroup.zero_variance goes away with the capture-side filtering. Mirroring the borrow on that path is its own change. Separate the budgets. A zero-variance drop is a data decision about a rollout that succeeded, not a rollout failure, so sharing the fault-tolerance knobs let each kind of drop redefine the other. generate_and_push now returns a distinct RolloutOutcome.FILTERED, counted against the new grpo.dynamic_sampling_max_replacements_per_slot rather than max_replacement_attempts. Filtering no longer reads on_dropped_prompt: a uniform group is never worth training on, so it is replaced under "shrink" too and the spare pool is filled either way. Exhausting the budget raises, as v1 does past dynamic_sampling_max_gen_batches, instead of crediting a shortfall against min_step_batch_fraction; an empty pool still shrinks, since that is the dataloader ending rather than a mis-tuned filter. The default budget of 32 is sized from the drop rate, not from how many resamples feel reasonable: a slot exhausts it with probability p**(N+1), and every slot of every step takes that draw. At 32 a 4096-group step over 1000 steps is quiet below roughly a 78% drop rate, well clear of the ~33% seen in practice, and raises above it. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: ruit <ruit@nvidia.com>

RayenTian · 2026-10-29T09:08:29Z

/ok to test 83f4879

…xpert IDs Signed-off-by: ruit <ruit@nvidia.com>

Signed-off-by: ruit <ruit@nvidia.com>

…ing, update docs Signed-off-by: ruit <ruit@nvidia.com>

…replay buffer A zero-variance group dropped on the token-capture path was already marked ready, and abort() ignores ready slots, so the sampler could still pick it after its rows were cleared. Remove the slot first. Also drop an unreachable config check and its test, remove a broken custom-sampler test, and tighten comments and docs. Co-Authored-By: Claude Code <noreply@anthropic.com> Signed-off-by: ruit <ruit@nvidia.com>

… mock - test_single_controller_actor: add pop_step_reward_metrics/filter_zero_variance_groups to the _train_pump_controller mock (the unconditional step-close call regressed 31 tests) - config: use per-algo (grpo/ppo) prefix in dynamic-sampling errors; note the 'custom' sampler and its stamping requirement in the sampler check message - rollout_manager: add docstring to the filter_zero_variance_groups property - test_ppo_setup: cover custom sampler acceptance (supplies the required target) Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: ruit <ruit@nvidia.com>

…mpler capability - Drop the align_reward_with_v1 reward/filtered_reward swap. On SC, `reward` stays the trained-rows mean (consistent whether or not dynamic sampling is on), and the pre-filter progress signal is carried by the always-reported unfiltered_reward / first_attempt_reward, rather than re-labeling an existing key to match the v1 path. - Gate _take_replacement on the sampler's _sampler_stamps_target_steps capability instead of inferring "no step is waiting on a replacement" from a per-call target_step==None. The sibling guards in _credit_shortfall/_promote_into_step keep their None checks (they use target_step as a key), so only _take_replacement changes. - Update docs and tests accordingly. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: ruit <ruit@nvidia.com>

… API Signed-off-by: ruit <ruit@nvidia.com>

…inventory Co-Authored-By: Claude Code <noreply@anthropic.com> Signed-off-by: ruit <ruit@nvidia.com>

…unter on zero-variance drops Two fixes on the single-controller dynamic-sampling path: - PPO + use_dynamic_sampling slipped through setup validation after use_dynamic_sampling was removed from the SC `unsupported` list: with on_dropped_prompt="replace", a stamping sampler and num_generations_per_prompt>=2 it validated, then failed a divisibility assert deep in the value workers (#2625) -- a dropped group shortens the step and the critic shards it against value.train_global_batch_size, which the separate PPO drop-budget guard does not catch. Reject it explicitly at setup; dynamic sampling is GRPO-only. - The non-capture zero-variance drop returned SKIPPED without resetting _consecutive_infra_drops, unlike a commit and unlike the finalizer path (which resets on a successful generation). A healthy generation that is policy-dropped now clears the consecutive-drop run so it cannot prematurely trip max_consecutive_dropped_prompts while the fleet is answering. Add tests covering both. Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com> Signed-off-by: ruit <ruit@nvidia.com>

Addresses both review comments on #4588. Scope it to the non-capture path. On the token-capture path a replacement cannot borrow a finished group from a later step the way the non-capture path does, so every drop waits out a full fresh rollout with the trainer idle. Dynamic sampling makes drops routine, so that cost would land on every step. Setup now rejects use_dynamic_sampling together with token_capture.enabled, and FinalizedGroup.zero_variance goes away with the capture-side filtering. Mirroring the borrow on that path is its own change. Separate the budgets. A zero-variance drop is a data decision about a rollout that succeeded, not a rollout failure, so sharing the fault-tolerance knobs let each kind of drop redefine the other. generate_and_push now returns a distinct RolloutOutcome.FILTERED, counted against the new grpo.dynamic_sampling_max_replacements_per_slot rather than max_replacement_attempts. Filtering no longer reads on_dropped_prompt: a uniform group is never worth training on, so it is replaced under "shrink" too and the spare pool is filled either way. Exhausting the budget raises, as v1 does past dynamic_sampling_max_gen_batches, instead of crediting a shortfall against min_step_batch_fraction; an empty pool still shrinks, since that is the dataloader ending rather than a mis-tuned filter. The default budget of 32 is sized from the drop rate, not from how many resamples feel reasonable: a slot exhausts it with probability p**(N+1), and every slot of every step takes that draw. At 32 a 4096-group step over 1000 steps is quiet below roughly a 78% drop rate, well clear of the ~33% seen in practice, and raises above it. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: ruit <ruit@nvidia.com>

Follow-up on the review of the previous commit. Eight findings, all verified against the code before changing anything. The budget sizing guidance was wrong twice over. It solved 4096*p**33 = 1, which is a single step, but presented the answer as covering a 1000-step run, and both of those numbers were invented rather than taken from any recipe. The guidance now gives the expression it was a worked example of, p**(N+1) * num_prompts_per_step * max_num_steps, and says to read p off the zero_variance_fraction a run reports. Same fix in the field comment on dynamic_sampling_max_replacements_per_slot. Filtering no longer reads on_dropped_prompt, but the guard that keeps the spare pool fillable still sat under on_dropped_prompt="replace". A "shrink" run with dynamic sampling on could therefore set replacement_reserve_prompts=0, which _divert_batch_to_reserve satisfies immediately, leaving every filtered group to shrink its step. Setup now requires it to be at least 1 whenever dynamic sampling is on. _discard_committed_group could turn a benign race into a fatal error. remove_group raises for a slot that is already gone, and mutation cuts do not exclude each other, so the train pump can claim a ready group while the discard waits on the barrier. It now reports that case instead of raising, and leaves the rows alone rather than clearing them out from under a step that is already running; the caller treats the claimed group as committed rather than seeking a replacement that would over-fill the step. Dynamic sampling also costs training steps, because replacements are ordinary dataloader prompts and max_num_steps is clamped to one epoch of the loader. A run yields about (1 - p) of its configured steps and then ends cleanly. That is documented, and setup now warns when dynamic_sampling_max_gen_batches or batch_multiplier is moved off its default, since neither is read on this path. The uint16 widening in _as_routed_experts_tensor had no test at all. It has three now, in uint16 specifically: uint8 and every signed dtype implement max(), so a test written with one of those passes against the unwidened code. One of the three pins that premise so the others cannot quietly stop testing anything. Two smaller ones: a comment pointed at a line number that its own PR had already invalidated, and the rl.rollout.attempt span attribute was left on the fault-tolerance counter, so a slot whose drops were all filtered reported attempt 0 on every retry. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: ruit <ruit@nvidia.com>

The overflow test fed id 70 against a torch.int8 carry dtype and expected a ValueError, but the guard only raises when the id exceeds torch.iinfo(dtype).max, which is 127 -- so pytest.raises ended in DID NOT RAISE. Use 200 instead: it is above the int8 limit and wraps to -56, which is the silent corruption the guard exists to catch. The fenced budget formula added to the Dynamic sampling section sat at column 0 inside a list item, which closed the "Setup requires:" list. GitHub's renderer returned two <ul> blocks with the code block and the paragraph after it loose at top level. Indenting both by two spaces restores one four-item list with them nested in the first item. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: ruit <ruit@nvidia.com>

The method returns three RolloutOutcome values but Returns: named only two. The ownership rules it states -- the caller owns the backpressure permit and the checkpoint-atomic ledger transition -- apply to FILTERED as well, yet were written as if only SKIPPED reached them. The enum comment does not close the gap either: it states FILTERED's budget obligation but not the permit or the ledger record. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: ruit <ruit@nvidia.com>

replacement_reserve_prompts is a low-water mark, not a size: the pool is refilled only once it drops below the mark, and one refill diverts exactly one batch of num_prompts_per_step. The default of 1 was sized for rollout failures, which are rare and capped at max_replacement_attempts=1 per slot. Dynamic sampling puts a second consumer on the same pool that is routine and budgeted at 32, so below one batch the pool reaches 0 between refills, and the pump cannot top it up while parked in _launch awaiting buffer capacity or the weight-sync pause. A group filtered in that window is credited short rather than resampled, and enough of them fail the step against min_step_batch_fraction. Warn rather than raise: how much headroom a run needs is num_prompts_per_step * p / (1 - p) at drop rate p, and p is only known once the run reports zero_variance_fraction. Name the pool as a third cause in the short-step RuntimeError too. That error lists only a failing fleet and over-large drop budgets, so an operator hitting this spends the investigation on the cluster; the setup warning fired hours earlier and is buried in the startup log. Correct two guide sentences that promised what the code does not hold at the shipped default: that a filtered group shrinks its step only once the dataloader is exhausted, and that a short epoch then ends cleanly. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: ruit <ruit@nvidia.com>

RayenTian · 2026-10-29T10:05:22Z

/ok to test ede988d

_target_groups_for_step now reads self._dynamic_sampling_enabled to name the spare pool as a third cause. TestTargetGroupsForStep builds the actor with object.__new__ and never set that attribute, and the class has no __getattr__, so the four tests that assert the raise got AttributeError instead of RuntimeError. Set it in the fixture and add a case pinning the suffix in both directions, since nothing covered it. The warning added for a small spare pool advised sizing it at num_prompts_per_step * p / (1 - p), which is below num_prompts_per_step for every p < 0.5 -- the range the warning fires in. Following the advice did not clear the warning. The threshold is right: the refill check runs once per dataloader batch, so the mark must cover one batch of filtered draws, and num_prompts_per_step is that quantity at p = 0.5. Keep the threshold, drop the formula from the message, and say to set it to at least num_prompts_per_step. The formula stays as a code comment explaining where the threshold comes from. "the spare pool reaches 0 between refills" over-claimed: the pool is topped up once it falls below the mark, so it only reaches 0 when the mark is 1. It can still empty between two checks at any mark, which is the real hazard. Say "can empty" and name the once-per-batch check. The test's pytest.warns match string moves with it. The FILTERED enum comment said the caller "must not let it count as a step shortfall", but the filtered path credits a shortfall whenever the spare pool is dry -- which is what the new warning exists to prevent. The replacement_reserve_prompts field comment still said the key is read only under on_dropped_prompt="replace"; filtering draws on the same pool under "shrink". Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: ruit <ruit@nvidia.com>

RayenTian · 2026-10-29T14:32:28Z

/ok to test af69573

Signed-off-by: ruit <ruit@nvidia.com> # Conflicts: # nemo_rl/experience/rollout_manager.py # tests/unit/experience/test_rollout_manager.py

…pool guidance The functional recovery hook took replacement_index to match the signature but dropped it when delegating, so every delegated call recorded as a first attempt and would have falsified first_attempt_reward for anyone asserting on it there. Inert today -- the scripts that use the hook do not enable dynamic sampling -- but wrong, and the delegate defaults the kwarg so it failed silently. Three wording fixes in the spare-pool guidance: - The guide stated the refill condition twice in consecutive clauses. Fold them into one sentence that also carries the check-frequency fact, which was the only new information the second clause added. - "covers any drop rate below 50%; below that the pool can empty" read as "below 50% the pool can empty" -- the nearest antecedent is the percentage, not the mark. Say "set it lower". - A test helper's default-argument comment still carried "anything smaller reaches 0 in between", the claim the previous commit retired everywhere else. Also note grpo.max_num_epochs: null in the shorter-epoch paragraph. Accepting null landed on main in #4975; the SingleController clamp already returned early for it, so the paragraph's advice now has a second option. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: ruit <ruit@nvidia.com>

RayenTian · 2026-10-30T08:52:30Z

/ok to test 8583e94

RayenTian force-pushed the ruit/dapo_dynamic_sampling_sc branch from d7434e7 to 46ce5b8 Compare October 13, 2026 03:18

RayenTian force-pushed the ruit/dapo_dynamic_sampling_sc branch from 46ce5b8 to 97782d0 Compare October 15, 2026 05:44

RayenTian added the CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) label Oct 15, 2026

github-actions Bot added the Documentation Improvements or additions to documentation label Oct 15, 2026

RayenTian force-pushed the ruit/dapo_dynamic_sampling_sc branch from 32cc804 to 77b759a Compare October 16, 2026 06:42

RayenTian force-pushed the ruit/dapo_dynamic_sampling_sc branch from 8e98c7c to e55ff23 Compare October 19, 2026 04:32

RayenTian marked this pull request as ready for review October 19, 2026 08:47

RayenTian requested a review from yuki-97 October 19, 2026 13:58

RayenTian changed the title ~~feat: enable single-controller dynamic sampling~~ feat(sc): enable single-controller dynamic sampling Oct 19, 2026

RayenTian requested review from pthombre and saumishr October 19, 2026 14:03

RayenTian mentioned this pull request Oct 20, 2026

feat(single-controller): support dynamic sampling for streaming sampler #4784

Open

3 of 4 tasks

yuki-97 reviewed Oct 28, 2026

View reviewed changes

Comment thread nemo_rl/algorithms/single_controller.py Resolved

Comment thread nemo_rl/algorithms/single_controller_utils/config.py Resolved

RayenTian force-pushed the ruit/dapo_dynamic_sampling_sc branch from 4a29761 to 3bfe02f Compare October 29, 2026 01:36

nemo-automation-bot Bot requested a review from a team October 29, 2026 09:09

RayenTian and others added 5 commits October 29, 2026 03:00

feat: enable single-controller dynamic sampling and handle unsigned e…

9728901

…xpert IDs Signed-off-by: ruit <ruit@nvidia.com>

fix: apply zero-variance filtering on the token-capture path

49077d0

Signed-off-by: ruit <ruit@nvidia.com>

feat(single-controller): log per-step reward before dynamic sampling

c1a3bfb

Signed-off-by: ruit <ruit@nvidia.com>

test(single-controller): cover dynamic sampling validation and filter…

2805324

…ing, update docs Signed-off-by: ruit <ruit@nvidia.com>

RayenTian and others added 10 commits October 29, 2026 03:00

test(single-controller): give test fakes the dynamic-sampling manager…

193dea5

… API Signed-off-by: ruit <ruit@nvidia.com>

test: pin the zero_variance field on FinalizedGroup in the RPC field …

5af7969

…inventory Co-Authored-By: Claude Code <noreply@anthropic.com> Signed-off-by: ruit <ruit@nvidia.com>

RayenTian force-pushed the ruit/dapo_dynamic_sampling_sc branch from a88bc72 to ede988d Compare October 29, 2026 10:02

RayenTian mentioned this pull request Oct 30, 2026

docs: point the AVQA reference at the dataset the recipe actually loads #4994

Closed

2 of 4 tasks

RayenTian requested a review from yuki-97 October 30, 2026 03:22

RayenTian and others added 3 commits October 30, 2026 01:48

Merge origin/main into ruit/dapo_dynamic_sampling_sc

95effa1

Signed-off-by: ruit <ruit@nvidia.com> # Conflicts: # nemo_rl/experience/rollout_manager.py # tests/unit/experience/test_rollout_manager.py

Merge branch 'main' into ruit/dapo_dynamic_sampling_sc

8583e94

Conversation

RayenTian commented Oct 8, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do ?

Issues

Result

Non token capture

token capture

Testing

Scope and limitations

Usage

Before your PR is "Ready for review"

Additional Information

Uh oh!

copy-pr-bot Bot commented Oct 8, 2026

Uh oh!

cmunley1 commented Oct 12, 2026

Uh oh!

smahdavi4 commented Oct 14, 2026

Uh oh!

RayenTian commented Oct 18, 2026

Uh oh!

RayenTian commented Oct 19, 2026

Uh oh!

RayenTian commented Oct 19, 2026

Uh oh!

RayenTian commented Oct 19, 2026

Uh oh!

RayenTian commented Oct 20, 2026

Uh oh!

yuki-97 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

RayenTian commented Oct 29, 2026

Uh oh!

RayenTian commented Oct 29, 2026

Uh oh!

RayenTian commented Oct 29, 2026

Uh oh!

RayenTian commented Oct 30, 2026

Uh oh!

This branch has not been deployed

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

RayenTian commented Oct 8, 2026 •

edited

Loading