feat(token-capture): carry media through Megatron capture rollouts

[lauradang]

What does this PR do ?

Stacked on NVIDIA-NeMo#4129 (laurad/minf-rollout-checkpointing), which is this PR's base branch.

Token capture on the Megatron backend was text-only. Receipt rows carried no pixel_values, so a VLM capture run trained image-blind, and every multi-turn VLM rollout failed at admission because the preparer spliced the expanded previous turn into the compact new render and Megatron expanded it twice.

• Compact chains. STAGING_FIELDS gains compact_token_ids_delta / compact_len. TQTokenSink pops Gym's compact delta out of extras into that column (like routed_experts); TQTokenSource.fetch_prefix_chains returns both token spaces, with compact falling back to expanded for text calls.
• Preparer. TQMegatronPromptPreparer splices with the compact chain, hands Gym the expanded chain to verify the engine prompt against, and records offload_params["ng_capture_minf"]["compact_prev_len"] for the stager.
• Media on the call row, as per-call deltas. Every chat request carries the whole conversation, so the engine hands the stager pixels for every image in the prompt. The preparer records how many media items the parent chain already staged (media_prev_count, next to compact_prev_len), the stager slices MInf's media_tensors at that boundary, derives the digest-covered geometry for Gym's adapter from the remainder, and stages those tensors as extra columns on the same staging key (MEDIA_STAGING_FIELDS), the way routed experts ride the row. Each row holds only the media new to that call, like the token columns.
• Finalizer. RolloutReassembler walks the terminal chain, checks each call's media columns against that call's staged geometry (media_columns_missing / media_mismatch), concatenates the deltas in chain order as it does for tokens, and publishes the engine's packed patches as pixel_values with imgs_sizes and num_frames. The Megatron-Bridge Omni model accepts the [total_patches, C*P*P] layout directly, so training projects exactly the pixels the policy generated against.
• Setup. Registers the media columns for multimodal runs; rejects capture with a multimodal policy on the vLLM backend (that path stages the pre-processor prompt and carries no media) and with grpo.deduplicate_multimodal_data.

docs/design-docs/token-capture-ledger.md gains a "Multimodal rollouts" section.

Worked example

A two-turn VLM rollout under capture, with toy values: the media placeholder is token 99, EOS is 2, one 4x4 image expands to three media tokens and to four packed patches of 12 features. Turn 1 sends text plus one image. Turn 2 sends the whole history plus a second image; the chat template re-tokenizes the turn-1 answer as 13 although the model generated 12, which is what the splice corrects. Each box is a process; the label says whose code runs in it.

sequenceDiagram
    box Gym server (Gym code)
        participant G as Agent + capture middleware<br/>+ MegatronWorkerCaptureHandler
    end
    box Megatron HTTP endpoint, in RL's Megatron worker (Megatron code)
        participant EP as chat_completions.py
    end
    box Engine TP0/PP0 (Megatron loop, RL+Gym hooks)
        participant P as TQMegatronPromptPreparer
        participant S as TQMegatronTokenStager<br/>+ Gym adapter
    end
    box Engine, all ranks (Megatron)
        participant E as add_request / _build_vlm_request
    end
    participant TQ as TransferQueue
    participant L as Gym ledger
    box RL actors (RL code)
        participant F as RolloutReassembler
        participant T as Trainer
    end

    Note over G,L: Turn 1 - text + image1
    G->>EP: chat req, offload_params.ng_capture = {c1, text, prev_len 0}
    EP->>P: compact [80,99,81] + image bytes (via DP coordinator)
    P->>E: text mode, pass through, broadcast
    E->>E: decode img1 to imgs[1,4,12], expand to [80,99,99,99,81], gen [12,2]
    E->>S: payload {expanded, gen, compact (new), media_tensors (new)}
    S->>TQ: r0/c1 - tokens, compact col [80,99,81,12,2], extras {imgs_sizes}, media cols
    S->>EP: ng_commit_coords {r0/c1, cum_len 7} via engine reply
    EP->>G: "a cat" + coords
    G->>L: CallRecord c1

    Note over G,L: Turn 2 - history + image2
    G->>L: resolve_parent finds c1
    G->>EP: ng_capture = {c2, token_in, prev_len 7, staging_chain [r0/c1]}
    EP->>P: compact [80,99,81,13,2,20,99,21], template_prefix [80,99,81,13,2], eos 2
    P->>TQ: fetch_prefix_chains([r0/c1])
    TQ->>P: expanded [80,99,99,99,81,12,2] / compact [80,99,81,12,2]
    P->>E: splice compact to [80,99,81,12,2,20,99,21], required_prefix = expanded, compact_prev_len = 5
    E->>E: 2 placeholders = 2 images OK, expand to [80,99,99,99,81,12,2,20,99,99,99,21], gen [30,2]
    E->>S: payload
    S->>S: Gym checks prompt[:7] == expanded chain, delta [20,99,99,99,21,30,2], compact [20,99,21,30,2]
    S->>TQ: r0/c2 tokens + compact + extras + media cols (image 2 only, media_prev_count 1)
    S->>EP: coords {r0/c2, prev_len 7, cum_len 14}
    EP->>G: answer + coords
    G->>L: CallRecord c2 (parent c1)

    Note over F,T: Rollout end
    F->>L: manifest to receipt {terminal c2, [c1, c2]}
    F->>TQ: fetch rows for 14 expanded tokens, fetch_media for r0/c1 and r0/c2, concat to imgs[1,8,12]
    F->>F: imgs_sizes cols == extras OK, pixel_values[8,12]
    F->>TQ: publish canonical row, clear r0/c1 and r0/c2
    T->>TQ: read input_ids + pixel_values, six 99s match six embeddings

Loading

Where the old code failed: at turn 2 the preparer spliced the expanded chain, so the engine saw four 99s for two images and rejected the request on the placeholder count. Splicing in compact space fixes that, and Gym's existing prefix check on the expanded prompt (prompt[:7]) verifies that re-expanding image 1 reproduced the same tokens.

Dependencies

• Megatron-LM: tdene/Megatron-LM PR on top of MInf: Add hooks to preprocess and stage NVIDIA/Megatron-LM#7015 (compact_prompt_token_ids, media_tensors on OffloadedRequestPayload), to be pinned through Megatron-Bridge.
• Gym: lauradang/Gym PR on top of feat(token-capture): add Megatron worker capture backend NVIDIA-NeMo/Gym#2823 (staging/media.py, Megatron adapter extras). Submodule and uv.lock re-pin pending, as with feat(rollout): integrate MInf ledger capture with checkpointing NVIDIA-NeMo/RL#4129.

Testing

• tests/unit/data_plane/test_tq_token_sink.py and test_rollout_reassembler.py against the in-process TransferQueue simple backend: 63 passed
• tests/unit/single_controller/test_setup.py -k "multimodal or megatron_token_capture": 7 passed

End-to-end Nemotron Omni video run

Completed a 42-step, 8-node × 4-GPU VSTAT GRPO run in 19m14s with this PR stack:

• RL: 79fa6856b5e40a456c8d932f434a490c9682e296
• Gym: fafecf7c71ef63d027bb2197895e93c9effced4a
• Megatron-LM: 0aa6cce27f75c2a4cf733b8eeb0e8180dc451707
• Megatron-Bridge: 8e4ccc34b67a79bca73cb76242479f9090856ff4

Run: PR stack (jtg9a7rp) · cye baseline (o4ed2p8n) · W&B overlay

The reward curves track closely: Pearson correlation 0.9169, mean absolute error 0.0486; mean reward over all 42 steps was 0.1692 for this PR stack vs. 0.1592 for the baseline. Steps 23–42 averaged 0.2229 vs. 0.2281.

The chart uses nemo_rl/step on the x-axis (not W&B's synthetic STEP). Duplicate history entries from continuation were reduced to the final value for each nemo_rl/step.

Important effective configuration:

policy:
  model_name: nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16
  is_vlm: true
  megatron_cfg:
    freeze_vision_model: false
    freeze_vision_projection: false
  generation:
    backend: megatron
    mcore_generation_config:
      megatron_inference_wrapper: megatron.core.inference.model_inference_wrappers.multimodal.nemotron_omni_inference_wrapper.NemotronOmniInferenceWrapper
      parsers: [nemotron-v3-reasoning, qwen3-coder-tool]
      video_num_frames: 16
      video_temporal_patch_size: 2
      video_target_num_patches: 1024

data:
  train:
    data_path: /opt/nemo-rl/workspace/datasets/vstat-8n4g/train-gym.jsonl
  default:
    dataset_name: NemoGymDataset
    processor: nemo_gym_data_processor
    num_frames: 16
    video_sampling_style: nemotron_vl
    video_temporal_patch_size: 2
    video_target_num_patches: 1024
    video_maintain_aspect_ratio: true

grpo:
  deduplicate_multimodal_data: false

data_plane:
  enabled: true
  impl: transfer_queue
  backend: simple

token_capture:
  enabled: true

checkpointing:
  enabled: true
  save_period: 10
  save_data_plane: true

The dataset rows contain real local video paths. The step-40 native TQ checkpoint contains pixel_values, imgs_sizes, and num_frames (plus their row-shape metadata) in replay_buffer_metadata.pt, and the same media columns in the TQ storage-unit pickle. A continuation read those columns from TQ and reached Nemotron Omni's _encode_images(images, ..., num_frames) path during policy logprob computation.

🤖 Generated with Claude Code

[lauradang]

Superseded by NVIDIA-NeMo#4332, which targets the upstream pranav/vlm_capture_support (NVIDIA-NeMo#4191) directly. Same head branch; it now merges the rebased NVIDIA-NeMo#4191 head 98e6a8cea.