feat(vlm): support combined video and tracking GRPO

[ehosseiniasl]

TL;DR

• Issue: CAPRL video-only and SA-V tracking-only training worked independently, but the combined loader flattened their different media representations and lost each row's media mode and processor-owned placeholder counts. In combined training, vLLM then emitted 8192 placeholders for a CAPRL video while Megatron reconstructed 8064 vision features—128 excess placeholders per video—and failed in policy.get_logprobs before an optimizer step.
• Fix: Preserve the row-specific media contract from preprocessing through Gym, rollout/replay, and policy scoring, and apply the checkpoint's nested vision settings consistently in vLLM and Megatron. R3 payloads are additionally kept as lazy Ray references and materialized only after microbatch slicing, avoiding the full-batch host-memory failure.
• Training semantics are unchanged: this PR does not alter GRPO, rewards, profile-band parameters, or the TMPE masking threshold.

Architecture

Scope clarification

• This is a combined CAPRL + SA-V integration issue, not a CAPRL video-only regression. The validated CAPRL video-only training path did not show the 8192-vs-8064 mismatch.
• CAPRL represents 64 tagged frames as one logical video and temporally patches frame pairs. SA-V represents frame-labeled images as independent media regions, applies spatial tiling per image, and uses a structured tracking verifier.
• A mixed batch therefore cannot safely reconstruct both row types from a generic flattened image list. It must retain each row's media mode, media-region boundaries, and exact processor-owned token counts.

Concrete examples

• CAPRL video QA (mcqa_simple_agent)
  • Media: 64 input_image items tagged with _is_video_frame: true and one _video_source.
  • Example prompt: At what exact timestamp does the third and final sparring sequence begin? (A) [00:00] (B) [00:04] (C) [00:09] (D) [00:13] Please put the final answer within \\boxed{...}.
  • Expected output: one multiple-choice answer in \\boxed{...}.
• SA-V tracking (sav_pt_tracks_agent)
  • Media: independent frame images interleaved with Frame N: labels; tracked points may first appear in later frames.
  • Example prompt: Track point 1 at (445, 624), point 2 at (34, 585) in frame 4, and point 3 at (32, 691) in frame 4. For every point and requested frame, report its 0-1000-grid location and visibility.
  • Expected output: strict per-frame tracking JSON such as <answer>{"tracks":[{"point":1,"frame":1,"xy":[x,y],"visible":true}]}</answer>.

Why a recipe-only fix is insufficient

• The old combined handoff dropped the original processed-media representation before policy.get_logprobs.
• vLLM expanded the affected CAPRL video as 32 tubelets × 256 placeholders = 8192 placeholders.
• Megatron's checkpoint processor reconstructed the same video through the aspect-preserving path as 32 × 252 = 8064 features.
• Batches containing 3, 4, or 5 CAPRL videos consequently failed with 384, 512, or 640 excess placeholders. By this point, YAML cannot recover the discarded row-level media mapping or exact region lengths.

Resolution

• Route each row by its actual representation: logical CAPRL video rows use the video processor, while frame-labeled SA-V rows use the image processor.
• Record each processed region's media token ID and exact contiguous placeholder-run length directly from checkpoint processor output.
• Carry that contract through NeMo Gym messages, trajectories, replay collation, and policy preprocessing.
• Reconcile SA-V image-region placeholder runs to their recorded lengths while requiring exact CAPRL video alignment; fail early instead of editing generated token sequences.
• Deep-merge the checkpoint's nested video settings into vLLM and Megatron: temporal patch size 2, target 1024 patches, and aspect-ratio preservation.
• Load the checkpoint's frozen vision weights with load_format: auto, and use the checkpoint tokenizer plus its default chat template on both paths.
• Pause/resume async vLLM generation around in-flight refits so requests do not cross a weight/cache transition.
• For R3, store routed-expert payloads as lazy per-message Ray references and materialize them only after policy microbatch slicing.
• Fully remove repeated-layer MTP metadata when mtp_num_layers=0, matching this checkpoint's actual non-MTP configuration.

Validation on latest super-v3.5-posttraining

Base: 3d4db49 (fix(vlm): recognize Super Omni placeholder processor (#3989)).

• 4-node sync: completed mixed rollout, reward processing, policy.get_logprobs, TMPE calculation (0 masks), and backward. It then exhausted the deliberately small 8-GPU policy pool while allocating the first Adam exp_avg_sq buffer; this is a topology capacity limit, not a media-alignment failure.
• 8-node sync: completed repeated full rollout → logprobs → backward → Adam → refit steps with zero masks. At step 7: reward 0.375, mean TMPE 1.0231, max TMPE 1.0335, and steady-state step time 51.0s. W&B
• 20-node async, R3 enabled, no profile-band: completed 11 steps with one masked sequence across 22,528 samples. Latest logged step: reward 0.4645, mean TMPE 1.0152, max TMPE 1.9197, masks 0, step time 545.1s. W&B
• 32-node async, profile-band enabled, R3 disabled: completed seven steps with five masked sequences across 14,336 samples; steps 1, 2, 4, 5, and 7 had zero masks. W&B
• TMPE spike diagnosis: retrospective per-sample dumps show the masks are sparse generated-text tails, not media placeholders. The first two were ordinary CAPRL word/subword tokens. The later three were coordinate digit tokens from three generations of one SA-V box-tracking prompt; the other 13 generations of the identical prompt/media prefix remained valid. Sequence-wide mean absolute logprob deltas stayed around 0.016–0.054, while one rare token dominated the exponential TMPE statistic. The async batches had trajectory age 1; the synchronous run remains mask-free.
• All 36 modified Python files parse cleanly, git diff --check passes, and all six parent commits include DCO sign-offs.

Scope

• The source combined-training commit 6bf9228 changes 63 files with 4,485 additions. This latest-Super backport changes 39 files with 1,882 additions and keeps only the end-to-end mixed-media, vLLM/refit, R3, non-MTP, tests, documentation, and portable recipe pieces required by the validated runs.
• Cluster-specific launch scripts and smoke recipes remain local and are not part of this PR.
• The SA-V verifier is tracked separately in NVIDIA-NeMo/Gym#3094; this PR updates the Gym submodule to that latest-main backport.

[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: fff7ebb (PR #3980 from ehsan/super35-latest-combined-video-minimal)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of super-v3.5-posttraining branch (fast-forward)

All submodule changes look good! ✨

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: 45979a8 (PR #3980 from ehsan/super35-latest-combined-video-minimal)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of super-v3.5-posttraining branch (fast-forward)

All submodule changes look good! ✨

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: 7d02a40 (PR #3980 from ehsan/super35-latest-combined-video-minimal)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of super-v3.5-posttraining branch (fast-forward)

All submodule changes look good! ✨

[github-actions]

✅ Submodule Fast-Forward Check Results

Check based on commit: 4eca9e4 (PR #3980 from ehsan/super35-latest-combined-video-minimal)

✅ Submodules that are properly updated:

Gym: ✅ PR branch is ahead of super-v3.5-posttraining branch (fast-forward)

All submodule changes look good! ✨