feat(rollout): integrate MInf ledger capture with checkpointing

[lauradang]

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]

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.

[copy-pr-bot]

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]

❌ 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:

1. 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).
2. 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).
3. PY_EXECUTABLES.MCORE_GYM fails the existing test_every_extras_py_executable_is_wired_to_an_actor (virtual_cluster.py).
4. The runtime ACTOR_ENVIRONMENT_REGISTRY swap is inert in the official container, whose prebuilt MegatronPolicyWorker venv lacks nemo_gym (setup.py).
5. 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]

/ok to test a9b17f2

[github-actions]

✅ 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! ✨

[github-actions]

✅ 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! ✨

[github-actions]

✅ Submodule Fast-Forward Check Results

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

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

[lauradang]

/ok to test 4f7beb9

[lauradang]

/ok to test 48861e3

[github-actions]

✅ Submodule Fast-Forward Check Results

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

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

[lauradang]

/ok to test afb4413

[github-actions]

✅ Submodule Fast-Forward Check Results

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

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

[lauradang]

/ok to test 1f9eb12

[github-actions]

✅ Submodule Fast-Forward Check Results

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

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

[lauradang]

/ok to test 0f38893

[github-actions]

✅ Submodule Fast-Forward Check Results

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

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

[terrykong]

Static review by a team of Claude Code agents. There is no GPU here, so linters and tests were not run locally; CI results are cited from the job logs.

Explainer page comparing the vLLM and Megatron capture paths: https://terrykong.github.io/gh-pages-poc/terryk/pr-4129-two-capture-paths.html

The shared pieces are well done. Both backends now resolve prefixes through one ChainPrefixCache / resolve_admission_prefix, splice through one replace_prefix_tokens, and write through one TQTokenSink. The GenerationInterface defaults give a third backend a clear error. Passing the sink and source into the Megatron stager and preparer is what lets their tests run on CPU. The fixes from the earlier review rounds hold at head. Thanks also for landing Megatron-Bridge #6180 and re-pinning to its merge commit.

All external dependencies have landed (Gym #2823, Megatron-LM #7015, Megatron-Bridge #6180). The main thing left before merge: the Megatron recovery test has never run. It is not tagged fast (L1_Functional_Tests_SingleController.sh#L202), and its self-skip can hide failures (grpo_async_gym_single_controller_sibling_recovery.sh#L21-L29). The other inline comments are smaller fixes, including an out-of-date design doc paragraph.

FYI, no action needed:

• fast_L1_Functional_Tests_Megatron_4 passed at the new head with token_mult_prob_error 1.093 (job), close to the 1.10 limit.
• Open PR #4139 (another Megatron-Bridge bump) changes the same pyproject.toml pins (flashinfer, cutlass-dsl, the flashinfer index URL). It does not include the Megatron-LM async-checkpoint removal that this PR adapts to in community_import.py and megatron_policy_worker.py. Whichever PR lands second will need a rebase.

Generated by Claude Code

[terrykong]

1 action item.

TL;DR — one bad manifest row makes the finalizer reject the whole receipt as invalid_receipt with no staging keys. The rollout's good rows then stay in TQ until restart, and this PR's new terminal_selection=None handling never runs.

When a row fails CallRecord validation (L987-L992), _assemble_receipt still ships every raw row (L1035). The finalizer's RolloutReceipt.model_validate fails on that same row (records.py#L365) and returns rejected(f"invalid_receipt:{error}", []) (rollout_reassembler.py#L179-L182).

AI-1

In _assemble_receipt, validate row by row and ship only the rows that parse (not a suggestion block: the fix also needs a test in another file):

        parsed_records: list[CallRecord] = []
        valid_rows: list[dict] = []
        for record in deduped.values():
            try:
                parsed_records.append(CallRecord.model_validate(record))
                valid_rows.append(record)
            except ValueError:
                selection_reason = "invalid_manifest_row"
        if selection_reason is None:
            ...  # attribution as today, over parsed_records
        ...
            "manifest": valid_rows,

Add a test that feeds _assemble_receipt output with one bad row into finalize_rollout and expects rollout_failed:invalid_manifest_row with the good row's key in staging_keys.

What is older than this PR, and what this PR adds

• The leak is older than this PR: the base shipped the same raw rows (nemo_gym.py#L1017 at the base). The rows stay until the restart-time sweep, _validate_rollout_recovery_inventory(clear_unreferenced=True). Bad rows come only from a Gym bug or corrupted data, so impact is low.
• This PR adds the terminal_selection=None handling (here, and in the metrics loop at rollout_reassembler.py#L454-L474), which cannot run today because such a receipt always fails validation first. The fix above makes it run. The premise in this reply — that None would change the rejection from rollout_failed:invalid_manifest_row to invalid_receipt — does not hold: it was already invalid_receipt. test_receipt_assembly_heuristic_masks_invalid_manifest_rows checks only the assembled dict, not the finalizer.

[terrykong]

@pthombre can you review this observation and the changes in this module/actor?

[pthombre]

The failure trace described above is valid. However the configured Gym implementation validates records before returning them. With matching Gym versions, the example malformed row should not reach RL through a successful manifest fetch. We can add this as pure defensive hardening, but I do not expect this path to occur in an actual run.

[terrykong]

1 action item.

TL;DR — this field exists only so Gym can tell which request shape to send: vLLM takes the capture fields at the top level, Megatron takes them inside offload_params. If the Megatron path accepted the same request as vLLM, nothing would need to know the backend, and this field (plus the setup code that copies policy.generation.backend into the user's config) could go away.

New in this PR. Everything after the request is already the same for both backends: the reply (text plus ng_commit_coords), the TQ row, the ledger row, and the finalizer. The one place they differ is how Gym builds the request, in its two handlers: VLLMWorkerCaptureHandler and MegatronWorkerCaptureHandler. Gym picks between them with make_external_capture_handler(backend), which is why RL has to pass the backend through _external_staging_backend and writes it into this config at setup.py#L1302.

AI-1

Make the Megatron path accept the same request Gym already sends to vLLM, since vLLM is the reference: the admission as a top-level ng_capture, with the same capture flags. Then:

1. In Gym, send that one request shape for both backends and drop the backend choice.
2. Move the Megatron-only limits that Gym checks today (n=1, text-only messages) into the Megatron side, for example as a raise in TQMegatronPromptPreparer.prepare_prompt (tq_token_sink.py#L431), so they still fail loudly.
3. Delete _external_staging_backend, this field, and the assignment at setup.py#L1302.

The Megatron endpoint is Megatron-LM's own chat server (RL starts it at megatron_worker.py#L717-L753), and it only reads offload_params (chat_completions.py#L797-L800). So reading a top-level ng_capture likely needs a small Megatron-LM change alongside the Gym one.

Not a suggestion block: this spans Megatron-LM, Gym, and RL.

Where the two requests differ today

	vLLM	Megatron
Admission	top-level ng_capture	offload_params.ng_capture
Extra fields	required_prefix_token_ids, return_tokens_as_token_ids=True	none
Checks in Gym	none	n must be 1; text-only messages

On the vLLM side, both fields are declared on RL's own request model (vllm_worker_async.py#L998-L1001).

[lauradang]

Agreed the backend selection should go. I'd flip the direction though: make vLLM accept the Megatron shape rather than the reverse. Megatron-LM main already forwards offload_params opaquely, so the Megatron request works upstream today with no Megatron-LM change, and it's the smaller payload (no engine flags, no prefix ids). Teaching the Megatron endpoint a top-level ng_capture would move the Gym-specific knowledge into Megatron-LM instead.

Plan: Gym sends one shape with the admission in offload_params.ng_capture; RL's vLLM request model gains offload_params and resolves the prefix from staging_chain like the Megatron path does; the n=1 and text-only checks move into TQMegatronPromptPreparer.prepare_prompt; generation_backend, _external_staging_backend, and the setup assignment get deleted. That covers all three of your items with changes in Gym and RL only.

[github-actions]

✅ Submodule Fast-Forward Check Results

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

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

[github-actions]

✅ Submodule Fast-Forward Check Results

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

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

[github-actions]

✅ Submodule Fast-Forward Check Results

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

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

[github-actions]

✅ Submodule Fast-Forward Check Results

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

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

[lauradang]

/ok to test bce284e

[github-actions]

✅ Submodule Fast-Forward Check Results

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

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

[github-actions]

✅ Submodule Fast-Forward Check Results

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

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨