feat(vlm): Nemotron Super Omni GRPO with MTP under context parallelism by DanialTaheri · Pull Request #3494 · NVIDIA-NeMo/RL

DanialTaheri · 2026-08-04T22:01:24Z

Summary

Enables Nemotron Super Omni GRPO on the multimodal NeMo-Gym path, including Multi-Token Prediction under context parallelism.

Fixes

Multimodal rollouts

nemo_gym.py — the Super chat template zeroes its structured-image count when prompt text already contains a literal <image>. 1,149 of 60,299 rows in the blended Super dataset do (100% of two MMPR/NLVR2 sources), so the template emitted 601 placeholders against 600 projected features and _merge_projected_media aborted the step. Adds sanitize_image_placeholders and pad_dynamic_image_shapes for ragged tiling.
vllm_worker_async.py — the embedded OpenAI server is built directly rather than through vLLM's CLI, where chat-template file paths are normally loaded. OpenAIServingRender expects literal Jinja, so a path was rendered as a template and every <image> marker was dropped.
nemo_rl/__init__.py, worker_groups.py — prefer an explicit HF_MODULES_CACHE and make it importable in the current interpreter, and forward it to isolated worker initializers. Without this a trust_remote_code actor cannot import transformers_modules while unpickling its constructor arguments.

MTP under context parallelism (three bugs, each masked by the previous)

data.py — Nemotron Omni consumes a full THD row and applies CP itself after rebuilding multimodal embeddings, so input_ids is passed unsharded. The MTP loss mask always took the CP-local shard, so mask and tokens had different layouts. Both call sites previously raised NotImplementedError rather than reconcile them; now the mask follows the same layout rule as input_ids.
train.py — a model that slices CP itself returns (output, sliced_loss_mask) when handed a full-sequence loss_mask. That tuple reached the loss wrapper, which called .narrow() on it.
run_multimodal_grpo_nemo_gym.py — did not pass trains_mtp to configure_generation_config (the text entrypoint does), so _mtp_weights_from_refit stayed False and the worker was routed into load_mtp_weights_from_disk. That path expects the MTP layer indexed positionally after the backbone; Nemotron checkpoints use an mtp.* namespace, so MTP speculative decoding failed with No MTP layer weights for layers [88] found. It also means the trained MTP head is now refit each step rather than the drafter staying frozen.

Tower controls — radio_force_cpe_eval_mode and the four freeze_* keys, applied only when present so provider defaults survive omission.

Launcher — Hydra's delete of compilation_config.cudagraph_capture_sizes errors when a recipe does not pin it; gated behind MTP_DROP_CUDAGRAPH_CAPTURE_SIZES (default unchanged).

Known dependency

MTP at CP>1 also needs a Megatron-LM fix, submitted separately: the MTP block calls the model's embedding directly and never applies the sequence-parallel scatter that multimodal models rely on their outer forward to perform, so _concat_embeddings fails (8192 vs 1024 at 16384 tokens, CP=2/TP=8). Carried here as patches/mtp-sp-scatter.patch until that lands.

Testing

Nemotron Super Omni 120B-A12B, TP8/EP16/CP2, 16 nodes, async GRPO:

50 steps at 32 prompts x 16 generations
50 steps at 128 prompts x 16 generations (2,048 rollouts/step), mean reward 0.45 over steps 1-25 rising to 0.53 over 26-50
MTP trained at mtp_loss_scaling_factor=0.1 with speculative decoding at ~75% acceptance, flat across the run
grad_norm stable in 0.035-0.145

Full metrics, curves and per-run configuration: W&B report — Upstream Super MTP to main NeMoRL

copy-pr-bot · 2026-08-04T22:01:27Z

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-08-09T23:18:45Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 605836d (PR #3494 from smohsenitahe/super-omni-mtp-cp)

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

github-actions · 2026-08-10T18:00:41Z

✅ Submodule Fast-Forward Check Results

Check based on commit: 3956216 (PR #3494 from smohsenitahe/super-omni-mtp-cp)

✅ Submodules that are properly updated:

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

All submodule changes look good! ✨

yfw

Thanks for this — enabling Nemotron Super Omni GRPO with MTP under context parallelism is a substantial piece of work, and the validation in the description (50 steps at 128x16, mean reward 0.45 -> 0.53, ~75% MTP acceptance held flat, grad_norm 0.035-0.145, plus the W&B report) is exactly the evidence we want on a change like this.

The core MTP-under-CP work holds up under scrutiny, and we want to say that clearly before the list below. We verified:

The MTP loss-mask layout change is correct, and it's the only layout the model accepts — _apply_context_parallel_sharding raises on inconsistent tensor lengths, so the previous CP-local read would have failed loudly at step 1 rather than silently mis-training. NemotronOmniModel is the only class in the pinned Bridge tree setting model_slices_context_parallel_inputs = True, so this is live code for a real model, not scaffolding.
The train.py tuple unwrap is necessary and safe. We hunted for a model that would lose element [1] and found none: HybridModel/GPTModel return Tensors, the fused-CE path returns a Tensor, Qwen3VL's dist-train PP path returns a dict, and mcore's output_processor hook (the one real candidate) is never installed anywhere in nemo_rl/.
All five tower-control fields exist on the Bridge provider, are non-deprecated, and are applied the way Bridge's own recipes apply them.
Your new unit test passes (pytest tests/unit/models/megatron/test_megatron_data.py --mcore-only -> 36 passed; the 3 failures are a missing transformer_engine in our environment, with zero assertion failures).
Config validation is green: 584 passed across test_config_validation, test_config_v2, and test_recipes_and_test_suites.
No cluster paths, account names, or usernames leak into any new file — only the intended /path/to/... placeholders, with a hard-fail preflight behind them.

Note that the review assumes a dependency on #3568 . We checked the coupling since this is meant to land on that bump: the Omni contract is unchanged at Bridge 0c565c9a0, and unwrap_model at Megatron-LM d12f6c8c9 already handles the FullyShardedDataParallelV1/V2 split. That last one is worth knowing — had it not, getattr(chunk, "model_slices_context_parallel_inputs", False) would have quietly returned False and the Omni path would have regressed to CP-local slicing with no error. There is also no merge conflict: your only hunk in megatron_policy_worker.py is ~2300 lines from #3568's. One thing we could not verify: _apply_context_parallel_sharding gains a padding_mask parameter at the bumped pin — almost certainly inert since NeMo-RL never passes it, but worth one grep after you rebase.

Most of what follows is mechanical. The one blocker is that all three shipped test drivers invoke a script that has never existed on main.

Reviewed by a team of six agents plus two independent adversarial passes, which dropped seven candidate findings (including two that would have been wrong to send you) before this review.

Generated by Claude Code

yfw · 2026-08-11T15:03:02Z

/ok to test df8a81a

yfw

Re-reviewed at fc11b496. Thanks for turning this around so quickly. The substantive work all landed, and I verified it rather than taking the commit messages at face value.

Confirmed fixed: the entrypoint blocker (all three drivers; zero repo-wide references to the dead path), the EXTRA_OVERRIDES -> EXTRA_HYDRA_ARGS rename, the MTP grad-clipping warning (the new text checks out against Megatron-LM d12f6c8c9), checkpoint selection on all three recipes, minimize-check (exit 0 across 197 configs, re-confirmed with a deliberate control probe), the toctree entry and MTP section, the ragged imgs_sizes reconstruction (verified on genuinely differently-shaped tiles), the MTP test hardening (now load-bearing), and the patches/ deletion — the SP-scatter fix really is in the pinned Megatron-LM, so dropping it is correct.

Three of the follow-ups below are my fault, not yours. Two of my earlier comments used ```suggestion blocks that weren't valid at the line they were anchored to, and applying them left nemo_rl unimportable (NameError at __init__.py:287) and nemo_rl/models/megatron/setup.py non-compiling. A third was bad advice from me about vLLM's default_chat_template_kwargs. Details in the threads.

Merge note: the only conflict is uv.lock. #3568 landed on main (46ab18ce1) and relocked it there, while this branch bumped Megatron-Bridge without relocking — merging main and taking its uv.lock should clear it, and your own bump commit is now redundant since main already pins 0c565c9a0.

Generated by Claude Code

yfw

Follow-up after the repair round. Everything from the last review is verified fixed — import restored, both VLM override sites aligned, the chat-template defaults now reaching all three serving consumers, the W&B opt-out exported, and new coverage for the placeholder sanitiser. One item from the original findings is still open; details inline.

Generated by Claude Code

yfw · 2026-08-13T19:55:00Z

/ok to test a635eee

Adds two env.nemo_gym options used by the Super Omni recipe: sanitize_image_placeholders - the Super chat template zeroes its structured-image count when the prompt text already contains a literal <image> token. 392 rows of the blended Super dataset carry one structured image plus two literal <image> strings, so the template emitted 601 placeholders against 600 projected features and _merge_projected_media aborted the step on the strict alignment check. Normalizing literal placeholders lets the structured items alone decide media placement. Validated across all 60,299 rows of the training file. pad_dynamic_image_shapes - Super's dynamic tiling yields ragged pixel_values across images in a microbatch; pad to max shape while keeping exact imgs_sizes so the projector still slices correctly. Signed-off-by: DanialTaheri <smohsenitahe@nvidia.com>

…s importable Adds the megatron_cfg keys the Super Omni recipe sets: radio_force_cpe_eval_mode, freeze_vision_model, freeze_vision_projection, freeze_sound_encoder, freeze_sound_projection. These are threaded through community_import/setup as VLM config overrides rather than hardcoded, and only keys present in the recipe are applied, so a provider's checkpoint defaults survive omission. worker_groups also forwards HF_HOME, HF_MODULES_CACHE and PYTHONPATH to isolated worker initializers; without them a trust_remote_code actor cannot import transformers_modules while deserializing its arguments. Signed-off-by: DanialTaheri <smohsenitahe@nvidia.com>

The RADIO LayerScale portion of this change is now upstream (18fb15b 'fix(vllm): isolate Nemotron Omni LayerScale handling'), so only the two remaining fixes are carried here: - The embedded OpenAI server is constructed directly rather than through vLLM's CLI, where chat-template file paths are normally loaded. OpenAIServingRender expects literal Jinja, so passing a path made Transformers render the path string and drop every <image> marker. - patch_transformers_module_dir now prefers an explicit HF_MODULES_CACHE and inserts it into the running interpreter's sys.path, so Ray actors can import trust_remote_code classes while unpickling arguments. Signed-off-by: DanialTaheri <smohsenitahe@nvidia.com>

…mers Handing default_chat_template_kwargs to the shared renderer does not cover the other endpoints, so popping the key out of the init bag starved them. Read against vLLM 0.25.1 as installed: ServingTokenization takes its own (serve/tokenize/serving.py:40), stores it (:54) and passes that -- not the renderer's -- into preprocess_chat (:89). OpenAIServingChat likewise stores its own (chat_completion/serving.py:142) and builds the reasoning parser from it (:190, :202). So /tokenize rendered with {}, reintroducing the divergence the deleted merge block existed to prevent, and OpenAIServingChat received {} where the native spelling used to pass straight through. Latent for the recipe that sets it (circle-click sets enable_thinking: true, matching the fallback), but false would silently misconfigure the parser. Normalize the legacy alias onto the native key and hand the same value to the renderer, the chat serving and the tokenization serving, which is what vLLM's own api_server does. Also drop the `pop(A) or pop(B)` idiom: it short-circuits, so a recipe setting both spellings left chat_template_kwargs in the bag and OpenAIServingChat(**kwargs) raised TypeError. Both keys are now consumed unconditionally, with the native one winning. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: DanialTaheri <smohsenitahe@nvidia.com>

The two options above it export; this one did not. A bare assignment stays in the reader's shell and never reaches the child `bash super_omni_launch.sh`, so the preflight sees an empty EXTRA_HYDRA_ARGS, the logger.wandb_enabled regex does not match, and the launcher aborts with the "W&B logging is on but WANDB_API_KEY is unset" error this line exists to avoid. Every other assignment in the guide is prefix form, which exports for that command; this was the only standalone-block one. Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com> Signed-off-by: DanialTaheri <smohsenitahe@nvidia.com>

packed_broadcast_producer concatenates the whole export stream into a single buffer and broadcasts it over a CUDA collective. torch.cat refuses to mix devices, so one host tensor anywhere in the stream fails the transfer. Exporters can legitimately yield host tensors: constants that are plain attributes rather than registered buffers never follow .to(device), and an export hook may drop to the host for a numpy round trip. For Nemotron Omni this was four constants -- an audio filterbank and window, and two image normalization statistics -- none of which are learnable, so nothing had a reason to place them on the device. Normalize in the producer rather than requiring every upstream exporter to agree: the producer owns its buffer's device. Signed-off-by: DanialTaheri <smohsenitahe@nvidia.com>

return_tensors=None makes the processor hand back every output as plain Python lists, not only the ragged pixel_values that mode was requested for. Downstream expects tensors -- input_ids in particular is rank-checked -- so a batch that took this path failed with "input_ids must be a one- or two dimensional tensor". Restore the rest of the batch to what return_tensors="pt" would have produced. Two images of equal resolution come back already stacked, so the restoration runs for every batch on this path rather than only behind the ragged branch. Values that resist conversion are genuinely ragged per-image metadata and are left for their own handling. Signed-off-by: DanialTaheri <smohsenitahe@nvidia.com>

DanialTaheri · 2026-08-13T20:25:30Z

/ok to test 4b28741

Sanitization rewrote rollout data so it would fit the model: it dropped a literal <image> when images were attached, and rendered it as the word "image" when they were not. That changed the author's prose -- "for any character <image> there is" became "for any character image there is" -- and it silently patched malformed rows, hiding a real defect rather than surfacing it. It also could not express the case it most needed to. There is no rewrite that both preserves the prose and tells the model that a given token is not an anchor. The validity mask says exactly that, so the rewrite is no longer the only option. BREAKING CHANGE: the sanitize_image_placeholders config flag is removed. A blend whose placeholder and image counts disagree now fails at the media merge instead of being rewritten at rollout time. That failure is the correct signal -- such rows are malformed -- but it means the blend must be filtered before training rather than repaired during it. Signed-off-by: DanialTaheri <smohsenitahe@nvidia.com>

The launcher resolves the recipe with tools/config_cli.py so it reads the same cluster block the driver will, since a raw scrape cannot follow `defaults:`. It called that resolver through `uv run`, which fails when the environment is not materialized on the submit host -- and then fell back to the raw scrape, which reports "could not read cluster.num_nodes" for any recipe that inherits its cluster block rather than declaring one. Try the project venv first, since it already has omegaconf and needs no resolution, and keep `uv run` as the documented path behind it. Signed-off-by: DanialTaheri <smohsenitahe@nvidia.com>

The guide and its toctree entry are removed. Main no longer carries the sibling Nemotron guides this one sat alongside, so registering it here would leave a lone entry in a section the rest of that family has left, and the recipe and launcher document themselves. No dangling references remain: docs/index.md held the only link. Signed-off-by: DanialTaheri <smohsenitahe@nvidia.com>

DanialTaheri · 2026-08-13T20:57:22Z

/ok to test d71b1e2

These tests set HF_HOME and assert on the directory derived from it, but HF_MODULES_CACHE now outranks HF_HOME, so a value in the developer's environment silently decides what they assert. The variable was never read before this PR, which is why an ambient value used to be harmless. Clear it once for the class rather than passing clear=True to each patch.dict, which would also wipe unrelated ambient variables. The file did not import pytest, so the fixture needs that import to avoid taking the whole module down at collection. Verified: 14 pass with the variable unset, set, and set alongside an ambient HF_HOME. Signed-off-by: DanialTaheri <smohsenitahe@nvidia.com>

…ize tests Removing those tests shortened the import block and left three blank lines before the next section. The repo runs ruff's import rules as a separate pre-commit hook with --fix, so it rewrote the file and the hook failed on the modification rather than reporting a diagnostic. Plain `ruff check` does not select the I rules, which is why this passed locally and failed in CI. Signed-off-by: DanialTaheri <smohsenitahe@nvidia.com>

yfw · 2026-08-13T21:25:04Z

/ok to test 5be93ae

Routing the chat template through vLLM's own load_chat_template added an import of vllm.entrypoints.chat_utils, but the fake module tree these tests install stops at vllm.entrypoints. Importing a submodule of a plain ModuleType fails with "'vllm.entrypoints' is not a package", so the server setup raised before the assertion it was meant to exercise. Register the leaf alongside the other stubbed entrypoints modules. The loader is not asserted on here; it only has to exist for the import to resolve. Signed-off-by: DanialTaheri <smohsenitahe@nvidia.com>

DanialTaheri · 2026-08-13T23:20:08Z

/ok to test dfc4e56

_postprocess_nemo_gym_to_nemo_rl_result is called unbound against lightweight stand-ins that define only the attributes they exercise, which is why the same method already reads the processor as getattr(self, "_processor", None). The pad flag was added as a bare attribute access and did not follow that, so a stand-in without it raises AttributeError instead of taking the default. Surfaced by test_nemo_gym_dedup_omits_actor_initial_tensor_and_preserves_later_media, which reached this branch for the first time after rebasing onto main: the test landed upstream on 2026-08-11 and the flag is new on this branch, so neither side saw the combination until now. Signed-off-by: DanialTaheri <smohsenitahe@nvidia.com>

…umers The async HTTP server builds three objects that each render chat messages -- OnlineRenderer, OpenAIServingChat and ServingTokenization -- and they do not share one copy of the template kwargs. Handing the value to only one makes /tokenize render differently from /v1/chat/completions on the same conversation. Nothing asserted that until now. Covers the four properties the wiring has to hold: both spellings reach all three consumers, the legacy chat_template_kwargs is renamed rather than merely read (the bag is splatted into OpenAIServingChat, which rejects an argument it does not declare), the native spelling wins when both are present, and an absent value renders as {} rather than None. Verified load-bearing by mutation: restoring the earlier `pop(native) or pop(legacy)` form -- which short-circuits on a truthy native value and leaves the legacy key in the bag -- fails test_native_spelling_wins_and_legacy_is_dropped. Kept in a new file with its own vLLM stub tree rather than extending test_vllm_generation.py, whose harness belongs to the reasoning-parser tests. Signed-off-by: DanialTaheri <smohsenitahe@nvidia.com>

DanialTaheri · 2026-08-14T04:45:28Z

/ok to test 6163e79

DanialTaheri requested review from a team as code owners August 4, 2026 22:01

github-actions Bot added the Documentation Improvements or additions to documentation label Aug 4, 2026

rohitrango force-pushed the rohit/gymv-mm-integration-v2 branch from d60004f to d3e02b4 Compare August 5, 2026 16:55

rohitrango requested a review from a team as a code owner August 5, 2026 16:55

Base automatically changed from rohit/gymv-mm-integration-v2 to main August 7, 2026 03:51

DanialTaheri force-pushed the smohsenitahe/super-omni-mtp-cp branch 2 times, most recently from e2c0e9c to 605836d Compare August 9, 2026 23:17

DanialTaheri force-pushed the smohsenitahe/super-omni-mtp-cp branch from 3956216 to df8a81a Compare August 11, 2026 05:26

yfw reviewed Aug 11, 2026

View reviewed changes

yfw added the CI:Lfast Runs a fast test suite and re-use nightly `main` container (but sync dependencies to PRs version) label Aug 11, 2026

DanialTaheri force-pushed the smohsenitahe/super-omni-mtp-cp branch from 4ef0e67 to 951f738 Compare August 11, 2026 22:01

yfw reviewed Aug 12, 2026

View reviewed changes

DanialTaheri force-pushed the smohsenitahe/super-omni-mtp-cp branch from fc11b49 to 60ff0bf Compare August 12, 2026 15:38

yfw reviewed Aug 13, 2026

View reviewed changes

Comment thread nemo_rl/__init__.py

DanialTaheri added 3 commits August 13, 2026 13:11

DanialTaheri and others added 4 commits August 13, 2026 13:13

DanialTaheri force-pushed the smohsenitahe/super-omni-mtp-cp branch from a635eee to 4b28741 Compare August 13, 2026 20:16

cuichenx mentioned this pull request Aug 13, 2026

ci: fix Ruff import hook arguments #3629

Closed

DanialTaheri added 3 commits August 13, 2026 13:49

DanialTaheri force-pushed the smohsenitahe/super-omni-mtp-cp branch from 4b28741 to d71b1e2 Compare August 13, 2026 20:52

DanialTaheri mentioned this pull request Aug 13, 2026

feat(vlm): carry a media-token validity mask through sequence packing #3631

Merged

DanialTaheri added 2 commits August 13, 2026 14:00

yfw previously approved these changes Aug 13, 2026

View reviewed changes

yfw enabled auto-merge (squash) August 13, 2026 22:41

DanialTaheri dismissed yfw’s stale review via dfc4e56 August 13, 2026 23:14

yfw previously approved these changes Aug 14, 2026

View reviewed changes

DanialTaheri dismissed yfw’s stale review via c1cb4f0 August 14, 2026 04:21

yfw approved these changes Aug 14, 2026

View reviewed changes

yfw merged commit f38b5a5 into main Aug 14, 2026
82 checks passed

yfw deleted the smohsenitahe/super-omni-mtp-cp branch August 14, 2026 07:15

nvilia mentioned this pull request Aug 14, 2026

feat(mopd): add Super Omni multimodal distillation #3430

Merged

5 tasks

Uh oh!

Conversation

DanialTaheri commented Aug 4, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Fixes

Known dependency

Testing

Uh oh!

copy-pr-bot Bot commented Aug 4, 2026

Uh oh!

github-actions Bot commented Aug 9, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

github-actions Bot commented Aug 10, 2026

✅ Submodule Fast-Forward Check Results

✅ Submodules that are properly updated:

Uh oh!

yfw left a comment

Choose a reason for hiding this comment

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 11, 2026

Uh oh!

yfw left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

yfw left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

yfw commented Aug 13, 2026

Uh oh!

DanialTaheri commented Aug 13, 2026

Uh oh!

DanialTaheri commented Aug 13, 2026

Uh oh!

yfw commented Aug 13, 2026

Uh oh!

DanialTaheri commented Aug 13, 2026

Uh oh!

DanialTaheri commented Aug 14, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

DanialTaheri commented Aug 4, 2026 •

edited

Loading