feat(sc): gate-authoritative token capture on the Gym tokidcap stack - #3455
Closed
pthombre wants to merge 44 commits into
Closed
feat(sc): gate-authoritative token capture on the Gym tokidcap stack#3455pthombre wants to merge 44 commits into
pthombre wants to merge 44 commits into
Conversation
Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: Yuki Huang <yukih@nvidia.com>
… reserve/commit slots Signed-off-by: Yuki Huang <yukih@nvidia.com> update unit test Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: Yuki Huang <yukih@nvidia.com>
…um_prompts_per_step Signed-off-by: Yuki Huang <yukih@nvidia.com>
…unctional test buffer sizing Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: Yuki Huang <yukih@nvidia.com>
…troller and grpo_async_gym_single_controller into it Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: Yuki Huang <yukih@nvidia.com>
Signed-off-by: Yuki Huang <yukih@nvidia.com>
…, weight-version stamping RL half of stage S2 (all dormant until the setup_token_capture fan-out, which S4 wires behind token_capture.enabled): - vllm_worker_async: install_token_capture (Gym CaptureHost seam), setup_token_capture fan-out target (in-worker DP client + TQTokenSink + the single install_capture call with the vLLM adapter; model owners only), _rollout_weight_version + set_rollout_weight_version. - vllm_generation: setup_token_capture / set_rollout_weight_version DP-leader fan-outs (async engine asserted). - single_controller: flag-gated set_rollout_weight_version rotation in _sync_weights beside RolloutManager.set_weight_version. - PY_EXECUTABLES.VLLM_GYM (--extra vllm --extra nemo_gym): capture-enabled worker env; verified nemo_gym capture core imports beside vllm 0.20.0. - pyrefly: tq_token_sink.py type-checked; nemo_gym.* replace-imports-with-any. - Gym submodule pin -> ccf9b6f6: S2 engine-blind capture core (staging/capture.py, fail-closed stage-then-respond ordering) + adapters/vllm.py (relocated replace_prefix_tokens splice, native id+logprob extraction) + 29 tests (75/75 green) + ruff format pass. - Tests: test_vllm_token_capture_hosting.py 6/6 (--nemo-gym-only). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>
Gym fork gains the gate: lineage hosting in the vLLM model server (memory token buffer, marker resolve + history fingerprint, § 3.3 serving rule, exact prefix serving, coords ingestion as the authoritative commit, token-free receipts, TTLs), register/seal/fail control routes + RolloutControlClient, and marker plumbing through the responses converter. All dormant behind token_capture_gate.enabled=false; Gym suite 1507 green. Also records the S1 functional pin-bump evidence in the implementation log (grpo_dp_single_controller.sh and grpo_async_gym_single_controller.sh both PASS flag-off after NRL_FORCE_REBUILD_VENVS=true; stale prebaked container venvs documented as environmental). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>
…tion
NeMo-RL side of the gate-authoritative pipeline (no Gym fork changes):
blackbox_finalizer (receipt → TQ fetch → § 5 verification → linearize →
always-N publish with placeholders, group_min/max_wv, group-drop and
mixed-wv policies), worker request-path capture (begin/finish/abort around
chat completions; coords ride ng_commit_coords, logprobs stripped),
receipt-mode run_rollouts + gate control plane in the NemoGym env
(register/seal/fail, NaN-retry hard error), capture dispatch in
RolloutManager (rollout-id minting {group_id}_g{i}, commit_finalized,
failure-path aborts), validity-aware GRPO baseline via sample_mask,
setup wiring (MVP-matrix validation, VLLM_GYM registry override,
setup_token_capture + weight-version fan-outs, finalizer threading),
exemplar YAML token_capture block. All dormant behind
token_capture.enabled=false; flag-off regression 522 passed.
Gate evidence (2×H100, 10 steps, capture enabled): both metric checks
PASS — median gen_kl_error 0.0375 (flag-off run: 0.038), max reward 0.5,
global_valid_seqs=8.0 every step (zero placeholders trained); cancelled
dispatches exercised the § 7 failure path loudly and leaked nothing.
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>
…pture L1 test, SWE A/B runbook - gate/* metrics into the SC logger per train step (token_in_rate derived; fetch failures logged, never fatal) via RolloutManager.gate_metrics() - env-gated train-row dump (NRL_SC_DUMP_TRAIN_ROWS) at both canonical publish sites for the legacy-vs-capture row diff - env-gated HTTP byte counters: RL vLLM-worker mirror (NRL_HTTP_BYTES_DIR) of the Gym middleware; Gym submodule pin 05986b04 -> e3b3eac6 (NG_HTTP_BYTES_DIR middleware) - fix: re-raise aiohttp.ClientResponseError as picklable RuntimeError in run_rollouts (CIMultiDictProxy headers cannot cross the Ray boundary; active flag-off, disclosed for the S5 gate) - capture-enabled SC gym functional test in L1 (full mode) - swe/: token-capture vs legacy perf A/B runbook + launch tooling - design docs (v2 + gate-authoritative + implementation log) into the docs tree S5 evidence recorded in docs/design-docs/tq-gym-gate-authoritative-implementation-log.md. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>
Env-agnostic SLURM launcher (ray.sub, 16 nodes: 8 train + 8 gen, SC_MODE=1 single-controller default) and Qwen3-30B-A3B-Thinking recipe, imported from the sc-test workspace (biguo) so the token-capture SWE A/B (swe/SWE_RUN.md) can snapshot from this branch. Site-specific values (account, container, secrets, data/model paths) come from a personal wrapper — the published yaml carries placeholder sandbox-image paths. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>
…b env The launcher predates the S5 byte counters; without forwarding, BYTES=1 in swe/launch_swe_ab.sh silently produced no per-hop accounting inside the sbatch job. Empty defaults keep the counters dormant. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>
Concrete experiment plan (exact commits/snapshot/compute shape/gates/ pre-flight checklist) for the run swe/SWE_RUN.md describes generically. Status: ready, not submitted. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>
A node-local GYM_VENV_DIR=/tmp rebuild only materializes on the NemoGym actor's node while Gym spawns servers cluster-wide — job 14542017 failed with missing venv pythons on every other node. The baked venvs are dep-compatible with the fork (swe_agents/vllm_model requirements unchanged vs the old pin) and editable installs carry the fork's code. Also record the uv-cache pre-warm requirement and the reaper's de facto 60-min idle kill in the runbook. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>
setup.py inits the reference model only when the KL penalty is positive, but AdvantageConfig.reference_logprobs_field defaults to set — the SC then requests reference logprobs from a worker that never built reference_state_dict and dies with AttributeError in the first train step (SWE A/B job 14545431; the exemplar configs mask this with penalty 0.01). Null the field under the same condition. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>
Covers recursive clone with the published Gym fork pin (NVIDIA-NeMo/Gym pthombre/tq-gate-capture @ e3b3eac6), the token_capture config surface, the 2-GPU capture-enabled functional smoke, the SWE A/B tooling, gate metrics to watch, and env-gated debug switches. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>
RL companion of the Gym branch tq-tokidcap-capture (submodule re-pinned here to b6051536 = upstream stack top 81ac2736/#2182 + our 7-commit staging/gate series). Executes docs/design-docs/tq-gym-tokidcap-migration.md; the working log with the workaround ledger and bump checklist is added as docs/design-docs/tq-gym-tokidcap-migration-log.md. Identity carrier switch (§3): drop the responses_create_params.metadata side-channel; rollout ids ride the run body as the opaque _ng_rollout_id key, agents stamp /ng-rollout/<id> on every model call via the stack's helpers, and the middleware-minted model_call_id becomes the call id — so the TQ sample id IS the capture key end to end. Registration stays create-only, pre-dispatch. Gate hosting config: TokenCaptureConfig grows lineage capacity (derived at setup from in-flight rollouts x max sequence length — finding M), a per-run control-plane bearer token (finding S), a hard per-call control deadline (S5 silent-stall finding), and the base capture dir the gate rides on (#2124-c1 workaround). nemo_gym.py injects the gate + capacity + auth into the policy model server, activates the base token capture layer, and opts every agent into correlation (token_id_capture_all_agents). Unchanged by design: TQTokenSink/TQTokenSource (keys are minted Gym-side), BlackboxFinalizer (the new terminal-aware linearize keeps its contract and surfaces unresolved retries as RebuildError -> the existing placeholder path), the worker capture hosting, and the weight-version fan-out. Dormant by default: everything is behind token_capture.enabled=false. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>
The new Gym base drops Gym's docs dependency-group, so the locked package set shrinks accordingly (934 deletions, no other changes; lockfile revision unchanged). Co-Authored-By: Claude Fable 5 <noreply@anthropic.com> Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>
github-actions
Bot
added
Documentation
Contributor
Author
|
Companion Gym PR (the gate/staging layer on the tokidcap stack): NVIDIA-NeMo/Gym#2278 /ok to test 624bb27 |
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
Token-in/token-out capture for NeMo-Gym rollouts in the async SingleController GRPO path, migrated onto the upstream Gym token-id-capture PR stack (NVIDIA-NeMo/Gym #2190→#2124→#2125→#2126→#2180→#2181→#2182) per
docs/design-docs/tq-gym-tokidcap-migration.md. The Gym submodule is pinned topthombre/tq-tokidcap-capture(= stack top81ac2736+ the 7-commit gate/staging series; companion Gym PR opened against the stack).With
token_capture.enabled=true: the Gym gate holds each rollout's token lineage (baseLineageIndex, capacity sized from the training config), vLLM workers stage per-call token deltas + logprobs directly into TransferQueue before the response releases, agent-facing messages and the Ray return are token-free, and theBlackboxFinalizerrebuilds verified always-N training rows from staged deltas. Withenabled=false(default) every legacy codepath is byte-identical.Highlights of this migration vs. the previous fork-based MVP:
RolloutManagerputs{group_id}_g{i}on the run body (_ng_rollout_id); agents stamp/ng-rollout/<id>on every model call via stack helpers (all agent impls unmodified); the middleware-mintedmodel_call_idis the call id. TQ sample ids ARE the capture keys.openai_utils/responses_converterplumbing); fallback causes are instrumented (ambiguityseparately — the data-driven trigger for a future marker).run_builder(prefix_merging); unresolved final-call retries publish placeholders instead of training on a guess.Working log with the workaround ledger (upstream asks → contained workarounds → what deletes when each ask lands) and a stack-bump checklist:
docs/design-docs/tq-gym-tokidcap-migration-log.md.Tests
--nemo-gym-only) + replay buffer / rollout manager / SC setup suites; 2-GPU flag-off functional (grpo_async_gym_single_controller.sh) as the pin-bump regression and the same with++token_capture.enabled=true(results recorded in the migration log).🤖 Generated with Claude Code