fix(vlm): match training-side image tiling to the rollout's placeholder runs

[pulkitkumar95]

What does this PR do?

Fixes a train/generation image-tiling mismatch that crashes async GRPO VLM training mid-run:

ValueError: Expanded-sequence media alignment failed: found 32160 valid
placeholders for 40800 projected features.

The bug

A VLM sample needs the number of vision features in the training tensors to equal the number of image-placeholder tokens in the sequence — the model forward merges features into placeholder positions one-to-one.

Those two counts are computed by different components:

• Generation (vLLM) sizes image tiles per request, shrinking them as the prompt approaches max_model_len. The rollout's token ids reflect that choice.
• Training (attach path) re-processes the original images through the HF processor, which only knows its static config budget and cannot see what the engine chose for this request.

On most rows both sides agree, so the bug is invisible. On budget-bound rows (large frames + long text — e.g. 16 frames at 2400×1080 in a near-32k prompt) the engine shrinks tiles (2010 tokens/image) while the training side produces its static count (2550/image) → 40800 features vs 32160 placeholders → hard crash on every rank, at whatever step first samples such a row.

The fix

Use the rollout itself as ground truth instead of trying to mirror the engine's budget arithmetic (which lives in engine internals and can drift):

1. multimodal_utils — parse per-image <img><image>*N</img> placeholder runs out of the rollout token ids; the run lengths are exactly the feature counts the model will demand. Verify the processor's native output against them; on mismatch, re-process each image pinned to the exact count (the image processor's max_model_len clamp is the budget lever, with a deterministic exact-grid resize fallback). If parity cannot be established, raise a clear error instead of attaching misaligned media.
2. nemo_gym actor — pass each turn's placeholder runs into the attach. With deduplicate_multimodal_data: true, keep the media omission only when the statically-budgeted pre-attached tensors provably match the rollout's first-turn runs (predicted via the processor's own grid math); otherwise attach rollout-matched tensors actor-side.
3. rollouts — the driver-side static reattach never overwrites media the actor already attached. Video rows are unaffected: their tensors attach at datum time and their turns carry no extracted images, so the parity logic never engages for them.

Non-binding rows (the common case) take a fast path — the batched processor call is verified against the runs and used as-is, so behavior and cost are unchanged.

Validation

• Offline against the checkpoint processor: exact-count repair across a 7-size × 7-budget sweep (processor state restored after each pinned call); the crash row (16×2400×1080 pinned to 2010/image) yields exactly 32160 feature tokens; non-binding rows byte-match the previous fast path; the dedup-omission predictor matches the native processor on uniform and mixed-size batches.
• In production: the failing 32-node run resumed from checkpoint, passed the previously fatal batch, and trained to completion (50 steps) with no alignment errors and healthy reward progression.

Notes for reviewers

• The placeholder-run grammar (<img><image>*N</img>) and the num_tokens output key follow the NemotronH Omni processor family; processors that don't expose these degrade gracefully (parity checks are skipped, current behavior preserved).
• Budget-bound dedup rows lose media dedup for that row (the actor ships rollout-matched tensors); a follow-up could restore sharing with a repair-at-reattach cache if such rows are common in a workload.

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

[aroshanghias-nvd]

Curious, why not just increase the context length temporarily? The mismatch should be properly fixed though.

[aroshanghias-nvd]

It seems the repair infers the learner-side resize from the number of image placeholder tokens. N only determines the feature count; it does not identify the exact 2-D grid or pixel transform selected by vLLM. Could vLLM return its exact per-image preprocessing decision—target grid/H×W and transform identity—and have a Nemotron-specific adapter replay it for Megatron?

The generic multimodal layer should only transport and validate this metadata. Nemotron placeholder parsing, grid calculations, and replay should live in the existing Nemotron helper or behind a ProcessorInterface/VLM adapter, rather than in multimodal_utils.py.

[aroshanghias-nvd]

[P1] Do not treat key presence as proof that Gym media is authoritative

A pre-existing target value can be a placeholder or stale payload, not rollout-matched media. test_reattach_original_multimodal_payloads_is_media_only_and_turn_aligned supplies "remote placeholder" under pixel_values and expects the original PackedTensor to replace it; this branch preserves the string, so the unchanged base test fails.

Please use explicit provenance or the omission marker to identify media produced from the rollout decision, and retain the existing reattachment behavior otherwise.

[pulkitkumar95]

Thanks for the detailed review. All points are addressed in 59d0db8. Summary below.

1. Reattach guard: key presence is not provenance (inline comment)

You are right, the previous guard treated any pre-existing value under a media key as authoritative, which broke the documented behavior of
test_reattach_original_multimodal_payloads_is_media_only_and_turn_aligned (a placeholder value must be replaced by the real tensor).

This is now done with an explicit provenance marker instead:

• The attach sets ROLLOUT_MATCHED_MEDIA_KEY on a user turn only when it actually re-processed that turn's images to match the rollout's placeholder runs. Turns where the native processor output
  already matched get no marker.
• The driver-side static reattach skips (and consumes) only marked turns. Everything unmarked, including placeholder or stale payloads, is replaced exactly as before.
• Because only repaired turns are marked, unrepaired turns still get the shared tensor set restored across a prompt group's repeated rows, so the existing memory dedup behavior is unchanged.
• Added a regression test: a marked turn keeps its rollout-matched tensors and the marker is consumed, while unmarked representations are still restored from the static source. The previously failing
  test passes again.

2. Code placement: Nemotron specifics out of multimodal_utils

Agreed. The placeholder-run parsing, the static budget prediction, the exact-count re-processing and the grid math now live in the existing Nemotron helper module,
nemo_rl/environments/nemotron_utils.py. multimodal_utils.py keeps only the generic contract on the attach (an optional expected-count parameter with verification and a loud failure) and delegates
the repair through a local import. Processors that do not expose the placeholder grammar are detected up front and skip the parity logic entirely, so non-Nemotron models keep the current behavior.

3. Why not just increase the context length

Raising max_model_len only moves the boundary, it does not remove it. vLLM shrinks image tiles whenever the per-request budget binds, at any limit. The crash row's 16 high-res frames want about 40k
image tokens under the static budget, which is beyond any practical context length, and a larger context also costs KV cache memory. The mismatch is structural: two components compute tile counts
independently, so the fix makes them agree instead of trying to avoid the case where they disagree.

4. Having vLLM return its exact preprocessing decisions

Agreed on the direction, and this change is compatible with it. The placeholder runs in the returned prompt token ids already are the engine's exact decision, one token per projected feature, so we
can read the decision back today without any change to the pinned vLLM fork's API. If the engine later exposes structured preprocessing metadata (target grid, height and width, transform identity),
the replay entry point in nemotron_utils is exactly where that would plug in, and the token-run parsing would become a fallback.

Validation after the changes

• All reattach tests pass, including the one cited in the review and the new marker regression test.
• Re-ran the offline parity checks against the checkpoint processor: the crash row (16 frames at 2400x1080, runs of 2010) repairs to exactly 32160 feature tokens with the marker set, the non-binding
  fast path is unchanged with no marker, and the static budget predictor matches the native processor output.
• Ruff check and format are clean.