feat(nemotron-omni): add Megatron Nano image GRPO support by aroshanghias-nvd · Pull Request #3290 · NVIDIA-NeMo/RL

aroshanghias-nvd · 2026-07-20T11:32:51Z

What does this PR do?

Add Megatron-based Nemotron Nano Omni image-GRPO support using NeMo-RL's existing sequence-packing data path and the dedicated MBridge NemotronOmniModel.

Fresh review requested: The branch is now synchronized with current main and uses the reviewed data-boundary packing design. The earlier model-owned-packing implementation is obsolete; please review this series from the beginning.

Cross-repository dependency: NVIDIA-NeMo/Megatron-Bridge#4885 has merged. The required Nemotron Omni support is included in the MBridge pin inherited from NeMo-RL main (573e088c); this PR does not change the submodule pin.

The ownership boundary is:

NeMo-RL groups and packs the processor-expanded token sequences using its existing Megatron data pipeline and constructs PackedSeqParams.
NeMo-RL passes the full packed THD row and multimodal inputs to MBridge.
MBridge NemotronOmniModel creates and inserts media embeddings, then selects the context-parallel portion of the merged embedding sequence.
Packing does not run inside model forward, and this path does not use the historical Nemotron Omni LLaVAModel collapse/expand behavior.

The PR includes:

capability-based dispatch for models that consume caller-packed full THD input and slice context-parallel inputs after embedding insertion;
a clear error for legacy LLaVA-layout Nemotron Omni checkpoints;
dynamic-resolution and historical tiled-image metadata handling;
generation compatibility for raw policy logprobs, context-length capping, and RADIO LayerScale initialization;
maintained single-node Clevr and four-node MMPR Nano image-GRPO recipes.

The branch is synchronized with NeMo-RL main's vLLM 0.25.1 dependency. This PR does not change external vLLM source or the vLLM dependency pin; its compatibility behavior is implemented in NeMo-RL's existing integration and runtime-patch layer.

Issues

No public issue.

Usage

The maintained recipes added by this PR are:

examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml
examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.yaml

Validation

A production-shaped four-node Nano image-GRPO parity run with vLLM 0.25.1 exercised generation, rewards, raw policy logprobs, training, repeated vLLM refit, and checkpointing through step 23. The PR #3290 validation run tracks the production packing-refactor reference closely, without the prior TMPE spikes.
The matched A/B used validation-only per-call deterministic seeding instrumentation; that instrumentation is not included in this PR.
The focused distributed Megatron suite passed: CP2 checkpoint, TP/SP, TP2+CP2, padded multi-row attention, and PP2 (5 passed).
Focused vLLM integration tests passed (8 passed).
uv lock --check and git diff --check passed.
Protected NeMo-RL CI is required before merge.

Scope and follow-up

This PR covers the Nano image/text path. Caller-packed context parallelism with MTP is deliberately unsupported until the separate Super/MTP integration adds and validates the required MTP-mask alignment. Audio and video support is layered on the MBridge foundation separately.

Before this PR is ready for review

Read and followed the contributor guidelines.
Added the necessary unit and functional tests.
Ran focused local and distributed validation.
Updated docs/guides/nemotron-3-nano-omni.md with the Megatron architecture, supported scope, checkpoint compatibility, recipes, and usage.

copy-pr-bot · 2026-07-20T11:32:54Z

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.

github-actions · 2026-07-20T11:33:01Z

✅ Submodule Fast-Forward Check Results

Check based on commit: df0db31 (PR #3227 from aroshanghias/nemotron-omni-main-migration)

✅ Submodules that are properly updated:

Megatron-Bridge: ✅ PR branch is ahead of main branch (fast-forward)

All submodule changes look good! ✨

Signed-off-by: Ali Roshan Ghias <aroshanghias@nvidia.com>

yfw · 2026-08-04T21:14:50Z

/ok to test a7e52a3

…model Addresses code-review findings on the Nemotron Omni PR: - Restore `logprobs_mode: processed_logprobs` in the root exemplars (grpo/ppo/distillation) and matching reference_configs. Commit bb89288 removed it from the configs and made the code optional, which silently deferred the default to vLLM 0.25.1's `raw_logprobs` for every recipe that does not set the key. - Reject mixed-resolution batches on the AutoModel path in `filter_multimodal_kwargs_for_model`: it drops `imgs_sizes` and cannot crop the `pad_to_max_shape` padding, so mixed-resolution inputs would feed padded pixels to the vision encoder and mismatch the placeholder count. Fail loud instead of corrupting silently, with unit tests covering the mixed-resolution (raises) and uniform-resolution (allowed) cases. - Run ruff-format on the affected files. Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

yfw · 2026-08-05T00:15:17Z

/ok to test 822d7b1

yfw

Approved @aroshanghias-nvd 's changes up until my last commit

…lies to Per review feedback, reference the concrete model (nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16) whose HF forward accepts pixel_values but not imgs_sizes, so it is clear which model the guard is for. Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com> Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

yfw · 2026-08-05T02:25:34Z

/ok to test b8d7523

yfw · 2026-08-05T02:45:16Z

/ok to test c65b1f0

Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

yfw · 2026-08-05T04:21:08Z

/ok to test aff74f9

yfw · 2026-08-05T06:14:06Z

/ok to test a9a9a87

aroshanghias-nvd requested review from a team as code owners July 20, 2026 11:32

github-actions Bot added the Documentation Improvements or additions to documentation label Jul 20, 2026

aroshanghias-nvd added the CI:L1 Run doctests, unit tests, and functional tests label Jul 20, 2026

aroshanghias-nvd requested a review from yfw July 20, 2026 11:33

aroshanghias-nvd mentioned this pull request Jul 20, 2026

feat(nemotron-omni): add Megatron Nano image GRPO support #3227

Closed

4 tasks

rohitrango reviewed Jul 20, 2026

View reviewed changes

Comment thread nemo_rl/data/multimodal_utils.py

yfw reviewed Jul 23, 2026

View reviewed changes

JeffPengCoder mentioned this pull request Jul 23, 2026

fix: harden multimodal Gym training lifecycle #3324

Open

6 tasks

rohitrango reviewed Jul 24, 2026

View reviewed changes

Comment thread nemo_rl/data/multimodal_utils.py

Comment thread nemo_rl/models/generation/vllm/vllm_backend.py Outdated

Comment thread nemo_rl/models/generation/vllm/vllm_worker.py

aroshanghias-nvd added 12 commits July 29, 2026 05:49

chore: update Megatron-Bridge pin

4c9b5a0

Signed-off-by: Ali Roshan Ghias <aroshanghias@nvidia.com>

feat(megatron): integrate canonical Nemotron Omni model

a174f29

Signed-off-by: Ali Roshan Ghias <aroshanghias@nvidia.com>

feat(multimodal): support Nemotron Omni image batches

3974f7c

Signed-off-by: Ali Roshan Ghias <aroshanghias@nvidia.com>

feat(vllm): support Nemotron Omni generation parity

7615435

Signed-off-by: Ali Roshan Ghias <aroshanghias@nvidia.com>

feat(nemotron-omni): add Nano GRPO recipes

e0446d0

Signed-off-by: Ali Roshan Ghias <aroshanghias@nvidia.com>

style(nemotron-omni): apply repository formatting

da1a684

Signed-off-by: Ali Roshan Ghias <aroshanghias@nvidia.com>

docs(nemotron-omni): document Megatron GRPO support

e26067e

Signed-off-by: Ali Roshan Ghias <aroshanghias@nvidia.com>

refactor: generalize multimodal tensor padding

4fc74f6

Signed-off-by: Ali Roshan Ghias <aroshanghias@nvidia.com>

refactor(megatron): pack Nemotron Omni inputs before forward

bbd5b18

Signed-off-by: Ali Roshan Ghias <aroshanghias@nvidia.com>

test(nemotron-omni): cover TP CP and sequence parallel

88110ed

Signed-off-by: Ali Roshan Ghias <aroshanghias@nvidia.com>

test(nemotron-omni): cover pipeline-parallel schedule

1c01cd4

Signed-off-by: Ali Roshan Ghias <aroshanghias@nvidia.com>

fix(nemotron-omni): select compatible attention backend

3b0776a

Signed-off-by: Ali Roshan Ghias <aroshanghias@nvidia.com>

aroshanghias-nvd added 3 commits August 4, 2026 10:49

Merge remote-tracking branch 'origin/main' into pr3290-main-merge

94365da

Signed-off-by: Ali Roshan Ghias <aroshanghias@nvidia.com>

Merge remote-tracking branch 'origin/main' into pr3290-main-merge

c0dc3f2

Signed-off-by: Ali Roshan Ghias <aroshanghias@nvidia.com>

Merge remote-tracking branch 'origin/main' into pr3290-main-merge

a7e52a3

Signed-off-by: Ali Roshan Ghias <aroshanghias@nvidia.com>

yfw added CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) and removed CI:L1 Run doctests, unit tests, and functional tests labels Aug 4, 2026

yfw reviewed Aug 4, 2026

View reviewed changes

Comment thread nemo_rl/models/generation/vllm/vllm_worker.py

Comment thread nemo_rl/data/processors.py

Comment thread nemo_rl/models/megatron/setup.py

yfw force-pushed the aroshanghias/nemotron-omni-main-migration branch from 79a4797 to 822d7b1 Compare August 5, 2026 00:14

yfw previously approved these changes Aug 5, 2026

View reviewed changes

rohitrango reviewed Aug 5, 2026

View reviewed changes

Comment thread nemo_rl/models/automodel/data.py

rohitrango previously approved these changes Aug 5, 2026

View reviewed changes

yfw dismissed stale reviews from rohitrango and themself via b8d7523 August 5, 2026 02:24

Merge branch 'main' into aroshanghias/nemotron-omni-main-migration

c65b1f0

raise nightly GPU-hour limit to 3800

aff74f9

Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>

terrykong approved these changes Aug 5, 2026

View reviewed changes

yfw enabled auto-merge (squash) August 5, 2026 05:10

Merge branch 'main' into aroshanghias/nemotron-omni-main-migration

a9a9a87

terrykong approved these changes Aug 5, 2026

View reviewed changes

yfw merged commit a704801 into main Aug 5, 2026
142 of 145 checks passed

yfw deleted the aroshanghias/nemotron-omni-main-migration branch August 5, 2026 09:11

zpqiu mentioned this pull request Aug 19, 2026

feat(models): add Qwen3.8-27B Megatron support #3678

Merged

4 tasks

rohitrango mentioned this pull request Aug 24, 2026

feat(vlm): route VLM GRPO through TQ trainer when data_plane.enabled #2957

Open

Uh oh!

Conversation

aroshanghias-nvd commented Jul 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

What does this PR do?

Issues

Usage

Validation

Scope and follow-up

Before this PR is ready for review

Uh oh!

copy-pr-bot Bot commented Jul 20, 2026

Uh oh!

github-actions Bot commented Jul 20, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yfw commented Aug 4, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yfw commented Aug 5, 2026

Uh oh!

yfw left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

yfw commented Aug 5, 2026

Uh oh!

yfw commented Aug 5, 2026

Uh oh!

yfw commented Aug 5, 2026

Uh oh!

yfw commented Aug 5, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

aroshanghias-nvd commented Jul 20, 2026 •

edited

Loading