feat(rollout): integrate MInf ledger capture with checkpointing by lauradang · Pull Request #4129 · NVIDIA-NeMo/RL

lauradang · 2026-09-14T20:04:45Z

Summary

unify vLLM and Megatron token capture around Gym's CaptureAdmission, token-free lineage ledger, and the shared TQ staging contract
resolve staged parent tokens before generation so both backends continue agent rollouts from the exact captured token history
carry capture state through Megatron generation and stage the exact OffloadedRequestPayload through RolloutTokenCapture and TQTokenSink
wait for durable TQ staging before returning commit coordinates; Gym writes a token-free CallRecord only after that acknowledgement
assemble a token-free RolloutReceipt at rollout end, then fetch, verify, and linearize staged records into the canonical GRPO sample during finalization
support both colocated and non-colocated SingleController Megatron-Inference deployments while retaining the vLLM capture path

Token capture flow

The two serving backends use different prompt-preparation and generation hooks, but converge at the same durability boundary. A model call becomes eligible for lineage resolution only after its canonical token delta is staged in TQ. Capture failures write a failure row instead of a committed call record, causing finalization to reject or mask the rollout rather than train on incomplete lineage.

The source diagram is tracked in docs/assets/token-capture-ledger-queue-data-flow.dot; the full invariants, terminal-selection rules, cleanup behavior, and fail-closed semantics are documented in docs/design-docs/token-capture-ledger.md.

Dependencies

Base stack: amahishi/partial-rollout-telemetry-v3
Gym capture contract, unified Megatron staging, and Optional RolloutReceipt.terminal_selection: NVIDIA-NeMo/Gym#2823 (submodule currently pinned to its head 37dc751f; must land on Gym main before this merges, then re-pin and regenerate uv.lock)
Megatron inference hooks (RequestPayloadStager / RequestPromptPreparer, typed RequestPromptPreparationResult, and offload_params carrying ng_capture, template_prefix_token_ids, eos_token_id): NVIDIA/Megatron-LM#7015, pinned transitively through NVIDIA-NeMo/Megatron-Bridge#6101. Both dependencies must land before this PR merges, then the Bridge submodule must be re-pinned to its merged commit and uv.lock regenerated.
Tracking: #4163 for the Megatron-backend recovery functional test once the dependencies land.

Testing

Ruff check and format verification pass for the newly touched RL files
Gym scoped pre-commit passes
Gym token-capture and focused model-server suites: 268 passed

Local test limitation

The RL lockfile supports Linux x86_64/aarch64 only, so the focused RL pytest files cannot run natively on this macOS host. An ephemeral dependency run progressed through PyTorch, Transformers, and pyzmq before reaching the NVIDIA-only pynvml runtime dependency; Linux CI provides the authoritative RL test coverage.

copy-pr-bot · 2026-09-14T20:04:48Z

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.

Signed-off-by: Laura Dang <laurad@nvidia.com>

…very.sh Signed-off-by: Laura Dang <lauradang.2000@gmail.com>

Signed-off-by: Laura Dang <lauradang.2000@gmail.com>

Signed-off-by: Laura Dang <laurad@nvidia.com>

… cache Extract the vLLM worker's _fetch_chain_prefix and _resolve_admission_prefix into tq_token_sink.py as ChainPrefixCache and resolve_admission_prefix, and make the worker methods one-line delegates. TQMegatronPromptPreparer now resolves staging chains through the same pair, so both backends share one cached TQ read (256 entries keyed by the chain's last staging key, deepest cached key bounds the fetch to the uncached suffix). The preparer reads the splice boundary from the request-metadata keys the Megatron chat endpoint writes (prefix_splice_suffix_token_ids, prefix_splice_boundary_token_id). The key strings are spelled out here rather than imported so this module stays importable in the vLLM worker and finalizer environments; a test asserts parity with Megatron's constants when Megatron is importable. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Laura Dang <laurad@nvidia.com>

Signed-off-by: Laura Dang <laurad@nvidia.com>

copy-pr-bot · 2026-09-14T20:14:23Z

Auto-sync is disabled for ready for review pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

github-actions · 2026-09-14T20:15:45Z

❌ Submodule Fast-Forward Check Failed

Check based on commit: 5bf90aa (PR #4129 from laurad/minf-rollout-checkpointing)

❌ Submodules that need attention:

Gym: ❌ Commits have DIVERGED from a common ancestor
TARGET (main branch): https://github.com/NVIDIA-NeMo/Gym/commits/fd5e84d6b1c485c80e7ae61553bbd485611c03b4/
CURRENT (PR #4129 from laurad/minf-rollout-checkpointing): https://github.com/NVIDIA-NeMo/Gym/commits/c3d7cf4ff484d0f4f07898aedeaba05eb4821438/

Please ensure all submodule commits are fast-forwards of the main branch before merging.

lauradang

Team review of PR #4129 (7 agents: rl-expert, expert-gym, expert-megatron-lm, bug-finder, test-agent, design-reviewer, devil-advocate). Static review only: this host is macOS without GPUs and the lockfile is Linux-only, so no tests or linters were run; every claim about upstream behavior is linked to source at the pinned SHAs.

Not landable at HEAD — five deterministic blockers, all inline:

The Megatron capture path programs against Megatron-LM hooks (payload_stager, prompt_preparer, PREFIX_SPLICE_*) that exist only on unmerged Megatron-LM #7015; the Megatron-Bridge pointer is unchanged, so setup_token_capture raises on every Megatron capture run (megatron_worker.py).
The Gym pointer c3d7cf4f is not reachable from any NVIDIA-NeMo/Gym ref and Gym main is wire-incompatible with this RL (3rdparty/Gym-workspace/Gym).
PY_EXECUTABLES.MCORE_GYM fails the existing test_every_extras_py_executable_is_wired_to_an_actor (virtual_cluster.py).
The runtime ACTOR_ENVIRONMENT_REGISTRY swap is inert in the official container, whose prebuilt MegatronPolicyWorker venv lacks nemo_gym (setup.py).
test_prefix_splice_keys_match_megatron_constants hard-fails the Nemo_Gym shard, and the bare checkpointing.save_data_plane=true override fails the streaming-recovery L1 test.

Lint: not run here; two findings (import order in megatron_generation.py, a stray blank line in test_checkpointing.py) would fail pre-commit run --all-files in CI.

The design itself is sound where it is checkable: the durability boundary is identical across backends, TQMegatronTokenStager.stage() fails closed, ChainPrefixCache / resolve_admission_prefix are a clean dedup, the set_generation_epoch rank-0 fan-out is correct, and all Gym API usage matches c3d7cf4f.

Informational (not staged as findings, below the confidence bar): at the #7015 tip both engine hooks run synchronously on the MP-coordinator engine loop, so each captured request costs one or two ray.get TQ round trips during which that DP replica does not step; vLLM offloads the same fetch via asyncio.to_thread. Worth a capture-on/off throughput number once the pin lands.

Devil's advocate: 24 confirmed, 1 disputed (a claimed isort collapse of the reassembler import), 7 downgraded below threshold.

Generated by Claude Code

lauradang · 2026-09-14T22:51:08Z

@@ -1 +1 @@
-Subproject commit fd5e84d6b1c485c80e7ae61553bbd485611c03b4
+Subproject commit c3d7cf4ff484d0f4f07898aedeaba05eb4821438


1 action item.

TL;DR — c3d7cf4f is not reachable from any NVIDIA-NeMo/Gym ref, and Gym main is wire-incompatible with this RL, so the pointer cannot fast-forward until the Gym side lands.

PR-introduced. git merge-base --is-ancestor c3d7cf4f fd5e84d6 → 1 (the bot's fast-forward check already fails). git ls-remote on NVIDIA-NeMo/Gym has no ref reaching c3d7cf4f; it is 40 commits stacked on open Gym #2278 on the lauradang/Gym fork, and .gitmodules points at NVIDIA-NeMo/Gym.git branch=main shallow=true. Against Gym main today, begin_call has no weight_version kwarg (RL passes it in TQMegatronTokenStager) and TokenIdCaptureSettings is extra="forbid" with no external_staging_backend, so the Gym config RL builds at nemo_gym.py:537 fails validation at Gym startup.

The open Gym PR for this work appears to be #2823, whose head is neither an ancestor nor a descendant of c3d7cf4f and which nests the Megatron admission under request_metadata.ng_capture (pinned c3d7cf4f writes it top-level).

AI-1

Action: Land Gym #2278 / #2823 (or a consolidated equivalent) on Gym main, then point this submodule at a main SHA. Before that, reconcile the Megatron wire contract (top-level ng_capture vs request_metadata.ng_capture) between RL, the Gym PR, and the Megatron-LM endpoint, and link the Gym PR instead of a fork SHA in the description.

Context — no action (Gym-side). At c3d7cf4f, _strip_capture_transport_fields does not strip compact_prompt_token_ids, which the Megatron endpoint emits under return_tokenized_data (chat_completions.py), so a prompt-length token array leaks to the agent hop per call. Gym #2823 already strips it, which is one more reason to follow that lineage.

partially fixed in a9b17f2 — the pointer now targets the head of Gym #2823 (fc08bf19), which is reachable from NVIDIA-NeMo/Gym, fast-forwards from Gym main, and nests ng_capture under request_metadata as the Megatron-LM #7015 endpoint requires. Still remaining: Gym #2823 must land on Gym main before this can merge, and the PR description should link Gym #2823 instead of the fork SHA.

- Declare the nemo_gym extra on MegatronPolicyWorker in actor_environments.py (the venv source of truth) instead of swapping ACTOR_ENVIRONMENT_REGISTRY at runtime, which the prebuilt container venv ignored; drop PY_EXECUTABLES.MCORE_GYM and the ModuleNotFoundError string-match remediation. - Gate backend=megatron token capture at setup on the MInf capture hook protocols (RequestPayloadStager / RequestPromptPreparer) from Megatron-LM PR #7015, failing with a NotImplementedError that names the dependency while the Megatron-Bridge pin predates it. - Point the Gym submodule at Gym PR NVIDIA-NeMo#2823 (fc08bf19), which is reachable from NVIDIA-NeMo/Gym, fast-forwards from Gym main, and nests ng_capture under request_metadata as the Megatron endpoint requires. - Skip test_prefix_splice_keys_match_megatron_constants when the pinned megatron-core lacks the constants, and importorskip megatron.core in the Megatron hosting test, so the Nemo_Gym shard skips instead of aborting. - Restore the ++ Hydra override for checkpointing.save_data_plane in the streaming recovery script (the key is absent from the Gym config chain). - Document the Megatron-LM #7015 dependency in the design doc. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Laura Dang <laurad@nvidia.com>

lauradang · 2026-09-15T07:57:22Z

/ok to test a9b17f2

github-actions · 2026-09-15T07:57:57Z

✅ Submodule Fast-Forward Check Results

Check based on commit: a9b17f2 (PR #4129 from laurad/minf-rollout-checkpointing)

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

- Add design-docs/token-capture-ledger.md to the docs toctree and drop links to guides that do not exist yet (Sphinx treats both as errors). - Read vllm_cfg / mcore_generation_config through a dict cast in the token-capture validation so pyrefly does not reject the TypedDict keys. - Apply ruff formatting to megatron_worker.py and test_checkpointing.py. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Laura Dang <laurad@nvidia.com>

github-actions · 2026-09-15T17:32:37Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 33bef60 (PR #4129 from laurad/minf-rollout-checkpointing)

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

…anup Stamp a Megatron request that spans a refit with its admission (oldest) policy epoch instead of failing capture and masking the whole rollout. This matches vLLM, which freezes the version at begin_call, and the finalizer's min-over-calls group tag; spans are counted (epoch_span_count) and logged at WARNING. Also fold in the mechanical review items: drop the undefined rollout_max_attempts_to_avoid_lp_nan knob, remove the dead KeyError branch in receipt parsing, route fetch_for_finalization through TQStagingStore, collapse the duplicate prefix predicate, call set_generation_epoch directly, fix import order, revert diff churn, fix the FIFO docstring and the manifest control route in the design doc, and remove the unused logical_request_id test parameter. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Laura Dang <laurad@nvidia.com>

github-actions · 2026-09-15T22:33:22Z

✅ Submodule Fast-Forward Check Results

Check based on commit: d678aa7 (PR #4129 from laurad/minf-rollout-checkpointing)

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

macandro96

Re-review: PR #4129 at `d678aa7`

Main ones are:

The targeted NeMo-Gym unit subset fails with 14 failed, 38 passed, 1 skipped. Two test fixtures have drifted from the pinned Gym schemas; see inline comments.
uv lock --check is currently failing. The Gym submodule changed package requirements, but the root uv.lock still records the previous Gym dependency metadata. The check currently fails during re-resolution on a missing flash-attention wheel hash.
Megatron capture/recovery still lacks an end-to-end functional test.
The external dependencies remain open: Gym #2823 and Megatron-LM #7015.

Generated by Claude Code

…dapter TQMegatronTokenStager now hands the offloaded payload to RolloutTokenCapture.complete_call_from_response with Gym's MegatronCaptureAdapter instead of extracting fields by hand. Malformed payloads poison the call with capture_failed coordinates, which Gym records as worker_capture_failed (matching vLLM) rather than worker_response_missing_commit_coordinates. Requires the Gym-side adapter from NVIDIA-NeMo/Gym#2823. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Laura Dang <laurad@nvidia.com>

tdene · 2026-09-16T15:58:15Z

Reminder that before merging, you will have to first merge the Gym side, bump the Gym commit hash, and re-generate uv.lock.

- add parent_chain_hash to the inline-prefix token_in admission - drop logical_request_id from _manifest_record (no such CallRecord field) Signed-off-by: Laura Dang <laurad@nvidia.com> Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

lauradang · 2026-09-16T21:55:21Z

/ok to test eadf9e5

Signed-off-by: Laura Dang <laurad@nvidia.com> Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

- reword _assemble_receipt docstring so the named failure reasons are examples, not an exhaustive list; add unresolved_parent and the capture_failed fallback for reason-less rows - terminal_selection is left unchanged: the pinned Gym RolloutReceipt (fc08bf19) types it Literal["declared","response_id","content", "heuristic"] with no None, so the invalid_manifest_row path cannot report "no heuristic ran" without a Gym-side schema change Signed-off-by: Laura Dang <laurad@nvidia.com> Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

…rics code Signed-off-by: Laura Dang <laurad@nvidia.com> Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

…wiring Signed-off-by: Laura Dang <laurad@nvidia.com> Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

lauradang · 2026-09-16T22:05:15Z

/ok to test 2c9baf1

…fload_params Megatron Inference renamed the opaque per-request dict it forwards to the payload stager and prompt preparer (NVIDIA/Megatron-LM PR #7015); follow it on the TQMegatronPromptPreparer / TQMegatronTokenStager side. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Laura Dang <laurad@nvidia.com>

…_prefix_tokens The Megatron chat endpoint now ships the chat-template render through the last assistant message (template_prefix_token_ids) and the EOS id in offload_params instead of a precomputed boundary (NVIDIA/Megatron-LM PR #7015). TQMegatronPromptPreparer feeds those, plus the prefix it resolved from the staging chain, to the same replace_prefix_tokens the vLLM worker uses, so one splicing algorithm serves both backends. replace_prefix_tokens gains a keyword-only eos_token_id override so callers that hold only token ids can use it without a tokenizer; the boundary failure message skips the detokenized reprs in that case. Existing callers are unchanged. Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com> Signed-off-by: Laura Dang <laurad@nvidia.com>

Signed-off-by: Laura Dang <laurad@nvidia.com> Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

- bump Gym to 37dc751f (RolloutReceipt.terminal_selection is now Optional) - _assemble_receipt no longer pre-labels parse failures as heuristic - finalizer metric loop skips the None member of the Optional annotation Signed-off-by: Laura Dang <laurad@nvidia.com> Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

lauradang · 2026-09-17T06:10:58Z

/ok to test 5a13a7c

Signed-off-by: Laura Dang <laurad@nvidia.com>

lauradang requested review from a team as code owners September 14, 2026 20:04

github-actions Bot added the Documentation Improvements or additions to documentation label Sep 14, 2026

lauradang and others added 10 commits September 14, 2026 13:09

feat(rollout): integrate MInf ledger capture with checkpointing

710c0c4

Signed-off-by: Laura Dang <laurad@nvidia.com>

Update tests/functional/grpo_async_gym_single_controller_sibling_reco…

bc5bd59

…very.sh Signed-off-by: Laura Dang <lauradang.2000@gmail.com>

Update nemo_rl/models/generation/megatron/megatron_generation.py

2297b18

Signed-off-by: Laura Dang <lauradang.2000@gmail.com>

Update tests/unit/single_controller/test_checkpointing.py

5eb7630

Signed-off-by: Laura Dang <lauradang.2000@gmail.com>

test(data-plane): parameterize Megatron capture cases

ca2d395

Signed-off-by: Laura Dang <laurad@nvidia.com>

fix(megatron): import logging for generation logger

e1b1415

Signed-off-by: Laura Dang <laurad@nvidia.com>

fix(gym): initialize rollout retry limit

3f8f7e2

Signed-off-by: Laura Dang <laurad@nvidia.com>

fix(data-plane): retain client in token source

615bd69

Signed-off-by: Laura Dang <laurad@nvidia.com>

chore(data-plane): trim prompt preparer comments

5bf90aa

Signed-off-by: Laura Dang <laurad@nvidia.com>

lauradang force-pushed the laurad/minf-rollout-checkpointing branch from 033f7a7 to 5bf90aa Compare September 14, 2026 20:14

lauradang commented Sep 14, 2026

View reviewed changes

lauradang added the CI:L0 Run doctests and unit tests label Sep 15, 2026

terrykong requested a review from macandro96 September 15, 2026 20:18

macandro96 reviewed Sep 15, 2026

View reviewed changes

lauradang mentioned this pull request Sep 16, 2026

feat(token-capture): add Megatron worker capture backend NVIDIA-NeMo/Gym#2823

Open

tdene mentioned this pull request Sep 16, 2026

MInf: Add hooks to preprocess and stage NVIDIA/Megatron-LM#7015

Open

6 tasks

tdene reviewed Sep 16, 2026

View reviewed changes

lauradang and others added 4 commits September 16, 2026 14:59

fix(token-capture): stamp the admission epoch explicitly for MInf spans

b8fafc3

Signed-off-by: Laura Dang <laurad@nvidia.com> Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

docs(token-capture): align ledger design doc with attribution and met…

e1b690c

…rics code Signed-off-by: Laura Dang <laurad@nvidia.com> Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

test(token-capture): cover Megatron capture setup guards and backend …

2c9baf1

…wiring Signed-off-by: Laura Dang <laurad@nvidia.com> Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

lauradang and others added 2 commits September 16, 2026 22:46

lauradang mentioned this pull request Sep 17, 2026

test: add Megatron-backend variant of Gym token-capture recovery functional test #4163

Open

lauradang and others added 2 commits September 16, 2026 23:04

refactor(generation): declare token-capture hooks on GenerationInterface

0150359

Signed-off-by: Laura Dang <laurad@nvidia.com> Co-Authored-By: Claude Fable 5.1 <noreply@anthropic.com>

fix(token-capture): consume typed prompt preparation result

68b63da

Signed-off-by: Laura Dang <laurad@nvidia.com>

lauradang requested a review from a team as a code owner September 17, 2026 16:35

		@@ -1 +1 @@
		Subproject commit fd5e84d6b1c485c80e7ae61553bbd485611c03b4
		Subproject commit c3d7cf4ff484d0f4f07898aedeaba05eb4821438

Conversation

lauradang commented Sep 14, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Token capture flow

Dependencies

Testing

Local test limitation

Uh oh!

copy-pr-bot Bot commented Sep 14, 2026

Uh oh!

copy-pr-bot Bot commented Sep 14, 2026

Uh oh!

github-actions Bot commented Sep 14, 2026

❌ Submodule Fast-Forward Check Failed

❌ Submodules that need attention:

Uh oh!

lauradang left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

lauradang Sep 14, 2026

Choose a reason for hiding this comment

AI-1

Uh oh!

lauradang Sep 15, 2026

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!

lauradang commented Sep 15, 2026

Uh oh!

github-actions Bot commented Sep 15, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Sep 15, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Sep 15, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

macandro96 left a comment

Choose a reason for hiding this comment

Re-review: PR #4129 at d678aa7

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

tdene Sep 16, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

lauradang commented Sep 16, 2026

Uh oh!

lauradang commented Sep 16, 2026

Uh oh!

lauradang commented Sep 17, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

lauradang commented Sep 14, 2026 •

edited

Loading

Re-review: PR #4129 at `d678aa7`