feat(vlm): mixed image+video Gym manifests and video media alignment verification - #4005
Open
pulkitkumar95 wants to merge 5 commits into
Conversation
…er runs Async GRPO VLM training crashed mid-run with "Expanded-sequence media alignment failed: found 32160 valid placeholders for 40800 projected features". The generation engine sizes image tiles per request (shrinking them as prompts approach max_model_len); the training-side attach path re-processed the original images under the processor's static config budget, so budget-bound rows (e.g. 16 frames at 2400x1080 in a near-32k prompt) produced more projected vision features than the rollout's placeholder tokens, and the model forward raised on every rank. Derive the truth from the rollout itself instead of mirroring the engine's budget arithmetic: - multimodal_utils: parse per-image <img><image>*N</img> placeholder runs out of the rollout token ids; verify the processor's native output against them and, on mismatch, re-process each image pinned to the exact run length (budget pin via the image processor's max_model_len clamp, with a deterministic exact-grid resize fallback). Raise rather than attach misaligned media. - nemo_gym actor: pass each turn's placeholder runs into the attach; with deduplicate_multimodal_data on, keep the 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), else attach rollout-matched tensors actor-side. - 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). Verified offline against the checkpoint processor (exact-count repair across a 7-size x 7-budget sweep; the crash row now yields exactly matching feature counts; non-binding rows byte-match the previous fast path) and in production: the failing run resumed through the previously fatal batch and trained to completion with no alignment errors. Signed-off-by: Pulkit Kumar <pulkitk@nvidia.com>
…s relocated - Replace the key-presence reattach guard with an explicit provenance marker (ROLLOUT_MATCHED_MEDIA_KEY): the attach sets it only on turns it actually REPAIRED to the rollout's placeholder runs, and the driver-side static reattach skips (and consumes) exactly those turns. Unmarked values — including placeholder or stale payloads — are replaced as before, restoring the documented behavior of test_reattach_original_multimodal_payloads_is_media_only_and_turn_aligned, and unrepaired turns keep the shared-tensor restore across a prompt group's repeated rows. - Move the Nemotron-specific parity logic (placeholder-run parsing, static budget prediction, exact-count re-processing and grid math) out of multimodal_utils into the existing Nemotron helper module (nemo_rl/environments/nemotron_utils.py). The generic attach keeps only the verification contract and delegates the repair via a local import; processors without the placeholder grammar degrade gracefully as before. - Add 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. Signed-off-by: Pulkit Kumar <pulkitk@nvidia.com>
Scoped port from ehsan/super35-video-rpb-tmpe-fix (6bf9228) onto the merged super-v3.5-posttraining tip, limited to what the combined-manifest experiment requires: - data/processors.py (+ small data/* hunks): allow one NeMo-Gym manifest to mix static-video rows with still-image rows (previously a hard failure), and fix the dtype mismatch that broke mixed batches; per-row image_max_num_tiles spec support. - multimodal_utils: add the uses_fixed_tile_image_processor predicate the ported processor code depends on. - generation/vllm/config.py: fail fast on the unsupported legacy vllm_cfg.video_loader key (the config typo behind the original failure). - algorithms/utils.py: keep the processor chat template in sync with the tokenizer when a template override is applied. - Gym submodule: current head plus a cherry-pick of the SAV tracking verifier (resources_servers/sav_tracks) — no pointer rewind. Deliberately excluded (tracked separately): refit pause/resume (superseded upstream per branch owner), vLLM private-API monkeypatches and the 0.20 serving shim, Megatron-Bridge registry surgery and other old-container affordances, tiling-parity/rollout attach changes, load_format handling, reference recipes and personal launchers.
The full_generalist_prod export ships this processor class; the NemoGym multimodal path asserts on _PLACEHOLDER_STYLE_PROCESSOR_NAMES and job 3515316 failed at startup with exactly that assertion. Add it to the placeholder-style and Nemotron video processor allowlists (same two-line hunks as the video-stack branch).
…attach Video rows attach statically-preprocessed tensors (one <img><image>*k</img> run per tubelet) to vLLM-authored rollout tokens. When vLLM's own expansion disagrees (different per-frame grid or tubelet count), training crashed deep in Megatron with an unattributable 'Expanded-sequence media alignment failed'. Compare the placeholder-run structure of the rollout tokens against the static source tokens before the attach and raise a diagnostic that names the mismatch: run counts, per-run lengths, first mismatching run, num_frames and frame size. Plumbed the tokenizer through attach_static_multimodal_payload's two live call paths (RolloutWorker.run_rollout, run_async_nemo_gym_rollout); verification is a no-op for image-only turns or when no tokenizer is passed.
Author
|
/ok to test 1ee3233 |
svcnvidia-nemo-ci
added
the
waiting-on-maintainers
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What does this PR do ?
Adds mixed image+video NeMo-Gym manifest support and a video media alignment verification that turns an unattributable Megatron crash into a named diagnostic.
Issues
image_max_num_tiles, fixed-tile processor predicate, fail-fast for the legacyvllm_cfg.video_loaderkey, processor/tokenizer chat-template sync).NEMOTRON_VIDEO_PROCESSOR_NAMESwas missingNemotronH_Omni_Reasoning_V3Processor(fix(vlm): recognize Super Omni placeholder processor #3989 covered only the placeholder-style list), so Super Omni video rows fell through to generic preprocessing.Usage
Before your PR is "Ready for review"
Pre checks:
Additional Information
Validated in 32-node async GRPO on a combined SA-V tracking + CapRL video blend (Nemotron Super Omni, 64-frame videos): the verification caught an aspect-ratio grid mismatch on the first batch, and with matched config the runs train 20+ steps with healthy rewards.