feat(sc): train path — StalenessSampler + _train_pump rewrite by yuki-97 · Pull Request #3220 · NVIDIA-NeMo/RL

yuki-97 · 2026-07-15T16:27:03Z

Part 2/4 of the #2819 split. See #2819 (reference) or #3267 (4/4) for the full code and context. Stack on #3219.

#3219 (1/4) -> #3220 (2/4) -> #3266 (3/4) -> #3267 (4/4)

Summary

Train-side refactor of the SC (Single Controller) async-GRPO path — swaps the trainer's data source from dp_client.claim_meta to the new TQReplayBuffer via a filter-only StalenessSampler.

StalenessSampler (nemo_rl/algorithms/async_utils/staleness_sampler.py): filter-only prompt-group selector over a TQReplayBuffer. Supports max_staleness_versions, sample_freshest_first, strict_weight_fifo, and force_in_order (matches legacy async_grpo target_weight semantics). Replaces the old top-level staleness_sampler.py (deleted).
SingleControllerActor._train_pump: rewritten to sampler.evict → sampler.select → _advantage_stage → TQPolicy split API (begin_train_step / train_microbatches_from_meta / finish_train_step) → dp_client.clear_samples. Drops orphaned pre-refactor helpers (_claim_available_meta, _claim_required_fields, _evict_stale_claimed, _flush_incomplete_groups) and their leftover config fields (consumer_task_name, claim_required_fields, max_claim_groups) + init attrs (_claimed_meta, _step_consumed_sample_ids).
single_controller_utils/utils.py: extracted SC helpers (aggregate_step_metrics, reduce_advantage_pump_metrics, fields_for_put, squeeze_trailing_unit_dim, tensor_field) — pulls pure functions out of single_controller.py to keep it focused on pump orchestration.
Tests: new test_staleness_sampler; new test_train_pump (real-Megatron TQPolicy end-to-end for _train_pump); shared _dp_fakes module (TQ data-plane fakes reused across pump tests); test_rollout_pump re-homed onto _dp_fakes.

Note: The production GRPO entrypoint (run_grpo) becomes end-to-end runnable only after PR#3 lands the driver-side setup + entrypoint wiring. The new test_train_pump already exercises SC end-to-end via a real Megatron TQPolicy on a tiny llama with the "simple" TQ backend.

Known Missing Features (Tracked in #2625)

Decouple logprob refresh from advantage config in the Single Controller: prev_logprobs_required / reference_logprobs_required gating is provisional.
Multi-mini-step inside a single RL step is not supported (one optimizer.step per RL step).
Add custom sampler support — spilt TQ impl out and let people could easily rewrite their own custom sampler for select and evict logic.

copy-pr-bot · 2026-07-15T16:27:06Z

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 · 2026-07-15T16:28:34Z

/ok to test bfb9075

terrykong

Reviewed in two passes with independent 5-agent teams (RL expert, test author, bug scanner, comment reviewer, devil's advocate); every finding was independently re-verified before staging, and the whole set was re-checked against the rebased head 283ba7f and against the upstack PRs (#3266/#3267) — findings resolved by the rebase or upstack (vllm marker, restored fast-lane tests, TaskGroup dispatch restructure, set_weight_version wiring, package __init__.py) were dropped rather than left as noise. Solid refactor overall — the sampler unit tests are thorough, the extracted helpers are byte-faithful moves of the originals, and the deleted nemo_rl/algorithms/staleness_sampler.py leaves zero dangling importers. 👍

One stack-aware note worth calling out: #3266 comments out the _sync_weights drain loop, which removes the invariant that keeps the evict-vs-commit race unreachable — that strengthens the case for the evict-skips-unready fix suggested on staleness_sampler.py.

Note: some comments below were staged against pre-rebase commit 6580a91 and may show as "outdated" threads — all were re-verified to still apply at 283ba7f.

CI: the last /ok to test bfb9075 predates the rebases, so head 283ba7f has no gated CI runs. Could you re-trigger with /ok to test 283ba7f3f8c366484a2d2ef1d142f15c7793010a? The mcore e2e (test_train_pump) couldn't be verified locally, so CI is the verification path for it.

pre-commit run --all-files passes on the branch, and the new source files (async_utils/staleness_sampler.py, single_controller_utils/utils.py) are in pyrefly.toml project-includes.

Generated by Claude Code

Split the monolithic StalenessSampler into a PromptGroupSampler interface (admit/select/evict + is_on_policy/required_buffer_capacity) with one named policy per behavior (WindowedSampler / WeightFifoSampler / InOrderSampler), selected via a discriminated-union config or a module:ClassName FQN. admit() moves the rollout-pump dispatch gate + target_step stamping into the sampler so the pump follows the selected algorithm; InOrderSampler keys evict on target_step so evict/select agree by construction. Legacy knobs still work via sampler_from_legacy_knobs(); StalenessSampler kept as a compat facade. DRAFT / design demonstration for the #3220 sampler discussion. Not merged; not run locally (no compatible env). Signed-off-by: Terry Kong <terryk@nvidia.com>

Replace the monolithic StalenessSampler (five mode-encoding flags, three disjoint select branches, a shared evict fitting only two of the modes) with a PromptGroupSampler interface — admit/select/evict + is_on_policy / required_buffer_capacity — and one named policy per behavior: WindowedSampler / WeightFifoSampler / InOrderSampler. The sampler is selected by a pydantic discriminated-union config (name = windowed|weight_fifo|in_order|custom) or a module:ClassName FQN, so invalid knob combinations are unrepresentable and there are no cross-field validations in __init__. admit() moves the rollout-pump dispatch gate + target_step stamping into the sampler (the pump follows whichever policy is selected); InOrderSampler keys evict on target_step so evict/select agree by construction. No legacy-knob path: over_sampling / force_in_order / batch_selection_strategy / max_weight_staleness_versions and warn_if_staleness_window_below_minibatches are removed; the actor no longer owns _max_rollout_version. test_staleness_sampler (monolith) is replaced by test_sampler_interface. DRAFT / design demonstration for the #3220 sampler discussion. Not merged; not run locally (no compatible env). Signed-off-by: Terry Kong <terryk@nvidia.com>

) (#3316) Signed-off-by: Terry Kong <terryk@nvidia.com>

RayenTian · 2026-07-23T08:55:52Z

/ok to test 4d0fa4e

Signed-off-by: Yuki Huang <yukih@nvidia.com>

…wire StalenessSampler init Signed-off-by: Yuki Huang <yukih@nvidia.com>

Signed-off-by: Yuki Huang <yukih@nvidia.com>

…ighbours Signed-off-by: Yuki Huang <yukih@nvidia.com>

Signed-off-by: Yuki Huang <yukih@nvidia.com>

…p attrs Signed-off-by: Yuki Huang <yukih@nvidia.com>

…s init arg Signed-off-by: Yuki Huang <yukih@nvidia.com>

…window Signed-off-by: Yuki Huang <yukih@nvidia.com>

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

) (#3316) Signed-off-by: Terry Kong <terryk@nvidia.com>

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

RayenTian · 2026-07-23T09:21:50Z

/ok to test e6f3d76

saumishr

Seems the PR is merged. Most of my comments can be addressed as a follow up

saumishr · 2026-07-23T17:53:48Z

+
+
+@runtime_checkable
+class PromptGroupSampler(Protocol):


nit Lets keep the interfaces and concrete implementations in their own folders for better maintainability.

saumishr · 2026-07-23T18:01:14Z

 # you may not use this file except in compliance with the License.
 # You may obtain a copy of the License at
 #
 #     http://www.apache.org/licenses/LICENSE-2.0


We should probably name this as rl_controller. Single Controller term we have been using to mostly communicate the design choice of having a single controller doing all of the orchestration. Please feel free to do that in a follow up PR so that this one doesn't get blocked.

saumishr · 2026-07-23T18:43:18Z

+    """Shared machinery for the built-in policies.
+
+    Owns the monotonic dispatch counter (the batch index formerly tracked as
+    ``SingleControllerActor._max_rollout_version``) and the common


This doc needs to be updated? We seem to have removed _max_rollout_version?

saumishr · 2026-07-23T19:15:14Z

@@ -0,0 +1,455 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.


init.py missing?

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

yuki-97 force-pushed the yukih/sc-split-2 branch from bfb9075 to f7a69a3 Compare July 17, 2026 16:10

yuki-97 changed the base branch from main to yukih/sc-split-1 July 17, 2026 16:19

yuki-97 commented Jul 18, 2026

View reviewed changes

Comment thread nemo_rl/algorithms/single_controller.py

yuki-97 mentioned this pull request Jul 18, 2026

[Single Controller / Async RL] cleanup tracking issue #2625

Open

22 tasks

yuki-97 force-pushed the yukih/sc-split-2 branch from a1d41bf to 6580a91 Compare July 18, 2026 13:19

This was referenced Jul 19, 2026

feat(sc): setup + entrypoint #3266

Open

feat(sc): NeMo-Gym path #3267

Open

feat(sc): rollout path — TQReplayBuffer, rollout_manager, _rollout_pump #3219

Merged

yuki-97 requested review from ZhiyuLi-Nvidia, saumishr and terrykong July 19, 2026 09:18

yuki-97 assigned RayenTian Jul 19, 2026

RayenTian force-pushed the yukih/sc-split-1 branch 4 times, most recently from f929aa0 to 48184ad Compare July 22, 2026 05:08

RayenTian force-pushed the yukih/sc-split-2 branch from 6580a91 to 283ba7f Compare July 22, 2026 05:41

RayenTian marked this pull request as ready for review July 22, 2026 05:42

RayenTian requested review from a team as code owners July 22, 2026 05:42

terrykong reviewed Jul 22, 2026

View reviewed changes

RayenTian force-pushed the yukih/sc-split-2 branch 2 times, most recently from ac394aa to d3188b2 Compare July 22, 2026 14:14

Base automatically changed from yukih/sc-split-1 to main July 22, 2026 23:47

terrykong requested a review from a team as a code owner July 22, 2026 23:47

terrykong mentioned this pull request Jul 23, 2026

[DRAFT] refactor(sc): interface-driven staleness samplers (design demo for #3220) #3316

Merged

RayenTian pushed a commit that referenced this pull request Jul 23, 2026

refactor(sc): interface-driven staleness samplers (design demo for #3220

f1589da

) (#3316) Signed-off-by: Terry Kong <terryk@nvidia.com>

RayenTian pushed a commit that referenced this pull request Jul 23, 2026

refactor(sc): interface-driven staleness samplers (design demo for #3220

7a6676e

) (#3316) Signed-off-by: Terry Kong <terryk@nvidia.com>

RayenTian force-pushed the yukih/sc-split-2 branch from f1589da to 4d0fa4e Compare July 23, 2026 08:50

yuki-97 and others added 17 commits July 23, 2026 02:14

feat(sc): train path — StalenessSampler + _train_pump rewrite

30aad15

Signed-off-by: Yuki Huang <yukih@nvidia.com>

refactor(sc): extract SC helpers to single_controller_utils.utils, re…

056c12e

…wire StalenessSampler init Signed-off-by: Yuki Huang <yukih@nvidia.com>

docs(sc): fix stale references in module/pump/warn docstrings

8ea11d2

Signed-off-by: Yuki Huang <yukih@nvidia.com>

test(sc): cover StalenessSampler force_in_order path

67ad762

Signed-off-by: Yuki Huang <yukih@nvidia.com>

refactor(sc): rename StalenessSampler.require_order → strict_weight_fifo

e990a7c

Signed-off-by: Yuki Huang <yukih@nvidia.com>

test(sc): drop redundant StalenessSampler.select cases subsumed by ne…

401196c

…ighbours Signed-off-by: Yuki Huang <yukih@nvidia.com>

pyrefly

7126ef7

Signed-off-by: Yuki Huang <yukih@nvidia.com>

refactor(sc): drop unused claim-state config fields and init attrs

25a8c5f

Signed-off-by: Yuki Huang <yukih@nvidia.com>

test(sc): add real-Megatron e2e train_pump test, fix stale _train_pum…

4effa0f

…p attrs Signed-off-by: Yuki Huang <yukih@nvidia.com>

refactor(sc): drop unused max_rollout_prompts config field and prompt…

71fe413

…s init arg Signed-off-by: Yuki Huang <yukih@nvidia.com>

test(sc): multi-step train_pump, cover force_in_order weight-outside-…

4fe817d

…window Signed-off-by: Yuki Huang <yukih@nvidia.com>

test(sc): fix stale replay buffer method reference

f721dc0

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

fix(sc): address train path review feedback

3af580e

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

fix(sc): restore advantage metrics and coverage

52e60e8

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

refactor(sc): interface-driven staleness samplers (design demo for #3220

00ac501

) (#3316) Signed-off-by: Terry Kong <terryk@nvidia.com>

fix(sc): enforce complete training steps

8257e1d

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

style(sc): format sampler files

e6f3d76

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

RayenTian force-pushed the yukih/sc-split-2 branch from 4d0fa4e to e6f3d76 Compare July 23, 2026 09:20

terrykong approved these changes Jul 23, 2026

View reviewed changes

terrykong merged commit fa81188 into main Jul 23, 2026
80 checks passed

terrykong deleted the yukih/sc-split-2 branch July 23, 2026 18:41

saumishr reviewed Jul 23, 2026

View reviewed changes

		@@ -0,0 +1,455 @@
		# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved.



		@runtime_checkable
		class PromptGroupSampler(Protocol):

Uh oh!

Conversation

yuki-97 commented Jul 15, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Known Missing Features (Tracked in #2625)

Uh oh!

copy-pr-bot Bot commented Jul 15, 2026

Uh oh!

yuki-97 commented Jul 15, 2026

Uh oh!

Uh oh!

terrykong left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

RayenTian commented Jul 23, 2026

Uh oh!

RayenTian commented Jul 23, 2026

Uh oh!

Uh oh!

saumishr left a comment

Choose a reason for hiding this comment

Uh oh!

saumishr Jul 23, 2026

Choose a reason for hiding this comment

Uh oh!

saumishr Jul 23, 2026

Choose a reason for hiding this comment

Uh oh!

saumishr Jul 23, 2026

Choose a reason for hiding this comment

Uh oh!

saumishr Jul 23, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

yuki-97 commented Jul 15, 2026 •

edited

Loading