[data] fix: isolate and unify Nemotron Omni collation

[cuichenx]

Summary

This PR has five goals:

1. Fix cross-sample contamination and multimodal metadata handling. adjust_image_tokens now adjusts each row and image independently using the correct num_patches, without flattening the batch or misaligning labels and loss masks. Nemotron Omni also distinguishes num_frames temporal grouping from num_image_tiles language-model replacement counts and preserves the tokenizer's real pad ID.
2. Consolidate the Direct-HF and Energon data paths. Both now use the same model-aware collator for tokenization, masking, image/video/audio processing, padding, and packing; the Energon task encoder only adapts source samples. Omni processors remain on this model-owned collator instead of selecting the generic text collator.
3. Make the public model contract explicit. Nemotron Omni supports only dynamic resolution, uses squared-ReLU projection and square-then-crop C-RADIO position embeddings, validates image/sound token and configuration IDs, and accepts only its model-owned pixel_values visual input.
4. Support correct multimodal online in-batch packing. Image-text and video/audio/text batches use THD boundaries computed after modality-token expansion, with row-major modality ownership, aggregate model-length guards, CP/SP alignment, fixed-width pipeline tails, and masked physical gaps that prevent cross-sample attention or supervision. Generic text packing retains its established emergent-width behavior.
5. Harden pipeline-parallel training and inference. Labels, expanded loss masks, exact vision replacement counts, and packed metadata reach the stages that consume them; model-returned loss masks are honored, receive buffers use the correct fixed width, and generation samples from the last real merged token rather than a padded position.

Full-data HF/Energon parity

Slurm job 13975195 trained on the complete MedPix splits (17,420 train / 3,080 validation) for 2,000 packed optimizer steps per path on the same 8xH100 node and source revision. The experiment matched sample and DP order to isolate data-pipeline behavior.

• All 312 traced Direct-HF/Energon runtime records were tensor-identical across tokens, labels, loss masks, positions, visual inputs, tile counts, and THD metadata, including the epoch boundary and train/validation transitions.
• LM-loss correlation was 0.99977; the worst validation-loss gap was 0.065%. Sequence load-balancing correlation was 0.99383, with zero skipped or NaN iterations.

Compatibility and scope

• Recipe/CLI usage is unchanged: --dataset vlm-hf, --dataset vlm-energon, and dataset.path=... remain valid.
• The dynamic_resolution provider field is retained for serialized/config compatibility, but Nemotron Omni validates that it is True; legacy fixed-tile processors fail explicitly.
• Omni's visual_keys setting is retained for configuration compatibility and must be ("pixel_values",). Previously ignored alternatives now fail explicitly.
• Non-temporal images retain public variable-resolution processing. Temporal mode uses an antialiased 512×512 compatibility canvas for every visual item because pinned MCore cannot pixel-shuffle ragged non-square tubelets.
• Low-level NemotronOmniTaskSample constructors must now provide the shared pre-collation example; batch visual_tensors and tokens compatibility adapters remain available.
• No dependency, CI workflow, or third-party submodule change.

[copy-pr-bot]

Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually.

Contributors can view more details about this message here.

[cuichenx]

more refactor coming to unify the hf and energon paths for nemotron omni. i'll absorb this fix to that refactor

[cuichenx]

/claude review

[claude]

Light review - Nemotron-3 Omni collation/inference refactor

Overall this is a solid, well-tested refactor. It removes the configurable dynamic_resolution flag (pinning Nemotron Omni to dynamic_resolution=True), consolidates the Energon and Direct-HF paths onto a single shared nemotron_omni_collate_fn, adds post-vision-merge-aware in-batch packing (with total_tokens plumbed through so MCore resets Mamba/SSM state at packed-row boundaries), and reworks inference next-token selection to skip PP padding. Test coverage is thorough, including an HF-vs-Energon parity test for image+video+audio packing and a randomized reference test for adjust_image_tokens.

No blocking correctness issues found. A few observations:

• Config-key contract change (heads-up, not blocking): dynamic_resolution is now Literal[True] on NemotronOmniModelProvider and finalize()/provide() raise if it is ever False. Any existing config or saved-checkpoint override that sets dynamic_resolution=False will now error rather than silently run the old static path. Intentional and validated by test_nemotron_omni_provider_rejects_static_resolution, but per CONTRIBUTING this is effectively a breaking change for the Omni recipe surface; consider the breaking-change label.
• Doc-accuracy nit (inline): the inference _patchify_frame docstrings claim they mirror the shared collator 'exactly', but the shared path uses antialiased bicubic interpolate on float32 while inference uses PIL.resize + ToTensor; results are close but not bit-identical.
• _add_audio_inputs hardcodes 16 kHz (waveform slice using 16000) while _audio_waveform takes a target_sampling_rate param. Not a bug today (target is always 16000), but the literal and the parameter can drift; minor.

Boundaries respected: no 3rdparty/Megatron-LM/ edits, no new dependencies, copyright headers present on new source files.

Suggested test cases

• test_nemotron_omni_finetune_recipe (tests/functional_tests/test_groups/recipes/test_nemotron_omni_recipes_finetune.py - exercises the recipe after dynamic_resolution removal)
• test_nemotron_omni_conversion_roundtrip (tests/functional_tests/test_groups/models/nemotron_omni/test_nemotron_omni_conversion.py)

No perf tests impacted (no scripts/performance/configs/ files touched).

[claude]

Minor doc-accuracy: this docstring now says it "Mirrors the shared nemotron_omni_collate_fn patchification exactly," but the two implementations diverge. The shared collator (_patchify_frame in collate_fn.py) resizes via torch.nn.functional.interpolate(mode="bicubic", align_corners=False, antialias=True) on a float32 tensor, while this inference helper resizes with PIL.Image.resize(...) + transforms.ToTensor(). Those produce subtly different pixels (PIL resize is not antialiased bicubic and rounds on uint8), so tensors are close but not identical. Consider softening "exactly" to avoid implying bit-parity. Same wording appears in valor32k_avqa_inference.py.