feat(sc): gate-authoritative token capture on the Gym tokidcap stack - #3456
Draft
pthombre wants to merge 3 commits into
Draft
feat(sc): gate-authoritative token capture on the Gym tokidcap stack#3456pthombre wants to merge 3 commits into
pthombre wants to merge 3 commits into
Conversation
Reproducible 6-node GB200 smoke that runs agentic SWE RL on Nemotron-3-Nano-30B-A3B with rollouts flowing through TransferQueue. Verified (Slurm 5648757): run_grpo_single_controller.py + nano_swe_teacher_sc.yaml completed train steps 1-5 with live TransferQueueController + 2 SimpleStorageUnit actors serving PUT_DATA / KV_RETRIEVE_META, no traceback and no OOM. Rewards are 0.0 throughout -- the base model solves no SWE instance in a 5-step smoke, so this validates the loop, not model quality. Only two entrypoints honour the data plane: grpo_train_sync and run_grpo_single_controller.py. run_grpo_nemo_gym.py's async path builds the in-memory ReplayBuffer (grpo.py:3898) and silently ignores data_plane, so async + real TQ means SingleController. Contents: - examples/configs/ultra/nano_swe_teacher_sc.yaml -- the recipe, plus its inheritance chain (swe_teacher, _nano_smoke_gb200.inc, nano_swe_teacher, nano_swe_teacher_qwen3mesh) and the sync+TQ variant - launchers: swe_nano.env plus interactive and batch wrappers over ultra_launch.sh, which gains an NRL_ENTRYPOINT override so a run can use run_grpo_single_controller.py - docs/guides/nano-swe-transferqueue.md -- setup, the read-only paths to reuse vs the per-user write paths to change, why each SingleController setting exists, and how to confirm TQ is actually engaged - nemo_rl/utils/config.py: skip mandatory (???) values during config merge. Required here: swe_teacher.yaml declares sif_dir: ??? and resolving it mid-merge raises MissingMandatoryValue before the child config supplies the value. - pin 3rdparty/Gym-workspace/Gym to v0.4.0 to match the container's prebuilt gym venvs (skip_venv_if_present) Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
A 6-node submission was held with Reason=QOSMaxNodePerUserLimit: the short QOS enforces a per-user node cap that a larger concurrent job of the same user had already consumed. short also caps walltime at 2h, which is only just enough for a cold start (~60 GB checkpoint download plus Megatron conversion) followed by five steps. Submit with no QOS and WALLTIME=3:59:00, the batch partition maximum. ultra_launch.sh only auto-selects short below 2h, so an empty SLURM_QOS means no --qos flag. Opting back in is one env var. Signed-off-by: Zhiyu Li <zhiyul@NVIDIA.com>
Token-in/token-out capture for NeMo-Gym rollouts in the async SingleController GRPO path, rebuilt on the nano SWE recipe branch (3fcc696, main + the TQ-honoured SWE recipe) per docs/design-docs/tq-gym-tokidcap-migration.md §9b.1a. The Gym submodule is pinned to pthombre/tq-tokidcap-capture (b6051536 = upstream tokidcap stack top 81ac2736/#2182 + the 7-commit gate/staging series; companion Gym PR NVIDIA-NeMo/Gym#2278). This squashes the capture series (S1-S5 primitives/worker/gate/receipts/ verification + the stack migration) into one commit ported onto the main-lineage SC: the granular history lives on the prior branch head 624bb27. Highlights: - TQTokenSink/TQTokenSource: per-call token deltas staged from the vLLM worker straight into TransferQueue, durable before the response releases; BlackboxFinalizer rebuilds digest-verified always-N training rows via Gym's terminal-aware linearize over run_builder. - Identity: rollout ids ride the run body as the opaque _ng_rollout_id key (agents stamp /ng-rollout/<id>; all agent impls unmodified); the middleware-minted model_call_id is the call id; TQ sample ids ARE the capture keys. - Gate hosting config: LineageIndex capacity derived from the training config; per-run control-plane bearer token; hard per-call control deadlines (gate death surfaces as failed dispatches + placeholders, not a silent retry stall); base capture dir under the run's log dir. - Receipt-mode dispatch in RolloutManager (_generate_and_finalize: reserve-with-ids, seal -> receipt, finalize_group, commit_finalized, always-N with placeholders), weight-version fan-out to workers, gate metrics logging (token_in_rate, fallback-by-cause). - Port adaptations onto main's structures: streaming run_rollouts (receipt postprocess branch + picklable aiohttp error re-raise), commit()'s pre-write and evicted-during-write guards keep main's semantics, legacy failure path keeps main's remove_group. Dormant by default: everything is behind token_capture.enabled=false. Working log: docs/design-docs/tq-gym-tokidcap-migration-log.md. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>
github-actions
Bot
added
the
Documentation
9 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do?
Adds gate-authoritative token-id capture for the async single-controller (SC) + NeMo-Gym rollout path, rebuilt on the upstream Gym token-id-capture PR stack:
vllm_modelgate custodies token lineage: a capacity-sizedLineageIndex, a bearer-authenticated control plane, and bounded client deadlines.TokenSink.run_builderlinearize.Everything is gated behind
token_capture.enabled(defaultfalse); flag-off byte-identity was verified.Note: this branch is based on the nano-SWE recipe branch, so the PR also carries the two nano-SWE recipe commits (87866ee, 3fcc696) that are not yet on
main.Stack dependency
Requires the Gym side: NVIDIA-NeMo/Gym#2278 (branch
tq-tokidcap-capture), pinned as the submodule gitlink in this PR.Supersedes
Supersedes closed #3455 — same work pre-rebase; this is the rebase onto the nano-SWE base, squashed, with docs excluded.
Validation (2026-07-31, 8xH100)
Known CI issue
tests/functional/L1_Functional_Tests_SingleController.shruns the flag-off leg before the capture leg. Worker venvs are cached by actor class name only, so the flag-off leg builds theVllmAsyncGenerationWorkervenv with the plain VLLM spec, and the capture leg then reuses it and fails loudly with "cached worker venv predates it" (nemo_gymmissing). Fix before merge: spec-aware venv naming, or prefetching the VLLM_GYM (superset) flavor.🤖 Generated with Claude Code