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/distillation_math.yaml

@@ -207,6 +207,7 @@ policy: &POLICY_BASE
             async_engine: false
             precision: ${...precision}
             kv_cache_dtype: "auto"
+            logprobs_mode: processed_logprobs
             tensor_parallel_size: 1
             pipeline_parallel_size: 1
             expert_parallel_size: 1  # When EP > 1, EP must be a multiple of TP since vLLM's EP = DP * TP

examples/configs/grpo_math_1B.yaml

@@ -382,6 +382,7 @@ policy:
       async_engine: false
       precision: ${policy.precision}
       kv_cache_dtype: "auto"
+      logprobs_mode: processed_logprobs
       tensor_parallel_size: 1
       pipeline_parallel_size: 1
       expert_parallel_size: 1  # When EP > 1, EP must be a multiple of TP since vLLM's EP = DP * TP

examples/configs/ppo_math_1B.yaml

@@ -249,6 +249,7 @@ policy:
       async_engine: false
       precision: ${policy.precision}
       kv_cache_dtype: "auto"
+      logprobs_mode: processed_logprobs
       tensor_parallel_size: 1
       pipeline_parallel_size: 1
       expert_parallel_size: 1

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

@@ -46,6 +46,8 @@ policy:
       tensor_parallel_size: 8
       enforce_eager: true
       max_model_len: 8192
+      cap_max_tokens_to_context: true
+      logprobs_mode: raw_logprobs
       gpu_memory_utilization: 0.5
       enable_prefix_caching: false
     vllm_kwargs:

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-automodel-ep8.v1.yaml

@@ -78,6 +78,8 @@ policy:
     vllm_cfg:
       tensor_parallel_size: 8
       enforce_eager: true
+      cap_max_tokens_to_context: true
+      logprobs_mode: raw_logprobs
       gpu_memory_utilization: 0.5
       enable_prefix_caching: false
     vllm_kwargs:

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

@@ -22,6 +22,7 @@
 
 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
@@ -81,7 +82,18 @@ def __init__(
         self,
         tensors: Union[torch.Tensor, list[Optional[torch.Tensor]], list[None]],
         dim_to_pack: int,
+        *,
+        pad_to_max_shape: bool = False,
     ) -> None:
+        """Wrap per-item tensors for concatenation along ``dim_to_pack``.
+
+        Args:
+            tensors: A tensor or list of per-item tensors. List entries may be
+                ``None`` for items without this modality.
+            dim_to_pack: Dimension along which ``as_tensor`` concatenates.
+            pad_to_max_shape: Pad every non-packing dimension to its batch-wide
+                maximum before concatenating. All tensors must have the same rank.
+        """
         assert tensors is not None, "Input tensors to PackedTensor cannot be None"
 
         if isinstance(tensors, torch.Tensor):
@@ -96,6 +108,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
@@ -108,8 +121,51 @@ 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)
+
+        # Some multimodal processors produce a different shape per prompt,
+        # such as dynamic-resolution images, variable-frame videos, or audio
+        # feature sequences. Concatenation already permits the packing
+        # dimension to vary; when explicitly requested, pad every other
+        # dimension to the largest size in the batch.
+        if self.pad_to_max_shape:
+            ranks = {tensor.ndim for tensor in non_none_tensors}
+            if len(ranks) != 1:
+                raise ValueError(
+                    "pad_to_max_shape requires tensors with the same rank, "
+                    f"but received ranks {sorted(ranks)}"
+                )
+
+            rank = ranks.pop()
+            pack_dim = (
+                self.dim_to_pack if self.dim_to_pack >= 0 else rank + self.dim_to_pack
+            )
+            if not 0 <= pack_dim < rank:
+                raise IndexError(
+                    f"dim_to_pack={self.dim_to_pack} is invalid for tensors with rank {rank}"
+                )
+            max_shape = [
+                max(tensor.shape[dim] for tensor in non_none_tensors)
+                for dim in range(rank)
+            ]
+
+            def pad_to_batch_shape(tensor: torch.Tensor) -> torch.Tensor:
+                padding = []
+                for dim in reversed(range(rank)):
+                    padding.extend(
+                        (
+                            0,
+                            0
+                            if dim == pack_dim
+                            else max_shape[dim] - tensor.shape[dim],
+                        )
+                    )
+                return F.pad(tensor, padding)
+
+            non_none_tensors = [
+                pad_to_batch_shape(tensor) for tensor 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
@@ -124,12 +180,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":
@@ -157,12 +221,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(
@@ -194,8 +266,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]: