[macandro96]

What does this PR do ?

Adds multimodal support to vLLM prefix-level rollout recovery. A checkpoint can now preserve an in-flight image or video generation prefix, record a SHA-256 digest of each of its media tensors, verify the reprocessed request against those digests, and continue decoding. Prefix rows never store pixels; the canonical terminal row stages them once, as before.

This is a stacked PR on top of #4507; the reviewable change is the seven commits on amahishi/prefix-recovery-multimodal. The stack is ordered correctness first, optimization last: #4507 → this PR → #4612 (router replay across prefixes) → #4508 (batched prefix TQ I/O).

Design

• Fingerprints each media tensor once, when the call is admitted (off the event loop): SHA-256 over its name, dtype, shape, and bytes.
• Stores media_spans and those tensor digests in the digest-covered extras of the first durable prefix row for a model call, with no pixel attachments, so a checkpoint cut writes no media under its deadline.
• Keeps later prefix chunks token-only and references the first row through the existing ordered staging-key chain.
• On recovery, reads the prefix rows (spans and digests) without any tensor read, reprocesses the reissued multimodal request, and verifies its media items and tensor digests (verify_restored_media, which names the mismatch: items, tensor set, dtype/shape, or bytes).
• Stages the reprocessed tensors, now proven byte-identical, in the canonical terminal row.
• Preserves the existing buffer-swap, deadline, and terminal-completion race behavior from the parent stack.
• Fails closed when media metadata/tensors are missing or changed. A cut before media preprocessing has completed is reported as durable_failure, so that model call is redispatched instead of restoring an incomplete prefix.

Checkpoint and recovery flow

flowchart TD
    A["Gym / agent sends request plus original media reference"] --> B["vLLM resolves and preprocesses media"]
    B --> C["Decode while token capture buffers token deltas"]
    C --> D{"Checkpoint cuts this active model call?"}
    D -- "No" --> E["Stage canonical terminal row with tokens and media"]
    D -- "Yes" --> F["Atomically swap the active token buffer"]
    F --> G{"First durable prefix row for this call?"}
    G -- "Yes" --> H["Stage tokens, media_spans, and media tensor digests in TQ"]
    G -- "No" --> I["Stage a token-only prefix chunk in TQ"]
    H --> J["Publish the coordinated TQ plus Gym / SC checkpoint"]
    I --> J
    J --> K["Process restarts and loads the checkpoint"]
    K --> L["Gym reissues the call with original media and the staging-key chain"]
    L --> M["Generation worker fetches the prefix rows (tokens, spans, digests) from TQ"]
    M --> N["vLLM reprocesses the original media"]
    N --> O{"Prompt tokens, media metadata, and tensor digests match?"}
    O -- "No" --> P["Fail closed; redispatch the model call from a safe boundary"]
    O -- "Yes" --> Q["Prefill the durable token prefix with verified media"]
    Q --> R["Decode the remaining suffix"]
    R --> S["Stage the canonical terminal row with the reprocessed media tensors"]
    S --> T["Clear obsolete prefix keys only after terminal staging succeeds"]

Loading

Constraints and assumptions

Constraints

• Prefix-level recovery remains vLLM-only. Megatron inference is not enabled by this PR.
• Router replay with prefix-level recovery remains unsupported and retains its existing setup guard.
• The original image or video must still be resolvable when Gym reissues the request. Prefix rows hold only digests of the packed tensors, so this PR does not provide source-free media prefill.
• This does not checkpoint vLLM's KV cache or media-processor cache. Recovery reprocesses the media and prefills the durable token prefix.
• A cut that arrives before media preprocessing and capture setup complete cannot save a media prefix. It returns durable_failure, and the call is redispatched from the preceding safe boundary.
• Only the first durable prefix row for a model call may carry media metadata and tensor digests. Later rows are token-only; missing, malformed, or conflicting media metadata fails closed.
• Missing or corrupt prefix-row metadata, changed media, or any metadata or digest mismatch rejects prefix recovery (retained_media_changed) rather than continuing with uncertain state.
• The configured data-plane checkpoint backend must restore tensor payload rows, not controller metadata alone.

Assumptions

• Restore uses a compatible model, tokenizer, renderer, media processor, and capture schema; the parent checkpoint stack enforces the broader configuration identity.
• Reprocessing the same source media is deterministic enough to reproduce the exact captured metadata and tensor bytes. The implementation intentionally requires byte equality, checked by SHA-256 digest.
• The ordered staging-key chain is preserved, and its first media-bearing row remains live until a canonical terminal row has been staged successfully.
• The parent stack publishes Gym state, SC recovery metadata, and the corresponding TQ snapshot as one coordinated recovery point.
• Gym/agent/environment state recovery is provided by the parent Gym v2 checkpoint protocol. This PR only adds multimodal generation-prefix payload handling.

Issues

None.

Usage

The existing prefix-recovery configuration now also accepts supported vLLM multimodal workloads:

token_capture:
  enabled: true

rollout_recovery:
  target_level: prefix

The generation backend must be vLLM and media capture must be supported by the configured VLM processor.

Testing

Added or extended unit coverage for:

• digest-only first prefix rows and later token-only chunks, for one and two cuts;
• the image and video (num_frames) cut/restore round trip through a real TQ prefix row, with no tensor read on restore;
• the prefix digest written at the cut matching the one restore rebuilds (text, image with one cut, image with two cuts);
• each restore rejection: missing, malformed, or mismatched spans and digests; item, tensor-set, dtype/shape, and byte mismatches; extras after the first chunk; a reprocessed prompt that differs from the durable prefix;
• terminal completion racing with an active prefix cut, terminal staging failure, and obsolete-prefix cleanup failure;
• checkpointing before media preprocessing completes;
• a recovery request whose source media disappeared failing before any TQ read or begin_call;
• setup accepting multimodal prefix recovery.

The shared prefix-recovery parity checks now find a restored call by counting back from its row's end, so positions the trainer masked inside the call (top-k/top-p -inf, #4154) no longer break them; the Workplace parity test uses the same fix.

Unit tests pass locally on Linux at the head commit:

• --nemo-gym-only (test_captured_media.py, test_vllm_token_capture_hosting.py): 149 passed;
• tests/unit/environments/test_gym_prefix_recovery_parity.py: 26 passed;
• tests/unit/single_controller/test_setup.py: 173 passed.

Functional: tests/functional/nemotron_omni_clevr_vllm_prefix_recovery_single_controller_1n2g.sh cuts active image generations at a checkpoint, hard-kills, restores, and requires each restored call to verify its image, finish once, and train on its exact restored prefix. It is registered in L1_Functional_Tests_GB200_Vllm_Omni_Single_Controller.sh (full mode). It needs Gym's string_match to declare checkpoint_mode = "stateless" (NVIDIA-NeMo/Gym#3887), which the stack's Gym pin includes. It passed locally on H100 with 3 GPUs, TP2 training, and gpu_memory_utilization=0.9 passed as overrides (4 cuts across both image groups restored and trained, media_row_rate 1). The committed defaults (2 GPUs, TP1, 0.6) match the sibling Omni token-capture L1 test and have not yet run on GB200.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
• Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

• Parent PR: test(checkpoint): add Workplace prefix-recovery parity #4507
• Head commit: 13d42ba2eaf118982aa6364e45482e729f44ee51

[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.

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: 01ef7df (PR #4583 from amahishi/prefix-recovery-multimodal)

✅ Submodules that are properly updated:

Megatron-Bridge: ✅ PR branch is ahead of amahishi/ckpt-v2-prefix-batched-io branch (fast-forward)

All submodule changes look good! ✨

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: 01ef7df (PR #4583 from amahishi/prefix-recovery-multimodal)

✅ Submodules that are properly updated:

Megatron-Bridge: ✅ PR branch is ahead of amahishi/ckpt-v2-prefix-batched-io branch (fast-forward)

All submodule changes look good! ✨

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: 0bf5789 (PR #4583 from amahishi/prefix-recovery-multimodal)

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

[terrykong]

This PR extends vLLM prefix-level rollout recovery to multimodal (Omni image/video) calls. At a checkpoint cut, the first prefix row of a model call stores its media_spans and packed media tensors in TQ. On restore, the worker reprocesses the original media, checks it against what was saved, and continues decoding. The PR also removes the setup guards that blocked multimodal prefix recovery.

Thanks for this. Good parts:

• MEDIA_SPANS_FIELD is imported from Gym where it is read, with no copied string.
• fetch_generation_cut is built on the existing fetch_for_finalization and fetch_media helpers instead of new TQ reads.
• The fail-closed design is the right call, and the "Constraints and assumptions" section in the description makes the limits easy to check.

Size. This PR's own source change is small (about 190 lines). The comment on return restored would cut about 55 of them, and the comment on the functional script would move about 500 test lines to the PR that can actually run them. The stack is already over 20k lines, so keeping each step tight matters. Later PRs in the stack (#4612, #4508) rewrite parts of this PR's new fetch and staging code, so landing the hash change and the getattr cleanup here first avoids carrying them up the stack.

CI. Only the lint and metadata checks have run on this head (0bf5789). No unit or functional CI has run yet.

Question for the author (no action required): does the Omni chat template render the same text on every call? A media restore now re-renders the first-turn prompt and fails closed if it differs, so a template that is not stable would turn every media restore into a redispatch.

Visual explainer: https://terrykong.github.io/gh-pages-poc/terryk/pr-4583-text-vs-media.html
Visual explainer: https://terrykong.github.io/gh-pages-poc/terryk/pr-4583-size-and-pixel-copy.html

Generated by Claude Code

[macandro96]

/ok to test 746b304

[macandro96]

/ok to test 8574280

[macandro96]

/ok to test 13d42ba

[macandro96]

/ok to test a4a4644

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: e0f827b (PR #4583 from amahishi/prefix-recovery-multimodal)

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

[macandro96]

/ok to test e0f827b

[macandro96]

/ok to test 19051ca