feat(model): enable Nemotron Omni audio, video, and packing

[cuichenx]

Note

Rebased onto main after #4885 merged as f6c9292c2fe1d144fa577562d681ebf1824f9ddb.

What changed

• enable canonical expanded-sequence Parakeet audio insertion
  • project [B, T, H] encoder output in the legacy-compatible layout
  • trim padded encoder rows using returned valid lengths
  • preserve sample order and cast FP32 encoder output to the Megatron projector dtype
  • use has_sound as the sole sound-capability and construction flag: HF import derives it from sound_config, CORD image-text SFT/PEFT set it false to omit the encoder/projector and roughly 1.13 GiB of replicated BF16 encoder weights per encoder-stage rank, and HF export writes sound_config=None so sound-free checkpoints stay sound-free
• enable canonical temporal multi-frame RADIO video insertion
  • remove the artificial single-frame guard
  • expand each tubelet placeholder to the exact post-pixel-shuffle feature count
• keep the expanded canonical model/collator as the default
  • retain the old LLaVA model/collator only through explicit compatibility providers and collapse_image_tokens=True
• move canonical packing completely into the collator
  • build final THD tokens, positions, labels, loss masks, true/padded cu-seqlens, max lengths, and total_tokens for text/image/video/audio rows
  • emit an explicit physical alignment-padding mask for media validation and future MoE routing
  • remove the canonical model-owned packing helper and set both model_owns_packing capabilities to false
• keep only post-media CP/SP localization in the model
  • insert media into the full, already packed stream without changing length
  • apply one MCore/Transformer Engine CP index to embeddings, IDs, positions, labels, loss masks, and padding masks
  • preserve global THD metadata on every PP stage
  • return the CP-local loss mask to the training step

Why

The canonical model uses one projected media feature per placeholder, so media insertion is length preserving. Packing before forward is therefore exact and avoids a second packing implementation inside the model. CP sharding remains after media insertion because all token-aligned tensors must use the same rank-local ordering.

For packed batches, padding_mask marks only synthetic CP/SP alignment gaps (True means padding). It is separate from loss_mask, which controls token supervision, and packed-sequence metadata, which controls attention and Mamba sequence boundaries. The collator creates this mask and the model keeps it aligned through CP/SP, but this PR deliberately does not forward it into the MCore language model yet: the pinned MCore expert-bias router crashes when it combines the flattened token mask with per-expert routing probabilities (MCore #6111). Until that upstream path is fixed, alignment gaps remain excluded from token loss but are knowingly counted by MoE router auxiliary losses and expert-bias statistics.

The #4885 canonical provider also retained sound/video modules but hard-failed their insertion. This stack enables both modalities while keeping the legacy collapse/expand contract explicit.

Why num_image_tiles was removed from canonical inference scripts

PR #4776 added num_image_tiles to the historical LLaVA inference path because that model received compact image placeholders and expanded each placeholder inside forward(). Every PP stage needed the replacement counts to reproduce the merged sequence length, and generation needed that length to avoid sampling from a padded tail.

The canonical model in this PR has the opposite contract: inference computes the same per-tile projected-feature counts with inference_num_image_tiles, aggregates them back to each source-media wrapper, and expands input_ids to one placeholder per projected feature before model forward. Media insertion is therefore one-for-one and length preserving, the model validates placeholder/feature equality, PP exchanges the actual variable sequence shape, and generation samples at input_ids.shape[1] - 1. Forwarding num_image_tiles to the canonical model would be unused legacy model-owned expansion metadata. The explicit deprecated NemotronOmniLlavaModel compatibility path still retains and consumes it.

This preserves the #4776 correctness invariant while moving ownership from model-forward expansion to canonical preprocessing. Dedicated regressions compare legacy merged widths against canonical pre-expanded widths for both dynamic multi-tile images and temporal tubelets.

Validation

Commands ran in the NeMo 26.06/26.08 development containers with real CUDA where applicable.

• focused packing/collator/Energon/step/model CPU matrix: 97 passed
• real single-GPU tiny model matrix: 22 passed
  • collator-owned packed RADIO image forward
  • temporal multi-frame video forward
  • packed Mamba boundary-reset parity
  • Parakeet subsampling/trim/count smoke
• real two-rank NCCL CP test: passed on both ranks
  • verified MCore/Transformer Engine per-document zigzag indices
  • verified identical sharding for every token-aligned tensor
  • verified global THD metadata remains unchanged
• Direct-HF/Energon canonical parity for mixed image, video, and audio packing: passed
• official Valor sample -uGHAvfqs2I, QA 9539, real clip + extracted audio
  • patched 10-frame/audio canonical forward: 1,280 video placeholders, 126 sound placeholders, finite loss 11.72488499
  • three-step BF16 AdamW legacy/canonical comparison: maximum loss delta 0.0
• real target 30B Bridge Trainer experiment on the equivalent no-mask path
  • two nodes / 16 H100s, TP2 / PP2 / EP4 / CP1 / SP, sequence length 4096, MBS2 / GBS8
  • expert bias remained enabled; collator-owned THD packing and physical alignment gaps remained unchanged
  • completed five optimizer iterations with finite losses and grad norms, zero skipped iterations, zero NaNs, and no runtime errors
• final workaround regression and adjacent packing/model/data matrix: 140 passed
• final real two-rank NCCL CP rerun: passed on both ranks
• uv run pre-commit run --all-files: passed
• git diff --check: passed

Fresh validation after rebasing onto main

On source/test commit a3c8d076f3fd6992715c2c9e30f21e33ef986bba (card-only head 7888d7aba2d04220d6bfc35fad7db4209e61d6ac):

• focused Nemotron Omni packing, data, model, conversion, recipe, and training-step matrix: 144 passed
• model verification card schema validation: passed
• uv run pre-commit run --all-files and git diff --check: passed
• git range-diff: all 13 patches preserved

The two-rank NCCL, official Valor replay, and three-step comparison remain historical evidence from source/test commit 7dff2df3e773adf346cd6018596909760679efcd; they were not rerun for this rebase.

The verification card is at examples/model_verification_cards/nemotron-3-nano-omni-30b-a3b-reasoning/card.yaml. It records this focused evidence without promoting any full end-to-end inventory item; the target-scale 30B result above remains historical evidence from the pre-restack series.

Follow-ups

• Compact variable-length packs do not yet restore original per-row batch structure for seq_aux_loss. This stack preserves attention/Mamba boundaries; exact per-row seq-aux semantics will use the boundary-aware Megatron-Core unflattening extension requested in MCore #6108.
• Restore MCore router padding_mask propagation after MCore #6111 is fixed and the Bridge MCore pin includes that fix.
• Packed MTP and target-scale Super/MTP conversion/training/checkpoint validation remain outstanding.
• NeMo-RL's previous model_owns_packing delegation path must move to data/collator-owned packing before downstream packed-RL validation.

Follow-up hardening

• audited the default end to end: AutoBridge, public exports, Direct-HF, Energon, recipes, training, and inference all select the canonical model and expanded collator
• deprecated the explicit LLaVA model/provider/bridge/collator compatibility path with documentation and FutureWarning notices
• fixed canonical PP inference to exchange variable sequence shapes instead of allocating fixed-width receive buffers
• fixed comma-separated multi-image inference for heterogeneous dynamic-resolution tile tensors
• independently reviewed the final diff and addressed both reported gaps

Additional validation on head 32905fdba: 152 focused tests passed in the NeMo 26.06 CUDA container, including an actual tiny packed multimodal forward/backward/AdamW parameter update. The final target-scale experiment used an equivalent disposable omission on prior head f69c73acf; commit bf6fb39f4 makes that omission explicit with a regression and an upstream TODO.

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

[claude]

Light review — PR #5007 (collator-owned sequence packing for Nemotron Omni)

LGTM. This is a well-scoped refactor moving THD packing ownership from the model to the collator, and enabling audio + multi-frame video modalities plus CP/SP sharding after length-preserving media insertion. Test coverage is thorough (expanded collate contract, sound encoder, dense + packed CP sharding, HF/Energon parity, 2-rank distributed CP, in-batch padding-mask emission).

Verified, no blocking issues:

• Tuple-return contract is consistent: forward() returns (output, loss_mask) only when a CP shard is applied (return_sliced_loss_mask), and nemotron_omni_step.forward_step guards with isinstance(model_output, tuple) before unpacking.
• Sound sentinel detection (sound_clips.shape == [1, 1] then value check) is guarded by numel() > 0 and handles the empty-sound placeholder correctly.
• padding_mask in in_batch.py is initialized all-True (physical alignment gaps) and set False at real token offsets; correctly added to reserved_keys.
• No stale references left behind: inference_merged_sequence_length still exists in nemotron_omni_utils.py and its remaining callers are untouched — only the valor32k_avqa_inference.py example intentionally switched to input_ids.shape[1].

Minor (non-blocking):

• New Python files should carry the 2026 NVIDIA copyright header per repo convention; please double-check any newly added source/test files use the current year.

Suggested test cases

• No perf tests impacted. (No scripts/performance/configs/ files are touched.)
• Relevant unit/functional coverage added by this PR:
  • test_nemotron_omni_registry_selects_canonical_expanded_contract
  • test_nemotron_omni_expanded_collate_emits_one_placeholder_per_temporal_feature
  • test_energon_temporal_video_defaults_to_expanded_contract
  • test_hf_and_energon_packing_are_identical_for_image_video_audio
  • test_packing_marks_only_physical_alignment_gaps_as_padding
  • test_collator_owned_thd_tensors_use_one_real_cp_partition_index
  • test_audio_forward_replaces_expanded_placeholders_without_changing_length
  • test_sound_encoder_drops_padded_rows_and_preserves_sample_order
  • test_dense_expanded_sequence_is_cp_sharded_after_media_insertion
  • test_collator_owned_packing_is_preserved_while_model_applies_cp_shard
  • test_real_radio_multiframe_video_forward

[cuichenx]

/ok to test 5aa4034

[cuichenx]

/ok to test 558e6b0