feat(rollout): integrate MInf ledger capture with checkpointing - #3869
Conversation
github-actions
Bot
added
the
Documentation
|
/ok to test 0ac1707 |
There was a problem hiding this comment.
Review of the lauradang commits only (08faaee37, 0c5e033ec, 0ac170793). Verified against the pinned Gym submodule (5cd5c0d), the pinned Megatron-Bridge → Megatron-LM (0c565c9a0 → d12f6c8c9), and the tdene:Megatron-LM:tde/ledger_capture branch. 4 inline comments below, plus one Gym-side issue on the submodule bump that I can't anchor inline:
Also verified OK, for the record: the exactly-one-rank expectation in flush_token_capture matches the MInf branch (ledger is populated only on the MP-coordinator rank per DP engine; fan-out-and-union is the documented pattern), and set_generation_epoch from global rank 0 propagates to all engines via the DP coordinator.
|
For other reviewers, note that NVIDIA/Megatron-LM#7015 is the "MInf metadata ledger" PR linked to this. |
tdene
left a comment
tdene
left a comment
There was a problem hiding this comment.
A functional test is also important to add.
| ) | ||
| ray.get(futures) | ||
|
|
||
| def flush_token_capture(self, receipt: dict[str, Any]) -> dict[str, Any]: |
There was a problem hiding this comment.
We enter this from _flush_and_seal after checking only pending_manifest, but we should also be checking capture_poisoned. Right now, poisoned rollouts that are still carrying pending_manifest will get flushed and sealed, even though it will just be thrown away in the end.
I think the correct fix is to early exit out of flush_token_capture when receipt["capture_poisoned"] or receipt.get("capture_poisoned").
| raise RuntimeError( | ||
| f"MInf request {uid!r} was present on {len(matches)} engine ranks; expected 1" |
There was a problem hiding this comment.
I don't think this should be handled right now, but I just want to bring attention to it.
Today, rollout checkpointing only happens on training boundaries. We lose everything within one training step.
But MLM and MRL have the ability to checkpoint continuously. Once that ability comes into play, we will have situations where matches==0 is valid, because we're loading checkpointed requests.
| setter(version) | ||
|
|
||
| @staticmethod | ||
| def _serialize_token_capture_record(record) -> dict: |
There was a problem hiding this comment.
This serializes both routing_indices and prompt_log_probs, but neither is consumed. In particular, routing_indices is absolutely huge (order of magnitude of 1KB per token).
| receipt=receipt, | ||
| reward=completion.reward, | ||
| ) | ||
| await _flush_and_seal() |
There was a problem hiding this comment.
This technically holds the barrier for longer than it strictly needs to. I'm not sure it matters though. verify is probably going to be the slowest thing that can be taken out of the barrier, but I really don't know.
| sink = TQTokenSink(dp_client, staging_partition=staging_partition) | ||
| self._token_capture = RolloutTokenCapture( | ||
| sink=sink, | ||
| weight_version_fn=lambda: 0, |
There was a problem hiding this comment.
This kind of "default to 0" has caused us trouble before in Megatron RL. I don't think we should be nice. I think we should raise if we don't have an actual weight_version. So, hah, maybe lambda: raise Error would be more appropriate.
6171c18 to
9c48021
Compare
70940a9 to
deabd07
Compare
❌ Submodule Fast-Forward Check FailedCheck based on commit: e10371b (PR #3869 from ❌ Submodules that need attention:Gym: ❌ Commits have DIVERGED from a common ancestor Please ensure all submodule commits are fast-forwards of the amahishi/partial-rollout-telemetry-v3 branch before merging. |
a4c10a7 to
23d1295
Compare
❌ Submodule Fast-Forward Check FailedCheck based on commit: 1de0d28 (PR #3869 from ❌ Submodules that need attention:Gym: ❌ Commits have DIVERGED from a common ancestor Please ensure all submodule commits are fast-forwards of the amahishi/partial-rollout-telemetry-v3 branch before merging. |
❌ Submodule Fast-Forward Check FailedCheck based on commit: a7cea3c (PR #3869 from ❌ Submodules that need attention:Gym: ❌ Commits have DIVERGED from a common ancestor Please ensure all submodule commits are fast-forwards of the amahishi/partial-rollout-telemetry-v3 branch before merging. |
a7cea3c to
fa4ce96
Compare
|
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. |
|
/ok to test fa4ce96 |
❌ Submodule Fast-Forward Check FailedCheck based on commit: fa4ce96 (PR #3869 from ❌ Submodules that need attention:Gym: ❌ Commits have DIVERGED from a common ancestor Please ensure all submodule commits are fast-forwards of the amahishi/partial-rollout-telemetry-v3 branch before merging. |
| @@ -1 +1 @@ | |||
| Subproject commit fd5e84d6b1c485c80e7ae61553bbd485611c03b4 | |||
| Subproject commit c3d7cf4ff484d0f4f07898aedeaba05eb4821438 | |||
There was a problem hiding this comment.
1 action item — merge-blocking; please fix in this PR. PR-introduced.
TL;DR — the Gym pin c3d7cf4 is a fork branch with no shared history with the base branch's pin; merging it reverts ~215 upstream Gym commits (730 deleted files).
Evidence (gh api repos/NVIDIA-NeMo/Gym/compare/...):
main...c3d7cf4→diverged, ahead 65, behind 215 (merge base3a5afd9, 2026-08-19)fd5e84d...c3d7cf4(base pin → PR pin) →diverged, ahead 65, behind 183;main...fd5e84d→behind 32, so the base pin is on main.git diff --diff-filter=D fd5e84d c3d7cf4 | wc -l→ 730 deleted files,+7365/−134155:nemo_gym/telemetry,nemo_gym/sandbox/providers/{e2b,local},nemo_gym/health,nemo_gym/exporters, a dozenresources_servers/*,tests/unit_tests, …c3d7cf4is only onlauradang/Gym:laurad/megatron-adapter(Gym PR feat: video + audio understanding GRPO training recipe #2823, which targetspthombre/tq-tokidcap-capture— already squash-merged to main as fix(megatron): honor policy.logprob_chunk_size in the training loss path #2872 = the base pin).
This is what the "Submodule Fast-Forward Check: Commits have DIVERGED" failure on this PR is reporting. RL is also coupled to fork-only surface (external_staging_backend — main's TokenIdCaptureSettings is extra="forbid"; begin_call(…, weight_version=); the megatron_worker handler in external_capture.py; CallRecord.logical_request_id), so re-pinning to main alone would break the RL side.
AI-1
Re-port Gym #2823 onto Gym main, merge it, then pin RL to a SHA where gh api repos/NVIDIA-NeMo/Gym/compare/main...<sha> reports identical or behind. Consider making the fast-forward check a required status.
| ) | ||
| missing = [ | ||
| name | ||
| for name in ("payload_stager", "prompt_preparer") |
There was a problem hiding this comment.
nemo_rl/models/generation/megatron/megatron_worker.py:778
1 action item, 1 follow-up. PR-introduced.
TL;DR — the pinned Megatron-LM has none of the engine hooks this path needs; they exist only in the still-open NVIDIA/Megatron-LM#7015, so Megatron capture is non-functional as pinned and the PR body understates the dependency.
Megatron-Bridge 5ed97996 pins Megatron-LM 1e7598cb. At that SHA dynamic_engine.py has is_mp_coordinator and inference_client.py has set_generation_epoch, but there is no payload_stager, prompt_preparer, request_metadata, or OffloadedRequestPayload, and chat_completions.py never forwards ng_capture / ng_prompt_suffix_token_ids or emits ng_commit_coords. All of those come from NVIDIA/Megatron-LM#7015 (OPEN). Result today: RuntimeError("… missing payload_stager, prompt_preparer") right here. This is the same merge-order blocker as the earlier local_metadata_ledger_* thread, under new names.
AI-1
State the dependency in the PR description (replace "No separate Megatron metadata-ledger dependency") and in the design doc's MInf section, and name #7015 in this RuntimeError so the remediation is one click away.
Follow-up
Gate merge on #7015 landing and a Megatron-Bridge bump whose 3rdparty/Megatron-LM includes it, then re-lock RL.
| assert result is None | ||
|
|
||
|
|
||
| def test_megatron_stager_declines_mixed_policy_epochs(tq_client, staging_partition): |
There was a problem hiding this comment.
tests/unit/data_plane/test_tq_token_sink.py:361
1 action item. PR-introduced.
stage() catches every exception and returns None (tq_token_sink.py:487-491), so assert result is None also passes for a bad admission, a sink failure, or a typo in _weight_version — it doesn't pin the mixed-epoch rejection. None of the other _weight_version branches (missing/empty/malformed/negative) are covered.
AI-1
Pin the logged reason and test the static method directly (one way to do it; not run here — no Linux env):
with caplog.at_level(logging.ERROR, logger="nemo_rl.data_plane.tq_token_sink"):
result = stager.stage(...)
assert result is None
assert "spans policy epochs [7, 8]" in caplog.text
@pytest.mark.parametrize(
("policy_epoch", "match"),
[
(None, "no policy_epoch boundaries"),
([], "no policy_epoch boundaries"),
([(0,)], "invalid policy_epoch metadata"),
([(0, 7), (1, 8)], r"spans policy epochs \[7, 8\]"),
([(0, -1)], "negative policy epoch -1"),
],
)
def test_megatron_stager_weight_version_rejects_invalid_metadata(policy_epoch, match):
with pytest.raises(ValueError, match=match):
TQMegatronTokenStager._weight_version(SimpleNamespace(policy_epoch=policy_epoch))| "nemo_rl.models.policy.workers.megatron_policy_worker.MegatronPolicyWorker" | ||
| ] = PY_EXECUTABLES.MCORE_GYM | ||
| else: | ||
| ACTOR_ENVIRONMENT_REGISTRY[ |
There was a problem hiding this comment.
nemo_rl/algorithms/single_controller_utils/setup.py:1148
1 action item — low severity. PR-introduced.
This else: re-registers VllmAsyncGenerationWorker → VLLM_GYM, which ray_actor_environment_registry.py:19 already sets unconditionally at import (at base and head). It is the only registry write outside the registry module and reads as if the vLLM capture path depends on it. (The megatron branch's MCORE_GYM registration just above, added in a7cea3c, is the one doing real work — that one should stay.)
AI-1
Delete the vLLM else: block, or turn it into an assert on get_actor_python_env(...) == PY_EXECUTABLES.VLLM_GYM if the intent was a guard.
| def _manifest_record( | ||
| call_id: str, | ||
| *, | ||
| logical_request_id: str | None = None, |
There was a problem hiding this comment.
tests/unit/environments/test_nemo_gym_token_capture.py:66
1 action item — trivial. PR-introduced.
No test passes logical_request_id=; the helper always emits None. Either drop the parameter and the dict entry, or add one assertion that a non-None value survives _assemble_receipt.
❌ Submodule Fast-Forward Check FailedCheck based on commit: 49a200a (PR #3869 from ❌ Submodules that need attention:Gym: ❌ Commits have DIVERGED from a common ancestor Please ensure all submodule commits are fast-forwards of the amahishi/partial-rollout-telemetry-v3 branch before merging. |
❌ Submodule Fast-Forward Check FailedCheck based on commit: f894379 (PR #3869 from ❌ Submodules that need attention:Gym: ❌ Commits have DIVERGED from a common ancestor Please ensure all submodule commits are fast-forwards of the amahishi/partial-rollout-telemetry-v3 branch before merging. |
❌ Submodule Fast-Forward Check FailedCheck based on commit: 87ffe07 (PR #3869 from ❌ Submodules that need attention:Gym: ❌ Commits have DIVERGED from a common ancestor Please ensure all submodule commits are fast-forwards of the amahishi/partial-rollout-telemetry-v3 branch before merging. |
1784afe to
195df04
Compare
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 <lauradang.2000@gmail.com>
87ffe07 to
708bede
Compare
❌ Submodule Fast-Forward Check FailedCheck based on commit: 708bede (PR #3869 from ❌ Submodules that need attention:Gym: ❌ Commits have DIVERGED from a common ancestor Please ensure all submodule commits are fast-forwards of the amahishi/partial-rollout-telemetry-v3 branch before merging. |
Signed-off-by: Laura Dang <laurad@nvidia.com>
❌ Submodule Fast-Forward Check FailedCheck based on commit: 18a3e14 (PR #3869 from ❌ Submodules that need attention:Gym: ❌ Commits have DIVERGED from a common ancestor Please ensure all submodule commits are fast-forwards of the amahishi/partial-rollout-telemetry-v3 branch before merging. |
Signed-off-by: Laura Dang <laurad@nvidia.com>
0262db0 to
0479389
Compare
❌ Submodule Fast-Forward Check FailedCheck based on commit: 0262db0 (PR #3869 from ❌ Submodules that need attention:Gym: ❌ Commits have DIVERGED from a common ancestor Please ensure all submodule commits are fast-forwards of the amahishi/partial-rollout-telemetry-v3 branch before merging. |
❌ Submodule Fast-Forward Check FailedCheck based on commit: 0479389 (PR #3869 from ❌ Submodules that need attention:Gym: ❌ Commits have DIVERGED from a common ancestor Please ensure all submodule commits are fast-forwards of the amahishi/partial-rollout-telemetry-v3 branch before merging. |
Signed-off-by: Laura Dang <laurad@nvidia.com>
❌ Submodule Fast-Forward Check FailedCheck based on commit: c5babff (PR #3869 from ❌ Submodules that need attention:Gym: ❌ Commits have DIVERGED from a common ancestor Please ensure all submodule commits are fast-forwards of the amahishi/partial-rollout-telemetry-v3 branch before merging. |
Signed-off-by: Laura Dang <laurad@nvidia.com>
❌ Submodule Fast-Forward Check FailedCheck based on commit: 812f29f (PR #3869 from ❌ Submodules that need attention:Gym: ❌ Commits have DIVERGED from a common ancestor Please ensure all submodule commits are fast-forwards of the amahishi/partial-rollout-telemetry-v3 branch before merging. |
2 participants
Summary
CaptureAdmission, token-free lineage ledger, and the shared TQ staging contractOffloadedRequestPayloadthroughRolloutTokenCaptureandTQTokenSinkCallRecordonly after that acknowledgementRolloutReceiptat rollout end, then fetch, verify, and linearize staged records into the canonical GRPO sample during finalizationToken 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 indocs/design-docs/token-capture-ledger.md.Dependencies
amahishi/partial-rollout-telemetry-v3NVIDIA-NeMo/Gym@c3d7cf4Testing
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
pynvmlruntime dependency; Linux CI provides the authoritative RL test coverage.