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

docs/guides/nemotron-3-nano-omni.md

@@ -1,6 +1,8 @@
 # Nemotron 3 Nano Omni
 
-This guide explains how to post-train the Nemotron 3 Nano Omni vision-language model with GRPO using NeMo RL on the AutoModel backend.
+This guide explains how to post-train the Nemotron 3 Nano Omni vision-language model with GRPO using NeMo RL. Both the AutoModel and Megatron backends are supported for image-and-text training.
+
+## AutoModel backend
 
 It covers two recipes:
 
@@ -9,7 +11,7 @@ It covers two recipes:
 
 Both share the same checkpoint, model code, and reward pipeline; they differ only in the dataset, reward functions, and node count.
 
-## Recipe 1 — CLEVR-CoGenT (single-node)
+### Recipe 1 — CLEVR-CoGenT (single-node)
 
 The CLEVR-CoGenT recipe uses [`examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-automodel-ep8.v1.yaml`](../../examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-automodel-ep8.v1.yaml). It expects 8 GPUs on a single node, EP=8 across the experts, and TP=8 in vLLM.
 
@@ -47,7 +49,7 @@ uv run examples/run_vlm_grpo.py --config examples/configs/recipes/vlm/vlm_grpo-n
     cluster.gpus_per_node=8 cluster.num_nodes=1
 ```
 
-## Recipe 2 — MMPR-Tiny (4-node Slurm)
+### Recipe 2 — MMPR-Tiny (4-node Slurm)
 
 The MMPR-Tiny recipe uses [`examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-automodel-ep8.v1.yaml`](../../examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-automodel-ep8.v1.yaml). Differences vs. the CLEVR recipe:
 
@@ -109,3 +111,37 @@ sbatch \
 ```
 
 To run on a different node count, change `NUM_NODES` and the `--nodes` flag.
+
+## Megatron backend
+
+The Megatron backend uses a dedicated `NemotronOmniModel` supplied by Megatron Bridge. The Hugging Face processor expands each image placeholder into the complete media-token sequence before the batch reaches the model. NeMo RL passes that expanded sequence and the image tensors to the model; `NemotronOmniModel` replaces the media-token positions with RADIO encoder outputs and then performs sequence packing and context-parallel sharding.
+
+This is the same model-owned packing boundary used by maintained Megatron VLM integrations. It differs from the historical Nemotron Omni `LLaVAModel` path, which collapsed the expanded media-token sequence before packing and expanded it again inside the model. The dedicated model removes that extra representation change and allows the integration to use Megatron Bridge and Megatron-LM from their maintained main branches.
+
+The current Megatron recipes cover Nano image-and-text GRPO. Super, video, and audio training are follow-up work and are not enabled by these recipes.
+
+### Checkpoint compatibility
+
+Use the `nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16` Hugging Face checkpoint or a checkpoint converted with the dedicated `NemotronOmniModel` integration. Legacy Megatron checkpoints whose parameter names use an `llava_model` prefix are not compatible with this model definition. Reconvert those checkpoints from the original Hugging Face checkpoint instead of loading them directly.
+
+### Maintained recipes
+
+| Workload | Recipe | Topology |
+|---|---|---|
+| CLEVR-CoGenT | [`vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml`](../../examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml) | 1 node, 8 GPUs, TP=8, EP=8 |
+| MMPR-Tiny | [`vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.yaml`](../../examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.yaml) | 4 nodes, 8 GPUs per node, TP=8, EP=16, vLLM TP=2 |
+
+Launch the single-node Megatron recipe from inside the container on an 8-GPU node:
+
+```bash
+uv run examples/run_vlm_grpo.py \
+    --config examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml
+```
+
+For a four-node Slurm run, use the `ray.sub` example above with the following configuration path and omit the AutoModel-specific `PYTHONPATH` addition:
+
+```bash
+CONFIG_PATH=examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.yaml
+```
+
+The recipes keep sequence packing enabled because the model owns the packing step after multimodal embedding insertion. They also request raw generation log probabilities so that vLLM and the Megatron policy compare the same pre-processor probability values when generation constraints such as `bad_words` are active. The generation context cap prevents the processor-expanded image prompt plus generated response from exceeding the configured 8192-token context length.

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

@@ -0,0 +1,55 @@
+defaults: ../../vlm_grpo_3B_megatron.yaml
+loss_fn:
+  reference_policy_kl_penalty: 0.0
+checkpointing:
+  checkpoint_dir: results/vlm_grpo_nemotron_omni_megatron
+policy:
+  model_name: nvidia/Nemotron-3-Nano-Omni-30B-A3B-Reasoning-BF16
+  train_global_batch_size: 8
+  logprob_batch_size: 1
+  max_total_sequence_length: 8192
+  sequence_packing:
+    enabled: true
+  megatron_cfg:
+    env_vars:
+      TORCH_CUDA_ARCH_LIST: '9.0'
+    tensor_model_parallel_size: 8
+    expert_model_parallel_size: 8
+    sequence_parallel: true
+    bias_activation_fusion: false
+    activation_checkpointing: true
+  generation:
+    max_new_tokens: 4096
+    bad_words:
+    - <image>
+    - <img>
+    - </img>
+    - <so_embedding>
+    - <so_start>
+    - <so_end>
+    vllm_cfg:
+      tensor_parallel_size: 8
+      enforce_eager: true
+      max_model_len: 8192
+      cap_max_tokens_to_context: true
+      gpu_memory_utilization: 0.5
+      enable_prefix_caching: false
+      logprobs_mode: raw_logprobs
+    vllm_kwargs:
+      limit_mm_per_prompt:
+        image: 2
+      max_num_batched_tokens: 16384
+      mamba_ssm_cache_dtype: float32
+      skip_mm_profiling: true
+      kernel_config:
+        enable_flashinfer_autotune: false
+        moe_backend: triton
+data:
+  default:
+    prompt_file: examples/prompts/clevr_cogent_cot_nemotron_omni.txt
+logger:
+  wandb:
+    project: grpo-vlm
+    name: nemotron-omni-megatron
+cluster:
+  gpus_per_node: 8

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

@@ -0,0 +1,65 @@
+defaults: ./vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml
+grpo:
+  num_prompts_per_step: 512
+  overlong_filtering: true
+  zero_variance_prompt_filtering: false
+  deduplicate_multimodal_data: false
+loss_fn:
+  ratio_clip_max: 0.28
+  use_on_policy_kl_approximation: true
+  sequence_level_importance_ratios: true
+  token_level_loss: false
+checkpointing:
+  checkpoint_dir: results/vlm_grpo_nemotron_omni_mmpr_megatron
+  keep_top_k: 4
+  checkpoint_must_save_by: 00:03:45:00
+policy:
+  train_global_batch_size: 2048
+  logprob_chunk_size: 1024
+  megatron_cfg:
+    empty_unused_memory_level: 2
+    expert_model_parallel_size: 16
+    optimizer:
+      lr: 3.0e-06
+      min_lr: 2.0e-09
+      weight_decay: 0.0
+      adam_beta2: 0.99
+    scheduler:
+      lr_decay_iters: null
+      lr_warmup_iters: 10
+      lr_warmup_init: 3.0e-08
+  generation:
+    max_new_tokens: ${policy.max_total_sequence_length}
+    vllm_cfg:
+      tensor_parallel_size: 2
+      load_format: auto
+      enforce_eager: false
+      max_model_len: ${policy.max_total_sequence_length}
+      gpu_memory_utilization: 0.75
+    vllm_kwargs:
+      max_num_batched_tokens: 32768
+      max_num_seqs: 512
+data:
+  train:
+    dataset_name: mmpr-tiny
+    download_dir: results/mmpr_tiny_processed
+    split_validation_size: 0.008
+    seed: 42
+  default:
+    prompt_file: null
+    env_name: mmpr-tiny
+env:
+  mmpr-tiny:
+    num_workers: 8
+    reward_functions:
+    - name: geo3k
+      weight: 1.0
+      kwargs:
+        format_score: 0.1
+logger:
+  wandb_enabled: true
+  wandb:
+    project: nemotron-omni-main-migration
+    name: nemotron-omni-mmpr-megatron
+cluster:
+  num_nodes: 4

nemo_rl/data/multimodal_utils.py

@@ -23,6 +23,7 @@
 import decord
 import requests
 import torch
+import torch.nn.functional as F
 from PIL import Image
 from transformers import PreTrainedTokenizerBase
 from transformers.audio_utils import load_audio
@@ -82,6 +83,8 @@ def __init__(
         self,
         tensors: Union[torch.Tensor, list[Optional[torch.Tensor]], list[None]],
         dim_to_pack: int,
+        *,
+        pad_to_max_shape: bool = False,
     ) -> None:
         assert tensors is not None, "Input tensors to PackedTensor cannot be None"
 
@@ -97,6 +100,7 @@ def __init__(
                 f"Unsupported type for input tensors to PackedTensor: {type(tensors)}"
             )
         self.dim_to_pack = dim_to_pack
+        self.pad_to_max_shape = pad_to_max_shape
 
     def as_tensor(
         self, device: Optional[torch.device] = None
@@ -109,8 +113,28 @@ def as_tensor(
         non_none_tensors = [t for t in self.tensors if t is not None]
         if len(non_none_tensors) == 0:
             return None
-        else:
-            return torch.cat(non_none_tensors, dim=self.dim_to_pack).to(device)
+
+        # Dynamic-resolution image processors can produce one raw-pixel batch
+        # per prompt with a different spatial extent, for example
+        # ``[1, 3, 448, 544]`` and ``[1, 3, 544, 448]``. These batches still
+        # pack along the image-count dimension, but they must first share a
+        # common canvas. The accompanying image-size tensor retains the true
+        # per-image dimensions so the model can crop away this padding before
+        # patchification.
+        if (
+            self.pad_to_max_shape
+            and self.dim_to_pack == 0
+            and all(t.ndim == 4 for t in non_none_tensors)
+            and len({t.shape[1] for t in non_none_tensors}) == 1
+        ):
+            max_height = max(t.shape[2] for t in non_none_tensors)
+            max_width = max(t.shape[3] for t in non_none_tensors)
+            non_none_tensors = [
+                F.pad(t, (0, max_width - t.shape[3], 0, max_height - t.shape[2]))
+                for t in non_none_tensors
+            ]
+
+        return torch.cat(non_none_tensors, dim=self.dim_to_pack).to(device)
 
     def __len__(self) -> int:
         # this is the number of tensors in this data wrapper
@@ -125,12 +149,20 @@ def to(self, device: str | torch.device) -> "PackedTensor":
     def slice(self, indices: Union[list[int], torch.Tensor]) -> "PackedTensor":
         idx = indices.tolist() if isinstance(indices, torch.Tensor) else indices
         tensors = [self.tensors[i] for i in idx]
-        return PackedTensor(tensors, self.dim_to_pack)
+        return PackedTensor(
+            tensors,
+            self.dim_to_pack,
+            pad_to_max_shape=self.pad_to_max_shape,
+        )
 
     @classmethod
     def empty_like(cls, other: "PackedTensor") -> "PackedTensor":
         """Return a new PackedTensor with same length and dim_to_pack as `other`, with all entries None."""
-        return cls([None] * len(other.tensors), other.dim_to_pack)
+        return cls(
+            [None] * len(other.tensors),
+            other.dim_to_pack,
+            pad_to_max_shape=other.pad_to_max_shape,
+        )
 
     @classmethod
     def concat(cls, from_packed_tensors: list["PackedTensor"]) -> "PackedTensor":
@@ -158,12 +190,20 @@ def concat(cls, from_packed_tensors: list["PackedTensor"]) -> "PackedTensor":
         assert len(set(dim_to_packs)) == 1, (
             "All packed tensors must have the same dim_to_pack"
         )
+        pad_to_max_shapes = [batch.pad_to_max_shape for batch in from_packed_tensors]
+        assert len(set(pad_to_max_shapes)) == 1, (
+            "All packed tensors must have the same pad_to_max_shape setting"
+        )
         # concatenate the tensors
         tensors = []
         for packed_tensor in from_packed_tensors:
             tensors.extend(packed_tensor.tensors)
         dim_to_pack = dim_to_packs[0]
-        return cls(tensors, dim_to_pack)
+        return cls(
+            tensors,
+            dim_to_pack,
+            pad_to_max_shape=pad_to_max_shapes[0],
+        )
 
     @classmethod
     def flattened_concat(
@@ -195,8 +235,16 @@ def flattened_concat(
         assert len(set(dim_to_packs)) == 1, (
             "All packed tensors must have the same dim_to_pack"
         )
+        pad_to_max_shapes = [batch.pad_to_max_shape for batch in from_packed_tensors]
+        assert len(set(pad_to_max_shapes)) == 1, (
+            "All packed tensors must have the same pad_to_max_shape setting"
+        )
         tensors = [p.as_tensor() for p in from_packed_tensors]
-        return cls(tensors, from_packed_tensors[0].dim_to_pack)
+        return cls(
+            tensors,
+            from_packed_tensors[0].dim_to_pack,
+            pad_to_max_shape=pad_to_max_shapes[0],
+        )
 
 
 def get_multimodal_keys_from_processor(processor) -> list[str]:

nemo_rl/data/processors.py

@@ -614,15 +614,39 @@ def vlm_hf_data_processor(
     user_message["token_ids"] = message["input_ids"][0]
     # add all keys and values to the user message, and the list of keys
     multimodal_keys = list(get_multimodal_keys_from_processor(processor))
-    # imgs_sizes is not declared in model_input_names by the NemotronOmni
-    # checkpoint's bundled image_processor, so append it explicitly when
-    # present. It packs along dim=0 (per-image).
+    # Current Nemotron Omni processors emit imgs_sizes. Historical MMPR
+    # checkpoints instead emit a batch of fixed-size image tiles and only
+    # declare pixel_values. Treat each tile as one dynamic-resolution image so
+    # the model-owned path can patchify it and preserve the processor's exact
+    # placeholder count.
+    if (
+        _uses_image_placeholder
+        and "pixel_values" in message
+        and "imgs_sizes" not in message
+        and message["pixel_values"].ndim == 4
+    ):
+        pixel_values = message["pixel_values"]
+        num_tiles, _, height, width = pixel_values.shape
+        message["imgs_sizes"] = torch.tensor(
+            [[height, width]] * num_tiles, dtype=torch.long
+        )
+
+    # imgs_sizes is not always declared in model_input_names by bundled image
+    # processors, so append it explicitly when present. RADIO uses temporal
+    # patching even for still images and requires one num_frames=1 entry per
+    # image/tile.
     if "imgs_sizes" in message and "imgs_sizes" not in multimodal_keys:
         multimodal_keys.append("imgs_sizes")
+    if "imgs_sizes" in message and "num_frames" not in message:
+        message["num_frames"] = torch.ones(len(message["imgs_sizes"]), dtype=torch.long)
+    if "num_frames" in message and "num_frames" not in multimodal_keys:
+        multimodal_keys.append("num_frames")
     for key in multimodal_keys:
         if key in message:
             user_message[key] = PackedTensor(
-                message[key], dim_to_pack=get_dim_to_pack_along(processor, key)
+                message[key],
+                dim_to_pack=get_dim_to_pack_along(processor, key),
+                pad_to_max_shape=_uses_image_placeholder and key == "pixel_values",
             )
 
     # specifically for gemma, we need to add token_type_ids to the user message as a sequence-type value

nemo_rl/models/generation/interfaces.py

@@ -126,6 +126,7 @@ class GenerationConfig(TypedDict):
     model_name: NotRequired[str]  # Not Required b/c GRPO writes this
     stop_token_ids: list[int] | None
     stop_strings: list[str] | None
+    bad_words: NotRequired[list[str] | None]
     colocated: NotRequired[ColocationConfig]
     port_range_low: NotRequired[int]
     port_range_high: NotRequired[int]

nemo_rl/models/generation/vllm/config.py

@@ -28,6 +28,15 @@ class VllmSpecificArgs(TypedDict):
     async_engine: bool
     load_format: NotRequired[str]
     precision: NotRequired[str]
+    # Whether vLLM returns logprobs before or after generation-time logit
+    # processors. RL policy recomputation uses raw model logits, so recipes
+    # with generation-time processors should request ``raw_logprobs`` when
+    # comparing generation and policy logprobs.
+    logprobs_mode: NotRequired[Literal["processed_logprobs", "raw_logprobs"]]
+    # Cap each request's generated tokens so the training prompt plus response
+    # fits within max_model_len. This is needed when multimodal processing makes
+    # the training prompt longer than its text-only representation.
+    cap_max_tokens_to_context: NotRequired[bool]
     # Use ModelOpt MXFP8 quantization when precision is fp8.
     is_mx: NotRequired[bool]
     kv_cache_dtype: Literal["auto", "fp8", "fp8_e4m3"]