feat(sc): gate-authoritative token capture via an external TransferQueue sink by pthombre · Pull Request #3837 · NVIDIA-NeMo/RL

pthombre · 2026-08-26T05:12:36Z

What this PR adds

Gate-authoritative token capture for NeMo-Gym rollouts on the SingleController path: instead of round-tripping token ids through the environment (token echo) or re-tokenizing text, the vLLM worker stages each model call's exact token delta into TransferQueue (the external sink), NeMo-Gym's policy model server custodies the lineage in a capture ledger, and a finalizer actor pool assembles verified receipts into training rows.

Enabled by token_capture.enabled=true (NeMo-Gym + vLLM async engine only, validated loudly at setup). The legacy token-echo path is unchanged and remains the default.

How it works (rollout → training row)

Staging (external sink). The vLLM async worker hosts Gym's capture core in-process. For every model call it stages the call's token delta (prompt/generation coordinates, logprobs, and optionally routed experts) into a TQ staging partition, keyed by a per-rollout id (/ng-rollout/<id>) that rides the Gym run body — the TQ sample id is the capture key end to end (nemo_rl/data_plane/tq_token_sink.py, vllm_worker_async.py).
Custody (capture ledger). Gym's policy model server records admission → staged-coordinates → terminal attribution as a token-free custody ledger (chained digest). Uncommitted off-chain calls poison only their own rollout, and undeclared rollouts get witness-based / heuristic terminal selection rather than silent acceptance.
Receipt assembly. After a rollout finishes, NeMo-RL fetches the manifest over the ledger control plane (bearer-token authenticated) and assembles a receipt: the verified prefix chain naming exactly which staged rows constitute the trajectory (nemo_rl/environments/nemo_gym.py, nemo_rl/experience/blackbox_finalizer.py).
Finalization (actor pool). The rollout pump hands a metadata-only FinalizationRequest (group id, rollout ids, receipts, rewards — never token tensors) to a finalizer actor pool, which publishes canonical training rows into the TQ training partition and commits the replay-buffer slot (nemo_rl/experience/finalizer_actor.py, route_plan.py, single_controller.py actor dispatch path).
Training. The train pump consumes committed groups as usual. Capture placeholders carry sample_mask=0, and the advantage estimator now takes a real valid_mask instead of hardwired all-ones (advantage_estimator.py).

Supporting changes

TQReplayBuffer token-capture mode: reserve(rollout_ids=…), commit_finalized, abort, staging-aware remove (nemo_rl/algorithms/async_utils/replay_buffer.py).
Routed-experts transport: nrlre1 envelope decode in the sink, int16 fields broadcast as int32 (NCCL has no Short), all-placeholder groups dropped when routed dims are unknown (data_plane/schema.py, preshard.py, worker_mixin.py).
token_capture config block with setup-time validation of the supported matrix; derived fields (control-plane auth token, capture dir) filled at setup (single_controller_utils/config.py, setup.py).
Build: Gym becomes an editable path dependency instead of a uv workspace member — both projects define a vllm extra with intentionally different server stacks, and uv applies a requested extra to every workspace member, making the RL-vLLM + base-Gym worker environment unselectable otherwise (pyproject.toml, regenerated uv.lock).
Gym submodule pinned to fd5e84d6 (feat(token-id-capture): external-sink token capture (worker-owned staging, ledger, terminal attribution) Gym#2872, the tokidcap capture stack, merged into Gym main 2026-09-04).

Tests

New unit suites: token sink, blackbox finalizer, finalizer actor + lifecycle, route plan, worker route assembly, TQ policy routes, nemo-gym token capture, vLLM capture hosting, advantage validity, plus TQ replay-buffer capture-mode coverage.
L1 functional gains a token-capture arm: the SC+Gym smoke rerun with ++token_capture.enabled=true.
Legacy-path invariance: all existing SC/experience/data-plane suites pass unchanged (RAY_ADDRESS=local run on a CPU node; final revalidation in flight).

Not in this PR

The nano SWE launch recipe (configs, launchers) and the recipe guides are kept out to scope this PR to core logic. History note: this supersedes #3456 (same feature, rebased onto current main with a clean history; the TQ actor runtime-env pin and the mtp refit fix from that branch were dropped as both are superseded upstream).

🤖 Generated with Claude Code

copy-pr-bot · 2026-08-26T05:12:40Z

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.

github-actions · 2026-09-01T03:08:43Z

✅ Submodule Fast-Forward Check Results

Check based on commit: b3e70a0 (PR #3837 from pthombre/tq-tokidcap-swe-v2)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

github-actions · 2026-09-01T03:25:28Z

✅ Submodule Fast-Forward Check Results

Check based on commit: f130e85 (PR #3837 from pthombre/tq-tokidcap-swe-v2)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

pthombre · 2026-09-01T05:30:27Z

/ok to test 14dde62

github-actions · 2026-09-01T05:31:29Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 14dde62 (PR #3837 from pthombre/tq-tokidcap-swe-v2)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

github-actions · 2026-09-01T19:14:23Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 65fa9b0 (PR #3837 from pthombre/tq-tokidcap-swe-v2)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

pthombre · 2026-09-02T01:53:33Z

/ok to test 82689d1

ZhiyuLi-Nvidia

Review finding: advantage-estimator valid_mask is only honored by GRPO.

ZhiyuLi-Nvidia · 2026-09-02T10:19:12Z

`_cleanup_consumed_metas` blocks the controller loop — rely on TQ's native async clear API

The end-of-step batching is a good change and should stay; it is fewer round trips than the per-chunk clears it replaced. The issue is only which TQ entry point it goes through.

_cleanup_consumed_metas                 single_controller.py:1314  (train pump, :2299)
└ self._call_dp("clear_samples")        :1057-1065 → local client → INLINE, no yield
  └ TQAdapter.clear_samples             adapters/transfer_queue.py:990-1015  (sync def)
    └ tq.kv_clear                       interface.py:558-579     ← sync veneer
      └ _run_coroutine(async_kv_clear(keys, partition_id))
        │                               client.py:1238-1241 → run_coroutine_threadsafe
        │                               future.result()   ← no timeout, not cancellable
        └ async_kv_clear                interface.py:933-959  ← the coroutine, already
          ├ async_kv_retrieve_meta      client.py:898
          ├ async_clear_samples         client.py:542
          └ storage_manager.clear_data
              simple → gather over SUs  simple_storage_manager.py:486   truly async
              KV     → NO await at all  managers/base.py:803-813        blocks the loop

_call_dp resolves to a local DataPlaneClient, so method(**kwargs) runs inline (:1063) with no suspension point, and the thread then blocks on future.result(). Called once per step from the train pump, that parks the rollout pump, the fleet-health probe and the finalizer pool — including _stall_watchdog_pump, which is a task on the same loop (:498) and so cannot report the stall it exists to catch.

Recommendation: call tq.async_kv_clear (interface.py:933) instead of tq.kv_clear. As the stack shows, kv_clear is _run_coroutine(async_kv_clear(...)) — the native API is the same coroutine without the thread hop and the un-cancellable wait. Same round trips; the loop keeps turning.

# nemo_rl/data_plane/adapters/transfer_queue.py
    async def async_clear_samples(self, sample_ids: list[str] | None, partition_id: str) -> None:
        """Async twin of clear_samples: awaits TQ's native coroutine directly."""
        self._mark_data_operation_started()
        await tq.async_kv_clear(keys=list(sample_ids), partition_id=partition_id)

# single_controller.py, _cleanup_consumed_metas
        await self._call_dp("async_clear_samples", sample_ids=unique_ids, partition_id=partition_id)

No plumbing needed — _call_dp (:1064) and call_data_plane (data_plane/async_utils.py:54) already await a returned coroutine.

Two things to settle first:

Safe on simple only today. KV-backend clear_data is an async def with no await in it (managers/base.py:803-813), so on mooncake_cpu the blocking batch_remove would land on SingleController's loop instead of TQ's. Needs the upstream await asyncio.to_thread(...) fix before that backend can use this path.
Cancellation becomes real. async_clear_samples marks-clearing before it clears (client.py:566-572), so a cancelled clear leaves rows unreadable but unfreed. Shield it, or accept that deliberately.

Interim one-liner if you'd rather not add the adapter's first async method in this PR: call_data_plane(self._dp_client, "clear_samples", offload_sync=True, ...) at both call sites — keeps the block off the loop on every backend, no upstream change.

Separately: _clear_data_plane_samples (:1068) now has no production callers — only tests/unit/single_controller/test_checkpointing.py:1192,1214 exercise it. Delete and repoint those tests, or reuse it here with a partition_id parameter.

pthombre · 2026-09-03T07:58:46Z

/ok to test 4e50901

Rollout token deltas are staged into TransferQueue by the vLLM worker (external sink), custody is tracked by Gym's capture ledger, and a finalizer actor pool assembles receipts into training rows — no token echo, no re-tokenization. Rebased onto main: the TQ runtime-env pin and the mtp refit fix were dropped (both superseded upstream), and uv.lock was regenerated against main's resolution. Signed-off-by: Prashant Thombre <pthombre@nvidia.com>

github-actions · 2026-09-05T01:19:44Z

✅ Submodule Fast-Forward Check Results

Check based on commit: f6305dc (PR #3837 from pthombre/tq-tokidcap-swe-v2)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

… lint after main merge The origin/main merge (f6305dc) auto-merged nemo_gym.py without conflict but dropped the token_capture kwarg from the new spinup_nemo_gym_actor wrapper introduced by #3367, while setup.py still passed it. Restore the parameter and forward it to build_nemo_gym_config; assert in the unit test that it reaches the actor config. Also fix the Lint CI job: point pyrefly.toml at the renamed rollout_reassembler*.py modules (the stale blackbox_finalizer / finalizer_actor entries made pyrefly abort), and apply ruff format/isort to the files the merge left unformatted. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Pranav Thombre <pthombre@nvidia.com>

github-actions · 2026-09-05T02:01:22Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 37cbd31 (PR #3837 from pthombre/tq-tokidcap-swe-v2)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

pthombre · 2026-09-05T02:01:25Z

/ok to test 37cbd31

…test setup_single_controller now indexes master_config.logger["wandb_enabled"] and ["wandb"] (main's full-result-table gating), so replacing the fixture's logger dict wholesale raised KeyError before the assertions ran. Extend the dict instead. Found while auditing the origin/main merge f6305dc. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Pranav Thombre <pthombre@nvidia.com>

github-actions · 2026-09-05T02:08:31Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 0e28303 (PR #3837 from pthombre/tq-tokidcap-swe-v2)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

The pre-commit `ruff check --select I` hook is stricter than the default rule set, so `ruff check` alone missed this after the RolloutReassembler rename. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Pranav Thombre <pthombre@nvidia.com>

github-actions · 2026-09-05T02:14:42Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 30d300c (PR #3837 from pthombre/tq-tokidcap-swe-v2)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

pthombre · 2026-09-05T02:37:07Z

/ok to test 30d300c

…t tests Two unit tests drifted from code this branch added and went red in CI (run 33939587768 on 30d300c): - test_train_pump_e2e::_prepopulate_buffer inserts a ready slot into TQReplayBuffer by appending to its parallel lists by hand. The branch added _rollout_ids_list and _staging_keys_list to the buffer, so the first remove() indexed past the end of the empty staging list (IndexError in _remove_unlocked). Append None to both, matching what begin_group/commit and checkpoint restore do for finalized groups. - test_rollout_pump_writes_expected_tq_data pins the TQ tag schema. The branch added num_routed_experts_backfilled to VIOLATION_TAG_KEYS so every row's tags now carry it; add it to the expected set. The remaining red shard, L0_Unit_Tests_Megatron_Policy_1, is a main-side conflict between #3881 (test expects train_microbatch to forward model_slices_context_parallel_inputs) and #2957 (guard raises on it) and is inherited unchanged. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Pranav Thombre <pthombre@nvidia.com>

github-actions · 2026-09-05T06:11:02Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 9286c3e (PR #3837 from pthombre/tq-tokidcap-swe-v2)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

github-actions · 2026-09-05T06:13:19Z

✅ Submodule Fast-Forward Check Results

Check based on commit: f54bbb1 (PR #3837 from pthombre/tq-tokidcap-swe-v2)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

pthombre · 2026-09-05T06:27:20Z

/ok to test f54bbb1

Catches the telemetry branch up to main, which had moved 22 commits and tripped pr-branch-up-to-date-check (max 10 behind). Eleven files conflicted. Resolutions, in the order a reviewer would want them: - single_controller.py: main split _dispatch_one_prompt into a token-capture branch (generate_for_finalization + the finalizer actor pool, #3837) and moved the legacy generate_and_push loop into an else. Git interleaved our per-prompt span into the new branch, so the function was restored to main's exactly and the per_prompt_scope() + rl.sc.generate_and_push umbrella re-applied to the legacy loop only. The token-capture branch is left uninstrumented: it commits through the finalizer pool rather than here, so its attempt boundary is a different shape than this span describes. Recorded in the coverage-gap table in docs/observability/span-groups.md. - run_grpo_single_controller.py: keeps our startup_span / setup_span phases and adopts main's VLM processor (#4009), with the processor-aware tokenizer construction moved inside setup_span ("tokenizer") and processor= threaded to setup_single_controller. - run_grpo.py: adopts main's make_policy_factory() helper in place of our inline factory selection, keeping setup_span("workers"). - nemo_gym.py: keeps @accepts_trace_context and the run_rollouts / _stream_rollouts span split, alongside main's ledger control plane; both signatures widen to main's 4-tuple yield. - rollout_manager.py: keeps dispatch_with_trace_context and unpacks main's added resolved_agent_ref. - factory.py: adopts main's LocalDataPlaneConfig-aware observability lookup, keeping the telemetry_enabled_in_env() arm that installs the wrapper for its spans. - vllm_worker.py: keeps umbrella_trace_fn(U_MODEL_INIT) on _load_model beside main's _refit_with_reload_api_enabled. - config.py, worker_mixin.py, virtual_cluster.py, vllm_generation.py: both sides added adjacent imports or fields; all kept. The textually merged uv.lock was corrupt (a missing source field on opentelemetry-instrumentation-aiohttp-client, which matched more than one package), so it was regenerated from main's with uv 0.11.28 to match CI's revision 3. The result is 52 insertions: the nemo-lens rev bump and the three otel aiohttp packages, nothing else. Both uv lock --check runs pass. Verified: all 54 changed Python files compile, ruff check and format are clean, and the seven test_source_drift.py guards pass -- including the carrier guard, which confirms run_rollouts is still dispatched with trace context after the signature change. Signed-off-by: Raj Singh <rajsin@nvidia.com>

github-actions Bot added the Documentation Improvements or additions to documentation label Aug 26, 2026

pthombre force-pushed the pthombre/tq-tokidcap-swe-v2 branch from fd8314c to fc6a551 Compare August 26, 2026 05:19

pthombre changed the title ~~Pthombre/tq tokidcap swe v2~~ feat(sc): gate-authoritative token capture via an external TransferQueue sink Aug 26, 2026

macandro96 mentioned this pull request Aug 31, 2026

feat(sc): sibling level rollout checkpointing at train boundary #3923

Merged

4 tasks

pthombre marked this pull request as ready for review August 31, 2026 23:56

pthombre requested review from a team as code owners August 31, 2026 23:56

pthombre force-pushed the pthombre/tq-tokidcap-swe-v2 branch 2 times, most recently from 9496ff6 to b3e70a0 Compare September 1, 2026 03:07

terrykong mentioned this pull request Sep 1, 2026

feat(sc): gate-authoritative token capture on the Gym tokidcap stack #3456

Closed

ZhiyuLi-Nvidia reviewed Sep 2, 2026

View reviewed changes

Comment thread nemo_rl/algorithms/single_controller.py

ZhiyuLi-Nvidia reviewed Sep 2, 2026

View reviewed changes

Comment thread nemo_rl/algorithms/async_utils/replay_buffer.py

ZhiyuLi-Nvidia reviewed Sep 2, 2026

View reviewed changes

Comment thread nemo_rl/algorithms/advantage_estimator.py

ZhiyuLi-Nvidia reviewed Sep 2, 2026

View reviewed changes

Comment thread nemo_rl/algorithms/async_utils/replay_buffer.py Outdated

terrykong reviewed Sep 5, 2026

View reviewed changes

Comment thread nemo_rl/data_plane/worker_mixin.py

Comment thread nemo_rl/algorithms/single_controller.py

terrykong enabled auto-merge (squash) September 5, 2026 04:56

Merge branch 'main' into pthombre/tq-tokidcap-swe-v2

f54bbb1

terrykong approved these changes Sep 5, 2026

View reviewed changes

terrykong disabled auto-merge September 5, 2026 19:25

terrykong merged commit 7036e5d into main Sep 5, 2026
68 of 70 checks passed

terrykong deleted the pthombre/tq-tokidcap-swe-v2 branch September 5, 2026 19:25

yfw mentioned this pull request Sep 10, 2026

build: bump vllm to 0.28.0 #4058

Closed

This was referenced Sep 10, 2026

perf(reassembler): batch group staging reads and expose fetch timings #4072

Draft

refactor(sc): unify capture retries and extract orchestration components #4073

Draft

macandro96 mentioned this pull request Sep 15, 2026

feat(rollout): integrate MInf ledger capture with checkpointing #4129

Open

Conversation

pthombre commented Aug 26, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What this PR adds

How it works (rollout → training row)

Supporting changes

Tests

Not in this PR

Uh oh!

copy-pr-bot Bot commented Aug 26, 2026

Uh oh!

github-actions Bot commented Sep 1, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Sep 1, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

pthombre commented Sep 1, 2026

Uh oh!

github-actions Bot commented Sep 1, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Sep 1, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

pthombre commented Sep 2, 2026

Uh oh!

Uh oh!

Uh oh!

ZhiyuLi-Nvidia left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

ZhiyuLi-Nvidia commented Sep 2, 2026

_cleanup_consumed_metas blocks the controller loop — rely on TQ's native async clear API

Uh oh!

pthombre commented Sep 3, 2026

Uh oh!

github-actions Bot commented Sep 5, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Sep 5, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

pthombre commented Sep 5, 2026

Uh oh!

github-actions Bot commented Sep 5, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Sep 5, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

pthombre commented Sep 5, 2026

Uh oh!

Uh oh!

Uh oh!

github-actions Bot commented Sep 5, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Sep 5, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

pthombre commented Sep 5, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

pthombre commented Aug 26, 2026 •

edited

Loading

`_cleanup_consumed_metas` blocks the controller loop — rely on TQ's native async clear API