Skip to content

feat(model): add canonical Nemotron Omni model - #4885

Merged
cuichenx merged 8 commits into
NVIDIA-NeMo:mainfrom
aroshanghias-nvd:aroshanghias/nemotron-omni-main-migration
Jul 30, 2026
Merged

feat(model): add canonical Nemotron Omni model#4885
cuichenx merged 8 commits into
NVIDIA-NeMo:mainfrom
aroshanghias-nvd:aroshanghias/nemotron-omni-main-migration

Conversation

@aroshanghias-nvd

@aroshanghias-nvd aroshanghias-nvd commented Jul 15, 2026

Copy link
Copy Markdown
Contributor

What does this PR do?

Add a dedicated NemotronOmniModel and conversion contract for processor-expanded multimodal sequences without using the historical Nemotron Omni LLaVAModel collapse/expand path.

Fresh review requested: The branch history has been replaced with a main-based implementation of the reviewed data-boundary packing design. The earlier model-owned-packing implementation is obsolete; please review this series from the beginning.

The ownership boundary is now:

  • The data pipeline packs processor-expanded token sequences and constructs PackedSeqParams. MBridge's canonical collator provides this behavior for native MBridge training; external callers such as NeMo-RL may provide the same caller-packed THD contract from their own data pipeline.
  • NemotronOmniModel receives the full packed THD sequence, builds and inserts media embeddings, and only then selects the context-parallel portion of the merged embedding sequence.
  • Packing does not run inside model forward.

Changelog

  • Add NemotronOmniModel, reusing the existing RADIO encoder, multimodal projector, and Nemotron-H language model components.
  • Preserve expanded image placeholders through the canonical collator, with packing performed at the data boundary.
  • Keep the legacy LLaVA collator and provider/bridge compatibility path explicit for compatible legacy checkpoints.
  • Add the caller-packed THD model contract needed to insert media embeddings before context-parallel slicing.
  • Register canonical Nano and Super Omni HF architectures with AutoBridge, with versioned contracts that reject ambiguous legacy checkpoints.
  • Add canonical import/export mappings for the top-level model namespace and reverse MTP QKV mapping support for Nemotron-H.
  • Add conversion, model-forward, media-alignment, packing, Mamba-boundary, activation, and compatibility tests.

Validation

  • Ruff lint and formatting checks passed.
  • Focused unit and functional validation passed, including the two-rank caller-packed context-parallel path.
  • The corresponding NeMo-RL distributed suite passed CP2 checkpoint, TP/SP, TP2+CP2, padded multi-row attention, and PP2 coverage against this exact MBridge head.
  • Nano image-GRPO validation covering generation, training, repeated vLLM refit, checkpointing, and probability agreement is recorded in W&B.

Scope and follow-up

This PR implements the canonical image/text foundation used by the current Nano migration. Audio and video support is layered separately in #5007. Super/MTP integration and validation remain follow-up work. This implementation does not require changes to Megatron-LM source.

GitHub Actions CI

See the CI section in the contributing guide for how to trigger CI. An NVIDIA developer may need to approve and trigger CI for external contributors.

Before this PR is ready for review

  • Read and followed the contributor guidelines.
  • Added the necessary tests.
  • Added or updated the necessary documentation.
  • No new optional dependency is introduced.

@copy-pr-bot

copy-pr-bot Bot commented Jul 15, 2026

Copy link
Copy Markdown

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
aroshanghias-nvd marked this pull request as ready for review July 16, 2026 01:13
@yaoyu-33 yaoyu-33 added area:model Model implementations and HF bridge logic feature New capabilities, enhancements, or enablement work full-test-suite needs-review PR is ready for code review and waiting on a reviewer labels Jul 16, 2026
@aroshanghias-nvd
aroshanghias-nvd force-pushed the aroshanghias/nemotron-omni-main-migration branch 2 times, most recently from 8b95417 to c9473b1 Compare July 16, 2026 23:35
Comment thread src/megatron/bridge/models/nemotron_omni/data/collate_fn.py Outdated
@yaoyu-33 yaoyu-33 added waiting-on-customer Waiting on the original author to respond and removed needs-review PR is ready for code review and waiting on a reviewer labels Jul 25, 2026
Signed-off-by: Ali Roshan Ghias <aroshanghias@nvidia.com>
Signed-off-by: Ali Roshan Ghias <aroshanghias@nvidia.com>
Signed-off-by: Ali Roshan Ghias <aroshanghias@nvidia.com>
Signed-off-by: Ali Roshan Ghias <aroshanghias@nvidia.com>
Signed-off-by: Ali Roshan Ghias <aroshanghias@nvidia.com>
Signed-off-by: Ali Roshan Ghias <aroshanghias@nvidia.com>
@aroshanghias-nvd
aroshanghias-nvd force-pushed the aroshanghias/nemotron-omni-main-migration branch from 545d6e8 to 7a9e305 Compare July 29, 2026 14:17
@aroshanghias-nvd

Copy link
Copy Markdown
Contributor Author

/ok to test 7a9e305

@cuichenx

Copy link
Copy Markdown
Contributor

audio/video capabilities in #5007

Signed-off-by: Chen Cui <chcui@nvidia.com>
@cuichenx

Copy link
Copy Markdown
Contributor

/ok to test 124b19b

cuichenx
cuichenx previously approved these changes Jul 29, 2026
@cuichenx cuichenx added the ready-to-merge PR is approved, current, and only waiting for CI to pass before merge label Jul 29, 2026
Signed-off-by: Chen Cui <chcui@nvidia.com>
@cuichenx

Copy link
Copy Markdown
Contributor

/ok to test f461eae

@cuichenx
cuichenx merged commit f6c9292 into NVIDIA-NeMo:main Jul 30, 2026
152 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

area:model Model implementations and HF bridge logic feature New capabilities, enhancements, or enablement work full-test-suite ready-to-merge PR is approved, current, and only waiting for CI to pass before merge waiting-on-customer Waiting on the original author to respond

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants