From e099b82486ec6fcae6e5e90b0ae07468bff2f8db Mon Sep 17 00:00:00 2001 From: Ali Roshan Ghias Date: Wed, 15 Jul 2026 16:59:42 -0700 Subject: [PATCH 1/9] feat(megatron): integrate canonical Nemotron Omni model Signed-off-by: Ali Roshan Ghias --- .../Megatron-Bridge-workspace/Megatron-Bridge | 2 +- nemo_rl/models/megatron/data.py | 33 +- nemo_rl/models/megatron/setup.py | 52 +++ .../policy/workers/megatron_policy_worker.py | 70 +++- .../models/megatron/test_megatron_data.py | 36 ++- .../models/megatron/test_megatron_setup.py | 64 +++- .../megatron/test_nemotron_omni_model.py | 300 ++++++++++++++++++ .../models/policy/test_megatron_worker.py | 65 ++++ 8 files changed, 592 insertions(+), 30 deletions(-) create mode 100644 tests/unit/models/megatron/test_nemotron_omni_model.py diff --git a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge index 554c7b9324..ccf4b09e73 160000 --- a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge +++ b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge @@ -1 +1 @@ -Subproject commit 554c7b9324225aa863eee52e8b8fdde7abced2b1 +Subproject commit ccf4b09e736a317346b277d42cfdccdea4a919c1 diff --git a/nemo_rl/models/megatron/data.py b/nemo_rl/models/megatron/data.py index bb216c8e29..df56fc3147 100644 --- a/nemo_rl/models/megatron/data.py +++ b/nemo_rl/models/megatron/data.py @@ -91,6 +91,7 @@ def make_processed_microbatch_iterator( straggler_timer: StragglerDetector, pad_full_seq_to: Optional[int], delegate_pack_to_model: bool = False, + delegate_mtp_loss_mask_to_model: bool = False, ) -> Iterator[ProcessedMicrobatch]: """Wrap a raw microbatch iterator to yield processed microbatches. @@ -124,6 +125,7 @@ def make_processed_microbatch_iterator( pad_full_seq_to=pad_full_seq_to, pack_sequences=pack_sequences, delegate_pack_to_model=delegate_pack_to_model, + delegate_mtp_loss_mask_to_model=delegate_mtp_loss_mask_to_model, straggler_timer=straggler_timer, ) @@ -148,6 +150,7 @@ def get_microbatch_iterator( straggler_timer: StragglerDetector, seq_length_key: Optional[str] = None, delegate_pack_to_model: bool = False, + delegate_mtp_loss_mask_to_model: bool = False, ) -> Tuple[Iterator[ProcessedMicrobatch], int, int, int, int]: """Create a processed microbatch iterator from a batch of data. @@ -212,6 +215,7 @@ def get_microbatch_iterator( pad_full_seq_to=pad_full_seq_to, straggler_timer=straggler_timer, delegate_pack_to_model=delegate_pack_to_model, + delegate_mtp_loss_mask_to_model=delegate_mtp_loss_mask_to_model, ) # Compute padded sequence length for pipeline parallelism @@ -246,6 +250,7 @@ def process_microbatch( pad_full_seq_to: Optional[int] = None, pack_sequences: bool = False, delegate_pack_to_model: bool = False, + delegate_mtp_loss_mask_to_model: bool = False, straggler_timer: Optional[StragglerDetector] = None, ) -> ProcessedInputs: """Process a microbatch for Megatron model forward pass.""" @@ -286,11 +291,10 @@ def process_microbatch( seq_lengths = data_dict[seq_length_key] if delegate_pack_to_model: - # The VLM packing path does not pack or propagate mtp_loss_mask, - # so MTP training would be silently dropped here. Fail loudly - # instead of producing wrong results. - assert "mtp_loss_mask" not in data_dict, ( - "MTP training is not supported with VLM sequence packing" + has_mtp_loss_mask = "mtp_loss_mask" in data_dict + assert not has_mtp_loss_mask or delegate_mtp_loss_mask_to_model, ( + "MTP training requires a self-packing VLM that advertises " + "model_owns_mtp_loss_mask_packing" ) # VLM path: model (e.g. mbridge Qwen3VL) does its own # preprocess_packed_seqs; NeMo-RL must NOT pre-pack + CP-shard, @@ -320,6 +324,25 @@ def process_microbatch( pad_individual_seqs_to_multiple_of, pad_full_seq_to=pad_full_seq_to, ) + if has_mtp_loss_mask: + source_mtp_loss_mask = data_dict["mtp_loss_mask"] + assert source_mtp_loss_mask.ndim == 2 + assert ( + source_mtp_loss_mask.shape[0] == input_ids_cp_sharded.shape[0] + ) + mtp_loss_mask = source_mtp_loss_mask.new_zeros( + input_ids_cp_sharded.shape + ) + copied_length = min( + source_mtp_loss_mask.shape[1], + input_ids_cp_sharded.shape[1], + ) + mtp_loss_mask[:, :copied_length] = source_mtp_loss_mask[ + :, :copied_length + ] + mtp_loss_mask = mtp_loss_mask * attention_mask.to( + dtype=mtp_loss_mask.dtype + ) position_ids = None else: token_identity = None diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index a88cd33411..4141592781 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -68,6 +68,11 @@ _HF_CONFIG_PATCHED = False +_NEMOTRON_OMNI_ARCHITECTURES = { + "NemotronH_Nano_Omni_Reasoning_V3", + "NemotronH_Super_Omni_Reasoning_V3", +} + def _patch_hf_config_double_instantiation(): """Patch HF config classes whose __post_init__ fails with Megatron's recursive instantiation. @@ -323,6 +328,44 @@ def _get_hf_config_overrides_hash(overrides: dict[str, Any]) -> str: return hashlib.sha256(canonical.encode("utf-8")).hexdigest()[:12] +def _validate_nemotron_omni_megatron_lm_layout( + hf_config: Any, + checkpoint_dir: str, +) -> None: + """Reject legacy LLaVA checkpoints before building the canonical Omni model. + + Raw ``megatron_lm`` checkpoints have no Bridge ``run_config.yaml`` in which + to persist the versioned model contract. Their tensor namespace is therefore + the only reliable way to distinguish the old LLaVA wrapper from the canonical + model-owned-packing implementation. + """ + + architectures = set(getattr(hf_config, "architectures", None) or ()) + if not architectures.intersection(_NEMOTRON_OMNI_ARCHITECTURES): + return + + from megatron.core import dist_checkpointing + + tensor_metadata = dist_checkpointing.load_tensors_metadata(checkpoint_dir) + legacy_keys = [ + str(key) + for key in tensor_metadata + if "llava_model" in str(key).split(".") + ] + if not legacy_keys: + return + + examples = ", ".join(repr(key) for key in legacy_keys[:3]) + raise RuntimeError( + "Cannot load this Nemotron Omni megatron_lm checkpoint as NemotronOmniModel: " + "its tensor namespace contains the legacy 'llava_model' wrapper " + f"({examples}). In this branch NemotronOmniModel uses model-owned packing and a " + "different top-level checkpoint layout. Reconvert the original HF checkpoint with " + "this branch, or use the explicit NemotronOmniLlavaModelProvider/" + "NemotronOmniLlavaBridge with the compatible legacy Megatron-LM pin." + ) + + def _resolve_iter_dir_from_root(path: str, not_found_msg: str) -> str: """Resolve the latest iteration directory under ``path``. @@ -485,6 +528,7 @@ def setup_model_config( hf_cfg = AutoConfig.from_pretrained( hf_model_name, trust_remote_code=True, **hf_config_overrides ) + _validate_nemotron_omni_megatron_lm_layout(hf_cfg, pretrained_path) bridge_obj = AutoBridge.from_hf_config(hf_cfg) model_cfg = bridge_obj.to_megatron_provider(load_weights=False) else: @@ -539,6 +583,14 @@ def setup_model_config( model_cfg = cfg_from_pretrained.model cfg_from_pretrained.logger = LoggerConfig() + # New and legacy Nemotron Omni checkpoints once serialized the same + # provider class name with different semantics. Providers that implement + # this hook must prove their explicit, versioned contract before setup + # applies runtime parallelism or constructs any model layers. + validate_model_contract = getattr(model_cfg, "validate_model_contract", None) + if callable(validate_model_contract): + validate_model_contract() + # Apply parallelism settings _apply_parallelism_config(model_cfg, config) diff --git a/nemo_rl/models/policy/workers/megatron_policy_worker.py b/nemo_rl/models/policy/workers/megatron_policy_worker.py index cba78d3acd..5f615c400a 100644 --- a/nemo_rl/models/policy/workers/megatron_policy_worker.py +++ b/nemo_rl/models/policy/workers/megatron_policy_worker.py @@ -128,16 +128,53 @@ def _model_self_packs_for_cp(model: Any) -> bool: Such models (mbridge VLM wrappers) call ``preprocess_packed_seqs`` in their forward, so NeMo-RL must hand them an unpacked ``[B, S]`` batch instead of - pre-packing + CP-sharding itself. The only such model today is mbridge's - Qwen3VL, which is also the only mbridge VLM that supports context - parallelism; classic mcore GPTModel and other VLMs do not self-pack. + pre-packing + CP-sharding itself. New wrappers advertise the capability + through ``model_owns_packing``. The Qwen3VL type check remains as a + compatibility fallback until that upstream model exposes the capability. """ from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.model import Qwen3VLModel from megatron.core.utils import unwrap_model unwrapped = unwrap_model(model) chunks = unwrapped if isinstance(unwrapped, (list, tuple)) else [unwrapped] - return any(isinstance(chunk, Qwen3VLModel) for chunk in chunks) + return any( + bool(getattr(chunk, "model_owns_packing", False)) + or isinstance(chunk, Qwen3VLModel) + for chunk in chunks + ) + + +def _model_self_packs_mtp_loss_mask(model: Any) -> bool: + """Whether a self-packing model also aligns and CP-shards MTP masks.""" + from megatron.core.utils import unwrap_model + + unwrapped = unwrap_model(model) + chunks = unwrapped if isinstance(unwrapped, (list, tuple)) else [unwrapped] + return any( + bool(getattr(chunk, "model_owns_mtp_loss_mask_packing", False)) + for chunk in chunks + ) + + +def _estimate_refit_tensor_size_in_bytes( + param: torch.Tensor, + *, + export_dtype: torch.dtype, + tp_size: int, + ep_size: int, +) -> int: + """Estimate the gathered tensor size produced by Bridge export. + + Floating-point model weights are exported at the policy dtype. Integral + state (for example BatchNorm ``num_batches_tracked`` buffers) keeps its + original dtype and must not be looked up in a floating-point-only table. + """ + element_size = ( + torch.empty((), dtype=export_dtype).element_size() + if param.is_floating_point() + else param.element_size() + ) + return param.numel() * tp_size * ep_size * element_size # Classes with @ray.remote can't be inherited from, so we split the implementation out. @@ -441,6 +478,12 @@ def __init__( # (mbridge VLM wrappers like Qwen3VL). If so, NeMo-RL must hand it an # unpacked [B, S] batch rather than pre-packing + CP-sharding itself. self.delegate_pack_to_model = _model_self_packs_for_cp(self.model) + self.delegate_mtp_loss_mask_to_model = _model_self_packs_mtp_loss_mask( + self.model + ) + assert ( + not self.delegate_mtp_loss_mask_to_model or self.delegate_pack_to_model + ), "A model cannot own MTP-mask packing without owning sequence packing" # vars used for refit ## will be initialized in prepare_refit_info @@ -641,6 +684,7 @@ def train( mbs, straggler_timer=self.mcore_state.straggler_timer, delegate_pack_to_model=self.delegate_pack_to_model, + delegate_mtp_loss_mask_to_model=self.delegate_mtp_loss_mask_to_model, ) # Track total microbatches for MoE aux-loss averaging total_num_microbatches += int(num_microbatches) @@ -1456,6 +1500,7 @@ def get_logprobs( logprob_batch_size, straggler_timer=self.mcore_state.straggler_timer, delegate_pack_to_model=self.delegate_pack_to_model, + delegate_mtp_loss_mask_to_model=self.delegate_mtp_loss_mask_to_model, ) use_fused_linear_logprobs = self.cfg["megatron_cfg"].get( @@ -1673,6 +1718,7 @@ def get_topk_logits( logprob_batch_size, straggler_timer=self.mcore_state.straggler_timer, delegate_pack_to_model=self.delegate_pack_to_model, + delegate_mtp_loss_mask_to_model=self.delegate_mtp_loss_mask_to_model, ) list_of_outputs = megatron_forward_backward( @@ -1817,17 +1863,11 @@ def calculate_size_in_bytes(param, tp_size, ep_size): # need to broadcast for other pp ranks size_in_bytes = None else: - # Calculate size for this parameter - prec_to_bytes = { - torch.bfloat16: 2, - torch.float16: 2, - torch.float32: 4, - torch.float8_e4m3fn: 1, - torch.float8_e5m2: 1, - } - scale = prec_to_bytes[self.dtype] / prec_to_bytes[param.dtype] - size_in_bytes = ( - param.element_size() * param.numel() * tp_size * ep_size * scale + size_in_bytes = _estimate_refit_tensor_size_in_bytes( + param, + export_dtype=self.dtype, + tp_size=tp_size, + ep_size=ep_size, ) # Broadcast size_in_bytes across pipeline parallel ranks diff --git a/tests/unit/models/megatron/test_megatron_data.py b/tests/unit/models/megatron/test_megatron_data.py index e72431ba11..c74b0a149b 100644 --- a/tests/unit/models/megatron/test_megatron_data.py +++ b/tests/unit/models/megatron/test_megatron_data.py @@ -566,13 +566,10 @@ def test_process_microbatch_delegate_pack_to_model(self, mock_prepare, mock_pack assert torch.equal(result.cu_seqlens_padded, mock_cu_seqlens_padded) def test_process_microbatch_delegate_pack_rejects_mtp_loss_mask(self): - """delegate_pack_to_model must reject a pre-computed mtp_loss_mask. + """Self-packing models must explicitly advertise MTP-mask ownership. - The VLM self-packing path does not pack/propagate mtp_loss_mask, so MTP - training would be silently dropped. process_microbatch must fail loudly - rather than produce wrong results. Regression guard for issue #2869: the - worker now only creates mtp_loss_mask when MTP is enabled, but if a mask - ever reaches this path it must raise instead of being silently ignored. + Qwen3-VL and other wrappers that have not implemented this contract stay + fail-closed rather than receiving a full-batch mask for CP-sharded tokens. """ from nemo_rl.models.megatron.data import process_microbatch @@ -596,8 +593,31 @@ def test_process_microbatch_delegate_pack_rejects_mtp_loss_mask(self): straggler_timer=MagicMock(), ) - assert "MTP training is not supported with VLM sequence packing" in str( - exc_info.value + assert "model_owns_mtp_loss_mask_packing" in str(exc_info.value) + + def test_process_microbatch_delegates_padded_mtp_loss_mask(self): + """A capable wrapper receives a padded full mask to pack with its IDs.""" + from nemo_rl.models.megatron.data import process_microbatch + + input_ids = torch.tensor([[1, 2, 3, 0, 0], [4, 5, 0, 0, 0]]) + mtp_loss_mask = torch.tensor([[0, 0, 1, 0, 0], [0, 1, 0, 0, 0]]) + result = process_microbatch( + { + "input_ids": input_ids, + "input_lengths": torch.tensor([3, 2]), + "mtp_loss_mask": mtp_loss_mask, + }, + seq_length_key="input_lengths", + pad_individual_seqs_to_multiple_of=4, + pack_sequences=True, + delegate_pack_to_model=True, + delegate_mtp_loss_mask_to_model=True, + ) + + assert result.input_ids_cp_sharded.shape == (2, 4) + assert torch.equal( + result.mtp_loss_mask, + torch.tensor([[0, 0, 1, 0], [0, 1, 0, 0]]), ) diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py index 02c634e3fe..638e128798 100644 --- a/tests/unit/models/megatron/test_megatron_setup.py +++ b/tests/unit/models/megatron/test_megatron_setup.py @@ -1,4 +1,4 @@ -# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved. +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -32,6 +32,68 @@ import torch +@pytest.mark.mcore +class TestNemotronOmniCheckpointContract: + """Tests for distinguishing raw Option 2 and legacy LLaVA checkpoints.""" + + def test_non_omni_checkpoint_does_not_read_tensor_metadata(self): + from nemo_rl.models.megatron.setup import ( + _validate_nemotron_omni_megatron_lm_layout, + ) + + hf_config = SimpleNamespace(architectures=["LlamaForCausalLM"]) + with patch( + "megatron.core.dist_checkpointing.load_tensors_metadata" + ) as load_metadata: + _validate_nemotron_omni_megatron_lm_layout(hf_config, "/unused") + + load_metadata.assert_not_called() + + def test_canonical_omni_checkpoint_layout_is_accepted(self): + from nemo_rl.models.megatron.setup import ( + _validate_nemotron_omni_megatron_lm_layout, + ) + + hf_config = SimpleNamespace( + architectures=["NemotronH_Nano_Omni_Reasoning_V3"] + ) + metadata = { + "language_model.embedding.word_embeddings.weight": object(), + "vision_model.decoder.layers.0.self_attention.linear_qkv.weight": object(), + } + with patch( + "megatron.core.dist_checkpointing.load_tensors_metadata", + return_value=metadata, + ): + _validate_nemotron_omni_megatron_lm_layout( + hf_config, "/checkpoints/iter_0000001" + ) + + def test_legacy_llava_checkpoint_layout_is_rejected(self): + from nemo_rl.models.megatron.setup import ( + _validate_nemotron_omni_megatron_lm_layout, + ) + + hf_config = SimpleNamespace( + architectures=["NemotronH_Super_Omni_Reasoning_V3"] + ) + metadata = { + "model.llava_model.language_model.embedding.word_embeddings.weight": object(), + "model.llava_model.vision_model.decoder.layers.0.self_attention.linear_qkv.weight": object(), + } + with patch( + "megatron.core.dist_checkpointing.load_tensors_metadata", + return_value=metadata, + ): + with pytest.raises( + RuntimeError, + match="tensor namespace contains the legacy 'llava_model' wrapper", + ): + _validate_nemotron_omni_megatron_lm_layout( + hf_config, "/checkpoints/iter_0002821" + ) + + @pytest.mark.mcore class TestValidateModelPaths: """Tests for validate_model_paths function.""" diff --git a/tests/unit/models/megatron/test_nemotron_omni_model.py b/tests/unit/models/megatron/test_nemotron_omni_model.py new file mode 100644 index 0000000000..e09b0ae328 --- /dev/null +++ b/tests/unit/models/megatron/test_nemotron_omni_model.py @@ -0,0 +1,300 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Distributed functional coverage for NeMo-RL's Nemotron Omni contract.""" + +import copy +import functools +import gc +import os +from dataclasses import dataclass + +import pytest +import torch +from megatron.core import dist_checkpointing, parallel_state +from megatron.core.distributed import DistributedDataParallelConfig +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed + +from megatron.bridge.models.nemotron_omni.nemotron_omni_provider import ( + NEMOTRON_OMNI_EXPANDED_SEQUENCE_CONTRACT, + NemotronOmniModelProvider, +) +from nemo_rl.distributed.model_utils import ( + from_parallel_logits_to_logprobs_packed_sequences, +) +from nemo_rl.models.megatron.data import process_microbatch + + +pytestmark = pytest.mark.mcore + +_IMAGE_TOKEN_ID = 18 + + +@dataclass +class _TinyOmniProvider(NemotronOmniModelProvider): + """Small real RADIO/NemotronH model for a two-rank functional test.""" + + has_sound: bool = False + language_model_type: str = "nemotron6-moe" + hidden_size: int = 128 + ffn_hidden_size: int = 256 + num_attention_heads: int = 4 + num_query_groups: int = 2 + kv_channels: int = 32 + mamba_num_heads: int = 4 + mamba_head_dim: int = 32 + mamba_num_groups: int = 2 + mamba_state_dim: int = 16 + hybrid_layer_pattern: str = "M" + vocab_size: int = 128 + seq_length: int = 32 + image_token_index: int = _IMAGE_TOKEN_ID + img_start_token_id: int = 21 + img_end_token_id: int = 22 + tokenizer_type: str = "nemotron6-moe" + dynamic_resolution: bool = True + use_vision_backbone_fp8_arch: bool = False + vision_proj_ffn_hidden_size: int = 256 + pipeline_model_parallel_size: int = 1 + use_cpu_initialization: bool = True + gradient_accumulation_fusion: bool = False + nemotron_omni_contract: str = NEMOTRON_OMNI_EXPANDED_SEQUENCE_CONTRACT + + def _build_vision_config(self, language_cfg): + vision_cfg = copy.deepcopy(language_cfg) + vision_cfg.sequence_parallel = False + vision_cfg.context_parallel_size = 1 + vision_cfg.tp_comm_overlap = False + vision_cfg.recompute_granularity = None + vision_cfg.recompute_method = None + vision_cfg.recompute_num_layers = None + vision_cfg.mtp_num_layers = None + vision_cfg.num_layers = 1 + vision_cfg.pipeline_model_parallel_size = 1 + vision_cfg.num_attention_heads = 4 + vision_cfg.add_bias_linear = True + vision_cfg.add_qkv_bias = True + vision_cfg.hidden_size = 128 + vision_cfg.ffn_hidden_size = 256 + vision_cfg.gated_linear_unit = False + vision_cfg.kv_channels = 32 + vision_cfg.num_query_groups = 4 + vision_cfg.normalization = "LayerNorm" + vision_cfg.qk_layernorm = False + vision_cfg.layernorm_epsilon = 1e-6 + vision_cfg.class_token_len = 10 + return vision_cfg + + +def _build_distributed_model(): + if parallel_state.model_parallel_is_initialized(): + parallel_state.destroy_model_parallel() + parallel_state.initialize_model_parallel( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + context_parallel_size=2, + ) + torch.manual_seed(123) + model_parallel_cuda_manual_seed(123) + + provider = _TinyOmniProvider( + freeze_language_model=True, + tensor_model_parallel_size=1, + context_parallel_size=2, + sequence_parallel=False, + ) + provider.finalize() + models = provider.provide_distributed_model( + ddp_config=DistributedDataParallelConfig( + grad_reduce_in_fp32=True, + overlap_grad_reduce=False, + use_distributed_optimizer=False, + check_for_nan_in_grad=True, + ), + wrap_with_ddp=True, + mixed_precision_wrapper=None, + ) + assert len(models) == 1 + return models[0] + + +def _expanded_fixture(device: torch.device): + input_ids = torch.tensor( + [ + [7, 21, 18, 18, 22, 9, 10, 0], + [11, 21, 18, 22, 12, 0, 0, 0], + ], + dtype=torch.long, + device=device, + ) + lengths = torch.tensor([7, 5], dtype=torch.long, device=device) + generator = torch.Generator(device=device) + generator.manual_seed(2026) + images = torch.randn(2, 3, 32, 64, generator=generator, device=device) + images[1, :, :, 32:] = 0 + image_sizes = torch.tensor( + [[32, 64], [32, 32]], dtype=torch.int32, device=device + ) + return input_ids, lengths, images, image_sizes + + +def _forward(model): + device = torch.device("cuda", torch.cuda.current_device()) + input_ids, lengths, images, image_sizes = _expanded_fixture(device) + processed = process_microbatch( + {"input_ids": input_ids, "input_lengths": lengths}, + seq_length_key="input_lengths", + pad_individual_seqs_to_multiple_of=4, + pack_sequences=True, + delegate_pack_to_model=True, + ) + output = model( + input_ids=processed.input_ids_cp_sharded, + attention_mask=processed.attention_mask, + packed_seq_params=processed.packed_seq_params, + pixel_values=images, + imgs_sizes=image_sizes, + ) + logprobs = from_parallel_logits_to_logprobs_packed_sequences( + output, + target=processed.input_ids, + cu_seqlens_padded=processed.cu_seqlens_padded, + unpacked_seqlen=input_ids.shape[1], + vocab_start_index=0, + vocab_end_index=output.shape[-1], + group=parallel_state.get_tensor_model_parallel_group(), + inference_only=False, + cp_group=parallel_state.get_context_parallel_group(), + ) + prediction_mask = torch.arange( + input_ids.shape[1] - 1, device=device + ).unsqueeze(0) < (lengths - 1).unsqueeze(1) + loss = -(logprobs * prediction_mask).sum() / prediction_mask.sum() + return loss, output + + +def _run_training_checkpoint_roundtrip( + rank: int, + world_size: int, + *, + checkpoint_dir: str, +) -> None: + assert world_size == 2 + model = _build_distributed_model() + model.train() + model.zero_grad_buffer() + + loss, output = _forward(model) + loss.backward() + model.finish_grad_sync() + + core_model = model.module + gradients = {} + before_update = {} + optimizer_parameters = [] + for name, parameter in core_model.named_parameters(): + if not parameter.requires_grad: + continue + assert name.startswith(("vision_model.", "vision_projection.")) + assert hasattr(parameter, "main_grad") + assert torch.isfinite(parameter.main_grad).all() + rank_zero_gradient = parameter.main_grad.detach().clone() + torch.distributed.broadcast(rank_zero_gradient, src=0) + torch.testing.assert_close( + parameter.main_grad, rank_zero_gradient, rtol=0, atol=0 + ) + gradients[name] = parameter.main_grad + before_update[name] = parameter.detach().clone() + parameter.grad = parameter.main_grad.to(parameter.dtype).clone() + optimizer_parameters.append(parameter) + assert gradients + + optimizer = torch.optim.SGD(optimizer_parameters, lr=1.0) + optimizer.step() + changed = { + name + for name, parameter in core_model.named_parameters() + if name in before_update and not torch.equal(parameter, before_update[name]) + } + assert any(name.startswith("vision_model.") for name in changed) + assert any(name.startswith("vision_projection.") for name in changed) + + model.eval() + with torch.no_grad(): + _, post_update_output = _forward(model) + post_update_output = post_update_output.detach().clone() + + metadata = { + "dp_cp_group": parallel_state.get_data_parallel_group( + with_context_parallel=True + ) + } + sharded_state = core_model.sharded_state_dict(metadata=metadata) + assert changed <= sharded_state.keys() + if rank == 0: + os.makedirs(checkpoint_dir, exist_ok=True) + torch.distributed.barrier() + dist_checkpointing.save({"model": sharded_state}, checkpoint_dir) + + provider = _TinyOmniProvider( + freeze_language_model=True, + tensor_model_parallel_size=1, + context_parallel_size=2, + sequence_parallel=False, + ) + provider.finalize() + restored_model = provider.provide().cuda().eval() + restore_template = restored_model.sharded_state_dict(metadata=metadata) + loaded_state = dist_checkpointing.load( + {"model": restore_template}, checkpoint_dir + ) + incompatible = restored_model.load_state_dict(loaded_state["model"]) + assert not incompatible.missing_keys + assert not incompatible.unexpected_keys + + restored_parameters = dict(restored_model.named_parameters()) + original_parameters = dict(core_model.named_parameters()) + for name in changed: + torch.testing.assert_close( + restored_parameters[name], original_parameters[name], rtol=0, atol=0 + ) + with torch.no_grad(): + _, restored_output = _forward(restored_model) + torch.testing.assert_close(restored_output, post_update_output, rtol=0, atol=0) + + if rank == 0: + print( + "NEMOTRON_OMNI_CP2_DCP_ROUNDTRIP " + f"loss={loss.item():.8f} changed_tensors={len(changed)} " + "post_restore_max_logit_abs_diff=0.00000000", + flush=True, + ) + + del output, post_update_output, restored_output, optimizer, model + del core_model, restored_model + gc.collect() + torch.cuda.empty_cache() + torch.distributed.barrier() + parallel_state.destroy_model_parallel() + + +def test_nemotron_omni_cp2_training_and_checkpoint_roundtrip( + distributed_test_runner, + tmp_path, +): + test_fn = functools.partial( + _run_training_checkpoint_roundtrip, + checkpoint_dir=str(tmp_path / "nemotron_omni_cp2_dcp"), + ) + distributed_test_runner(test_fn, world_size=2) diff --git a/tests/unit/models/policy/test_megatron_worker.py b/tests/unit/models/policy/test_megatron_worker.py index 793cc52bd9..694b5d43e7 100644 --- a/tests/unit/models/policy/test_megatron_worker.py +++ b/tests/unit/models/policy/test_megatron_worker.py @@ -45,6 +45,71 @@ pytestmark = pytest.mark.mcore +def test_model_owned_packing_capability_is_detected(): + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + _model_self_packs_for_cp, + ) + + class ModelOwnedPackingModel: + model_owns_packing = True + + assert _model_self_packs_for_cp(ModelOwnedPackingModel()) + + +def test_model_owned_mtp_loss_mask_packing_capability_is_detected(): + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + _model_self_packs_mtp_loss_mask, + ) + + class ModelOwnedPackingModel: + model_owns_mtp_loss_mask_packing = True + + assert _model_self_packs_mtp_loss_mask(ModelOwnedPackingModel()) + assert not _model_self_packs_mtp_loss_mask(object()) + + +def test_regular_model_does_not_delegate_packing(): + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + _model_self_packs_for_cp, + ) + + assert not _model_self_packs_for_cp(object()) + + +def test_refit_size_estimate_preserves_integral_buffer_dtype(): + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + _estimate_refit_tensor_size_in_bytes, + ) + + param = torch.zeros(3, dtype=torch.int64) + + assert _estimate_refit_tensor_size_in_bytes( + param, export_dtype=torch.bfloat16, tp_size=2, ep_size=4 + ) == 3 * 8 * 2 * 4 + + +def test_refit_size_estimate_casts_floating_weight_to_export_dtype(): + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + _estimate_refit_tensor_size_in_bytes, + ) + + param = torch.zeros(3, dtype=torch.float32) + + assert _estimate_refit_tensor_size_in_bytes( + param, export_dtype=torch.bfloat16, tp_size=2, ep_size=4 + ) == 3 * 2 * 2 * 4 + + +def test_qwen3vl_type_fallback_still_delegates_packing(): + from megatron.bridge.models.qwen_vl.modelling_qwen3_vl.model import Qwen3VLModel + + from nemo_rl.models.policy.workers.megatron_policy_worker import ( + _model_self_packs_for_cp, + ) + + assert _model_self_packs_for_cp(Qwen3VLModel.__new__(Qwen3VLModel)) + + class _FakeTrainableModel: def __init__(self): self.train_called = False From 42e4b350a23699e9a3f8207e20a33a288ecf02f2 Mon Sep 17 00:00:00 2001 From: Ali Roshan Ghias Date: Wed, 15 Jul 2026 17:00:57 -0700 Subject: [PATCH 2/9] feat(multimodal): support Nemotron Omni image batches Signed-off-by: Ali Roshan Ghias --- nemo_rl/data/multimodal_utils.py | 65 ++++++++++++++++++++-- nemo_rl/data/processors.py | 34 +++++++++-- tests/unit/data/datasets/test_mmpr_tiny.py | 44 ++++++++++++++- tests/unit/data/test_multimodal_dict.py | 16 ++++++ 4 files changed, 146 insertions(+), 13 deletions(-) diff --git a/nemo_rl/data/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py index 6a16e18fec..7d029af578 100644 --- a/nemo_rl/data/multimodal_utils.py +++ b/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,29 @@ 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 +150,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 +191,22 @@ 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 +238,18 @@ 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]: diff --git a/nemo_rl/data/processors.py b/nemo_rl/data/processors.py index 734ab6cf0e..eae85a6b78 100644 --- a/nemo_rl/data/processors.py +++ b/nemo_rl/data/processors.py @@ -614,15 +614,41 @@ 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 diff --git a/tests/unit/data/datasets/test_mmpr_tiny.py b/tests/unit/data/datasets/test_mmpr_tiny.py index 8f1dce267c..a2c61d5dee 100644 --- a/tests/unit/data/datasets/test_mmpr_tiny.py +++ b/tests/unit/data/datasets/test_mmpr_tiny.py @@ -128,7 +128,7 @@ def test_missing_download_dir_raises_value_error(self): MMPRTinyDataset(download_dir="") -def _make_stub_nemotron_processor(): +def _make_stub_nemotron_processor(*, include_imgs_sizes=True, num_tiles=1): """Build a minimal stub whose class name is NemotronNanoVLV2Processor. The stub implements just enough of the AutoProcessor interface for @@ -164,10 +164,13 @@ def apply_chat_template(self, messages, **kwargs): def __call__(self, text=None, images=None, **kwargs): self.captured_call_text = text - return { + result = { "input_ids": fake_input_ids, - "pixel_values": torch.randn(1, 3, 224, 224), + "pixel_values": torch.randn(num_tiles, 3, 224, 224), } + if include_imgs_sizes: + result["imgs_sizes"] = torch.tensor([[224, 224]] * num_tiles) + return result return NemotronNanoVLV2Processor() @@ -233,6 +236,41 @@ def test_processor_produces_valid_datum_spec(self, tiny_image_path): assert "vllm_images" in result assert len(result["vllm_images"]) == 1 assert result["task_name"] == "mmpr-tiny" + user_message = result["message_log"][0] + assert torch.equal( + user_message["num_frames"].as_tensor(), torch.tensor([1]) + ) + + def test_historical_tiled_processor_gets_media_metadata(self, tiny_image_path): + from nemo_rl.data.interfaces import TaskDataSpec + from nemo_rl.data.processors import vlm_hf_data_processor + + task_data_spec = TaskDataSpec(task_name="mmpr-tiny") + task_data_spec.prompt = _TEST_PROMPT_TEMPLATE + processor = _make_stub_nemotron_processor( + include_imgs_sizes=False, num_tiles=3 + ) + result = vlm_hf_data_processor( + datum_dict={ + "images": [tiny_image_path], + "question": _RAW_QUESTION, + "answer": "A", + "task_name": "mmpr-tiny", + }, + task_data_spec=task_data_spec, + processor=processor, + max_seq_length=8192, + idx=0, + ) + + user_message = result["message_log"][0] + assert torch.equal( + user_message["imgs_sizes"].as_tensor(), + torch.tensor([[224, 224], [224, 224], [224, 224]]), + ) + assert torch.equal( + user_message["num_frames"].as_tensor(), torch.ones(3, dtype=torch.long) + ) def test_prompted_text_contains_boxed_literal_and_no_raw_dataset_string( self, tiny_image_path diff --git a/tests/unit/data/test_multimodal_dict.py b/tests/unit/data/test_multimodal_dict.py index a94412222a..ea56ac79c9 100644 --- a/tests/unit/data/test_multimodal_dict.py +++ b/tests/unit/data/test_multimodal_dict.py @@ -349,3 +349,19 @@ def test_packedtensor_as_tensor_with_mixed_none_and_tensors(): out = pt.as_tensor() expected = torch.cat([t1, t3], dim=0) assert torch.equal(out, expected) + + +def test_packedtensor_pads_mixed_dynamic_resolution_images(): + """Raw image batches pad spatial dimensions before packing on dim 0.""" + first = torch.ones(1, 3, 2, 4) + second = 2 * torch.ones(1, 3, 4, 2) + + packed = PackedTensor( + [first, second], dim_to_pack=0, pad_to_max_shape=True + ).as_tensor() + + assert packed.shape == (2, 3, 4, 4) + torch.testing.assert_close(packed[0, :, :2, :4], first[0]) + torch.testing.assert_close(packed[0, :, 2:, :], torch.zeros(3, 2, 4)) + torch.testing.assert_close(packed[1, :, :4, :2], second[0]) + torch.testing.assert_close(packed[1, :, :, 2:], torch.zeros(3, 4, 2)) From 393fe23d6f74f079938233a46c18ca5774a96463 Mon Sep 17 00:00:00 2001 From: Ali Roshan Ghias Date: Wed, 15 Jul 2026 17:01:07 -0700 Subject: [PATCH 3/9] feat(vllm): support Nemotron Omni generation parity Signed-off-by: Ali Roshan Ghias --- nemo_rl/models/generation/interfaces.py | 1 + nemo_rl/models/generation/vllm/config.py | 9 + nemo_rl/models/generation/vllm/patches.py | 162 ++++++++++++++++++ .../models/generation/vllm/vllm_backend.py | 72 ++++++++ nemo_rl/models/generation/vllm/vllm_worker.py | 56 +++++- .../models/generation/test_vllm_backend.py | 79 +++++++++ .../models/generation/test_vllm_generation.py | 48 ++++++ .../models/generation/test_vllm_patches.py | 116 +++++++++++++ 8 files changed, 534 insertions(+), 9 deletions(-) create mode 100644 tests/unit/models/generation/test_vllm_patches.py diff --git a/nemo_rl/models/generation/interfaces.py b/nemo_rl/models/generation/interfaces.py index 025707c700..57ae6e6928 100644 --- a/nemo_rl/models/generation/interfaces.py +++ b/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] diff --git a/nemo_rl/models/generation/vllm/config.py b/nemo_rl/models/generation/vllm/config.py index 44eb5d5d89..a38ee75ac1 100644 --- a/nemo_rl/models/generation/vllm/config.py +++ b/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"] diff --git a/nemo_rl/models/generation/vllm/patches.py b/nemo_rl/models/generation/vllm/patches.py index 1d7fbfbdbb..61e896f162 100644 --- a/nemo_rl/models/generation/vllm/patches.py +++ b/nemo_rl/models/generation/vllm/patches.py @@ -282,6 +282,167 @@ def _patch_vllm_hermes_tool_parser_thread_safety(logger) -> None: logger.info("Successfully patched hermes tool parser for thread-safety.") +def _patch_vllm_bad_words_tokenization_cache(logger) -> None: + """Cache vLLM ``bad_words`` tokenization under a process-wide lock. + + vLLM 0.20.0 tokenizes the same string list for every request. Besides the + avoidable cost, concurrent calls into a Hugging Face fast tokenizer can + fail with ``RuntimeError: Already borrowed``. The Omni vLLM fork carries + this cache; apply the same narrow behavior to the supported stock wheel. + + The patch is deliberately fail-closed on source shape. If vLLM changes the + method, leave the installation untouched and require the new version to be + audited instead of applying a partial textual rewrite. + """ + try: + file_to_patch = _get_vllm_file("sampling_params.py") + except RuntimeError: + logger.warning("Could not locate sampling_params.py for bad_words patch.") + return + + old_import = "import json as json_mod\nfrom dataclasses import field" + new_import = ( + "import json as json_mod\nimport threading\nfrom dataclasses import field" + ) + constants_anchor = "_SAMPLING_EPS = 1e-5\n_MAX_TEMP = 1e-2\n" + cache_definition = ( + "_SAMPLING_EPS = 1e-5\n" + "_MAX_TEMP = 1e-2\n\n" + "# Cache tokenized bad_words across requests. Fast tokenizers are not\n" + "# safe for concurrent encode calls on the same tokenizer instance.\n" + "_BAD_WORDS_TOKEN_IDS_CACHE: dict[\n" + " tuple[int, tuple[str, ...]], list[list[int]]\n" + "] = {}\n" + "_BAD_WORDS_TOKEN_IDS_CACHE_LOCK = threading.Lock()\n" + "_BAD_WORDS_TOKEN_IDS_CACHE_MAX_ENTRIES = 1024\n" + ) + old_method = """ def update_from_tokenizer(self, tokenizer: TokenizerLike) -> None: + if not self.bad_words: + return + self._bad_words_token_ids = [] + for bad_word in self.bad_words: + # To prohibit words both at the beginning + # and in the middle of text + # (related to add_prefix_space tokenizer parameter) + for add_prefix_space in [False, True]: + prefix = " " if add_prefix_space else "" + prompt = prefix + bad_word.lstrip() + prompt_token_ids = tokenizer.encode( + text=prompt, add_special_tokens=False + ) + + # If no space at the beginning + # or if prefix space produces a new word token + if (not add_prefix_space) or ( + add_prefix_space + and prompt_token_ids[0] != self._bad_words_token_ids[-1][0] + and len(prompt_token_ids) == len(self._bad_words_token_ids[-1]) + ): + self._bad_words_token_ids.append(prompt_token_ids) + + invalid_token_ids = [ + token_id + for bad_words_token_ids in self._bad_words_token_ids + for token_id in bad_words_token_ids + if token_id < 0 or token_id > tokenizer.max_token_id + ] + if len(invalid_token_ids) > 0: + raise VLLMValidationError( + f"The model vocabulary size is {tokenizer.max_token_id + 1}," + f" but the following tokens" + f" were specified as bad: {invalid_token_ids}." + f" All token id values should be integers satisfying:" + f" 0 <= token_id <= {tokenizer.max_token_id}.", + parameter="bad_words", + value=self.bad_words, + ) +""" + new_method = """ def update_from_tokenizer(self, tokenizer: TokenizerLike) -> None: + if not self.bad_words: + return + cache_key = (id(tokenizer), tuple(self.bad_words)) + cached = _BAD_WORDS_TOKEN_IDS_CACHE.get(cache_key) + if cached is not None: + self._bad_words_token_ids = cached + return + with _BAD_WORDS_TOKEN_IDS_CACHE_LOCK: + cached = _BAD_WORDS_TOKEN_IDS_CACHE.get(cache_key) + if cached is not None: + self._bad_words_token_ids = cached + return + self._bad_words_token_ids = self._tokenize_bad_words(tokenizer) + if ( + len(_BAD_WORDS_TOKEN_IDS_CACHE) + >= _BAD_WORDS_TOKEN_IDS_CACHE_MAX_ENTRIES + ): + _BAD_WORDS_TOKEN_IDS_CACHE.clear() + _BAD_WORDS_TOKEN_IDS_CACHE[cache_key] = self._bad_words_token_ids + + def _tokenize_bad_words(self, tokenizer: TokenizerLike) -> list[list[int]]: + bad_words_token_ids: list[list[int]] = [] + for bad_word in self.bad_words: + # To prohibit words both at the beginning + # and in the middle of text + # (related to add_prefix_space tokenizer parameter) + for add_prefix_space in [False, True]: + prefix = " " if add_prefix_space else "" + prompt = prefix + bad_word.lstrip() + prompt_token_ids = tokenizer.encode( + text=prompt, add_special_tokens=False + ) + + # If no space at the beginning + # or if prefix space produces a new word token + if (not add_prefix_space) or ( + add_prefix_space + and prompt_token_ids[0] != bad_words_token_ids[-1][0] + and len(prompt_token_ids) == len(bad_words_token_ids[-1]) + ): + bad_words_token_ids.append(prompt_token_ids) + + invalid_token_ids = [ + token_id + for token_ids in bad_words_token_ids + for token_id in token_ids + if token_id < 0 or token_id > tokenizer.max_token_id + ] + if len(invalid_token_ids) > 0: + raise VLLMValidationError( + f"The model vocabulary size is {tokenizer.max_token_id + 1}," + f" but the following tokens" + f" were specified as bad: {invalid_token_ids}." + f" All token id values should be integers satisfying:" + f" 0 <= token_id <= {tokenizer.max_token_id}.", + parameter="bad_words", + value=self.bad_words, + ) + return bad_words_token_ids +""" + + with _locked_file_patch(file_to_patch) as (content, write_back): + if "_BAD_WORDS_TOKEN_IDS_CACHE_LOCK" in content: + logger.info("vLLM bad_words tokenization cache already applied.") + return + if ( + old_import not in content + or constants_anchor not in content + or old_method not in content + ): + logger.warning( + "Could not apply vLLM bad_words tokenization cache: expected " + "vLLM 0.20.0 source shape was not found in %s.", + file_to_patch, + ) + return + + content = content.replace(old_import, new_import, 1) + content = content.replace(constants_anchor, cache_definition, 1) + content = content.replace(old_method, new_method, 1) + write_back(content) + + logger.info("Successfully patched vLLM bad_words tokenization cache.") + + def _apply_vllm_patches( py_executable: str, *, extra_env_vars: list[str] | None = None ) -> None: @@ -295,3 +456,4 @@ def _apply_vllm_patches( _patch_vllm_llama_eagle3_own_lm_head(patch_logger) _patch_vllm_hermes_tool_parser_thread_safety(patch_logger) + _patch_vllm_bad_words_tokenization_cache(patch_logger) diff --git a/nemo_rl/models/generation/vllm/vllm_backend.py b/nemo_rl/models/generation/vllm/vllm_backend.py index f6b2e5adb9..de9d172b9d 100644 --- a/nemo_rl/models/generation/vllm/vllm_backend.py +++ b/nemo_rl/models/generation/vllm/vllm_backend.py @@ -165,6 +165,78 @@ def prepare_refit_info(self, state_dict_info: dict[str, Any]) -> None: e.g. {tensor_name: (shape, dtype)} """ self.state_dict_info = state_dict_info # pyrefly: ignore[implicitly-defined-attribute] This class does not define __init__ so assignments like this should be ignored + self._validate_nemotron_omni_radio_layerscale_refit(state_dict_info) + + def _is_nemotron_omni(self) -> bool: + architectures = set( + self.model_runner.vllm_config.model_config.architectures or [] + ) + return not architectures.isdisjoint( + { + "NemotronH_Nano_Omni_Reasoning_V3", + "NemotronH_Super_Omni_Reasoning_V3", + } + ) + + def _validate_nemotron_omni_radio_layerscale_refit( + self, state_dict_info: dict[str, Any] + ) -> None: + """Reject explicit LayerScale state that stock vLLM would ignore. + + This method can run while a colocated vLLM engine is asleep. It must not + read or mutate model tensors because vLLM's level-1 sleep allocator has + temporarily released their CUDA storage. + """ + if not self._is_nemotron_omni(): + return + + explicit_layerscale = [ + name + for name in state_dict_info + if name.startswith("vision_model.radio_model.model.blocks.") + and name.endswith((".ls1", ".ls2")) + ] + if explicit_layerscale: + raise RuntimeError( + "Nemotron Omni refit contains explicit RADIO LayerScale tensors, " + "but the stock vLLM 0.20 loader ignores them. Refusing to replace " + "checkpoint values with the folded-checkpoint identity behavior." + ) + + def _initialize_nemotron_omni_radio_layerscale(self) -> int: + """Set folded RADIO LayerScale parameters to identity while vLLM is awake. + + Nano/Super Omni checkpoints fold RADIO LayerScale into the adjacent + projection weights and therefore do not export ``ls1``/``ls2``. Stock + vLLM 0.20 leaves those parameters at dummy-initialized values during + direct load or refit, corrupting image inference. Initialize the + vLLM-only parameters once, immediately after engine creation and before + colocated sleep can release their CUDA storage. + """ + if not self._is_nemotron_omni(): + return 0 + + model = self.model_runner.model + vision_model = getattr(model, "vision_model", None) + if vision_model is None: + raise RuntimeError( + "Nemotron Omni vLLM model has no vision_model during initialization." + ) + + initializer_factor = getattr(vision_model.config, "initializer_factor", 1.0) + initialized = 0 + with torch.no_grad(): + for name, parameter in vision_model.named_parameters(): + if name.rsplit(".", 1)[-1] in {"ls1", "ls2"}: + parameter.fill_(initializer_factor) + initialized += 1 + + if initialized == 0: + raise RuntimeError( + "Nemotron Omni vLLM model exposes no RADIO ls1/ls2 parameters; " + "the expected vLLM 0.20 model layout may have changed." + ) + return initialized def _maybe_process_fp8_kv_cache(self) -> None: """Process weights after loading for FP8 KV cache (static scales).""" diff --git a/nemo_rl/models/generation/vllm/vllm_worker.py b/nemo_rl/models/generation/vllm/vllm_worker.py index a2a59292ca..dd90a668e1 100644 --- a/nemo_rl/models/generation/vllm/vllm_worker.py +++ b/nemo_rl/models/generation/vllm/vllm_worker.py @@ -48,6 +48,19 @@ logger = logging.getLogger(__name__) +def _context_capped_max_new_tokens( + *, configured_max_new_tokens: int, input_length: int, max_model_len: int +) -> int: + """Cap generation so the training prompt and response fit the context.""" + remaining_context = max_model_len - input_length + if remaining_context <= 0: + raise ValueError( + "Cannot generate from an input whose training length exhausts the " + f"model context: input_length={input_length}, max_model_len={max_model_len}." + ) + return min(configured_max_new_tokens, remaining_context) + + def _resolve_enable_prefix_caching(vllm_cfg: dict[str, Any]) -> bool: enable_prefix_caching = vllm_cfg.get("enable_prefix_caching", None) if enable_prefix_caching is None: @@ -471,11 +484,21 @@ def _load_model(self, bundle_indices, seed): enable_sleep_mode=True, # Set disable_log_stats=False so that self.llm.get_metrics() works. disable_log_stats=False, - logprobs_mode="processed_logprobs", + # Keep the main default, while allowing an RL recipe to request + # raw model logprobs. This is required when policy logprobs are + # compared with vLLM generation under a logits processor. + logprobs_mode=self.cfg["vllm_cfg"].get( + "logprobs_mode", "processed_logprobs" + ), **vllm_kwargs, ) self._create_engine(llm_kwargs) + # Nemotron Omni checkpoints fold RADIO LayerScale into adjacent weights, + # while stock vLLM still allocates ls1/ls2 parameters. Initialize those + # parameters before colocated level-1 sleep releases their CUDA storage; + # mutating them later from prepare_refit_info corrupts sleep/wake state. + self.llm.collective_rpc("_initialize_nemotron_omni_radio_layerscale") log_gpu_memory_diagnostics( label="after_engine_create", worker_type="VllmGenerationWorker", device_id=0 ) @@ -532,6 +555,7 @@ def _build_sampling_params( stop_token_ids=self.cfg["stop_token_ids"], stop=stop_strings, include_stop_str_in_output=True, + bad_words=self.cfg.get("bad_words") or None, ignore_eos=self.cfg.get("ignore_eos", False), ) @@ -682,10 +706,6 @@ def generate( input_lengths = data["input_lengths"] batch_stop_strings: list[list[str]] = data.get("stop_strings", []) stop_strings = self._merge_stop_strings(batch_stop_strings) - sampling_params = self._build_sampling_params( - greedy=greedy, - stop_strings=stop_strings, - ) # verify inputs have correct padding verify_right_padding(data, pad_value=self.cfg["_pad_token_id"]) @@ -693,13 +713,31 @@ def generate( # Original input length with padding padded_input_length = input_ids.size(1) - # Convert inputs to vLLM format - prompts = format_prompt_for_vllm_generation(data) - - # Generate outputs assert self.llm is not None, ( "Attempting to generate with either an uninitialized vLLM or non-model-owner" ) + if self.cfg["vllm_cfg"].get("cap_max_tokens_to_context", False): + max_model_len = int(self.llm.llm_engine.model_config.max_model_len) + sampling_params = [ + self._build_sampling_params( + greedy=greedy, + stop_strings=stop_strings, + max_new_tokens=_context_capped_max_new_tokens( + configured_max_new_tokens=int(self.cfg["max_new_tokens"]), + input_length=int(input_length), + max_model_len=max_model_len, + ), + ) + for input_length in input_lengths.tolist() + ] + else: + sampling_params = self._build_sampling_params( + greedy=greedy, + stop_strings=stop_strings, + ) + + # Convert inputs to vLLM format and generate outputs. + prompts = format_prompt_for_vllm_generation(data) use_tqdm = self.cfg["vllm_cfg"].get("use_tqdm", True) outputs = self.llm.generate(prompts, sampling_params, use_tqdm=use_tqdm) diff --git a/tests/unit/models/generation/test_vllm_backend.py b/tests/unit/models/generation/test_vllm_backend.py index 2fd8327811..254a9c2710 100644 --- a/tests/unit/models/generation/test_vllm_backend.py +++ b/tests/unit/models/generation/test_vllm_backend.py @@ -259,3 +259,82 @@ def test_load_mtp_weights_from_disk_raises_when_mtp_weights_missing( with pytest.raises(ValueError, match="No MTP layer weights"): ext.load_mtp_weights_from_disk(str(model_dir)) ext._load_draft_weights.assert_not_called() + + +class _TinyRadioVisionModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.config = SimpleNamespace(initializer_factor=1.0) + self.ls1 = torch.nn.Parameter(torch.full((4,), 0.001)) + self.ls2 = torch.nn.Parameter(torch.full((4,), -0.001)) + self.projection = torch.nn.Parameter(torch.full((4,), 0.25)) + + +def _make_extension_with_radio(architecture): + from nemo_rl.models.generation.vllm.vllm_backend import ( + VllmInternalWorkerExtension, + ) + + ext = VllmInternalWorkerExtension.__new__(VllmInternalWorkerExtension) + ext.model_runner = SimpleNamespace( + vllm_config=SimpleNamespace( + model_config=SimpleNamespace(architectures=[architecture]) + ), + model=SimpleNamespace(vision_model=_TinyRadioVisionModel()), + ) + return ext + + +@pytest.mark.vllm +def test_initialize_folded_nemotron_radio_layerscale_while_awake(): + ext = _make_extension_with_radio("NemotronH_Nano_Omni_Reasoning_V3") + + initialized = ext._initialize_nemotron_omni_radio_layerscale() + + vision_model = ext.model_runner.model.vision_model + assert initialized == 2 + assert torch.equal(vision_model.ls1, torch.ones_like(vision_model.ls1)) + assert torch.equal(vision_model.ls2, torch.ones_like(vision_model.ls2)) + assert torch.equal( + vision_model.projection, torch.full_like(vision_model.projection, 0.25) + ) + + +@pytest.mark.vllm +def test_prepare_refit_info_does_not_mutate_folded_layerscale(): + ext = _make_extension_with_radio("NemotronH_Nano_Omni_Reasoning_V3") + + state_dict_info = {"language_model.weight": ((4, 4), torch.bfloat16)} + ext.prepare_refit_info(state_dict_info) + + vision_model = ext.model_runner.model.vision_model + assert ext.state_dict_info is state_dict_info + assert torch.equal(vision_model.ls1, torch.full_like(vision_model.ls1, 0.001)) + assert torch.equal(vision_model.ls2, torch.full_like(vision_model.ls2, -0.001)) + + +@pytest.mark.vllm +def test_prepare_refit_info_rejects_explicit_nemotron_radio_layerscale(): + ext = _make_extension_with_radio("NemotronH_Super_Omni_Reasoning_V3") + + with pytest.raises(RuntimeError, match="explicit RADIO LayerScale"): + ext.prepare_refit_info( + { + "vision_model.radio_model.model.blocks.0.ls1": ( + (4,), + torch.bfloat16, + ) + } + ) + + +@pytest.mark.vllm +def test_prepare_refit_info_leaves_non_nemotron_model_unchanged(): + ext = _make_extension_with_radio("SomeOtherArchitecture") + + ext.prepare_refit_info({}) + assert ext._initialize_nemotron_omni_radio_layerscale() == 0 + + vision_model = ext.model_runner.model.vision_model + assert torch.equal(vision_model.ls1, torch.full_like(vision_model.ls1, 0.001)) + assert torch.equal(vision_model.ls2, torch.full_like(vision_model.ls2, -0.001)) diff --git a/tests/unit/models/generation/test_vllm_generation.py b/tests/unit/models/generation/test_vllm_generation.py index bb4f72feda..4582aec55c 100644 --- a/tests/unit/models/generation/test_vllm_generation.py +++ b/tests/unit/models/generation/test_vllm_generation.py @@ -37,6 +37,8 @@ ) from nemo_rl.models.generation.vllm import VllmConfig, VllmGeneration from nemo_rl.models.generation.vllm.vllm_worker import ( + VllmGenerationWorkerImpl, + _context_capped_max_new_tokens, _resolve_enable_prefix_caching, ) from nemo_rl.models.generation.vllm.vllm_worker_async import ( @@ -138,6 +140,52 @@ } +def test_context_capped_max_new_tokens(): + assert ( + _context_capped_max_new_tokens( + configured_max_new_tokens=8192, + input_length=3058, + max_model_len=8192, + ) + == 5134 + ) + assert ( + _context_capped_max_new_tokens( + configured_max_new_tokens=256, + input_length=3058, + max_model_len=8192, + ) + == 256 + ) + with pytest.raises(ValueError, match="exhausts the model context"): + _context_capped_max_new_tokens( + configured_max_new_tokens=8192, + input_length=8192, + max_model_len=8192, + ) + + +def test_sampling_params_preserve_bad_words(): + worker = object.__new__(VllmGenerationWorkerImpl) + worker.cfg = { + "top_k": None, + "temperature": 1.0, + "top_p": 1.0, + "max_new_tokens": 128, + "stop_token_ids": None, + "bad_words": ["", ""], + "ignore_eos": False, + } + worker.SamplingParams = lambda **kwargs: kwargs + + sampling_params = worker._build_sampling_params( + greedy=False, + stop_strings=None, + ) + + assert sampling_params["bad_words"] == ["", ""] + + def test_resolve_enable_prefix_caching_respects_explicit_config(monkeypatch): def raise_if_called(): raise AssertionError("CUDA capability should not be queried") diff --git a/tests/unit/models/generation/test_vllm_patches.py b/tests/unit/models/generation/test_vllm_patches.py new file mode 100644 index 0000000000..78803028f5 --- /dev/null +++ b/tests/unit/models/generation/test_vllm_patches.py @@ -0,0 +1,116 @@ +# Copyright (c) 2026, NVIDIA CORPORATION. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +from pathlib import Path + +from nemo_rl.models.generation.vllm import patches + + +class _Logger: + def __init__(self): + self.info_messages = [] + self.warning_messages = [] + + def info(self, message, *args): + self.info_messages.append(message % args if args else message) + + def warning(self, message, *args): + self.warning_messages.append(message % args if args else message) + + +def _stock_sampling_params_source() -> str: + return """import copy +import json as json_mod +from dataclasses import field + +_SAMPLING_EPS = 1e-5 +_MAX_TEMP = 1e-2 + +class SamplingParams: + def update_from_tokenizer(self, tokenizer: TokenizerLike) -> None: + if not self.bad_words: + return + self._bad_words_token_ids = [] + for bad_word in self.bad_words: + # To prohibit words both at the beginning + # and in the middle of text + # (related to add_prefix_space tokenizer parameter) + for add_prefix_space in [False, True]: + prefix = " " if add_prefix_space else "" + prompt = prefix + bad_word.lstrip() + prompt_token_ids = tokenizer.encode( + text=prompt, add_special_tokens=False + ) + + # If no space at the beginning + # or if prefix space produces a new word token + if (not add_prefix_space) or ( + add_prefix_space + and prompt_token_ids[0] != self._bad_words_token_ids[-1][0] + and len(prompt_token_ids) == len(self._bad_words_token_ids[-1]) + ): + self._bad_words_token_ids.append(prompt_token_ids) + + invalid_token_ids = [ + token_id + for bad_words_token_ids in self._bad_words_token_ids + for token_id in bad_words_token_ids + if token_id < 0 or token_id > tokenizer.max_token_id + ] + if len(invalid_token_ids) > 0: + raise VLLMValidationError( + f"The model vocabulary size is {tokenizer.max_token_id + 1}," + f" but the following tokens" + f" were specified as bad: {invalid_token_ids}." + f" All token id values should be integers satisfying:" + f" 0 <= token_id <= {tokenizer.max_token_id}.", + parameter="bad_words", + value=self.bad_words, + ) +""" + + +def test_bad_words_patch_adds_bounded_thread_safe_cache(tmp_path, monkeypatch): + sampling_params = tmp_path / "sampling_params.py" + sampling_params.write_text(_stock_sampling_params_source()) + monkeypatch.setattr(patches, "_get_vllm_file", lambda _: str(sampling_params)) + logger = _Logger() + + patches._patch_vllm_bad_words_tokenization_cache(logger) + patched = sampling_params.read_text() + + assert "import threading" in patched + assert "_BAD_WORDS_TOKEN_IDS_CACHE_MAX_ENTRIES = 1024" in patched + assert "with _BAD_WORDS_TOKEN_IDS_CACHE_LOCK:" in patched + assert "def _tokenize_bad_words" in patched + assert "return bad_words_token_ids" in patched + assert not logger.warning_messages + + patches._patch_vllm_bad_words_tokenization_cache(logger) + assert sampling_params.read_text() == patched + assert ( + logger.info_messages[-1] == "vLLM bad_words tokenization cache already applied." + ) + + +def test_bad_words_patch_leaves_unknown_vllm_source_untouched(tmp_path, monkeypatch): + sampling_params = Path(tmp_path) / "sampling_params.py" + sampling_params.write_text("# newer vLLM source\n") + monkeypatch.setattr(patches, "_get_vllm_file", lambda _: str(sampling_params)) + logger = _Logger() + + patches._patch_vllm_bad_words_tokenization_cache(logger) + + assert sampling_params.read_text() == "# newer vLLM source\n" + assert "expected vLLM 0.20.0 source shape" in logger.warning_messages[-1] From b8f0e9755d7b5ef36684b3c4b238515db3b275d6 Mon Sep 17 00:00:00 2001 From: Ali Roshan Ghias Date: Wed, 15 Jul 2026 17:01:15 -0700 Subject: [PATCH 4/9] feat(nemotron-omni): add Nano GRPO recipes Signed-off-by: Ali Roshan Ghias --- ...-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml | 80 ++++++++++++ ...-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.yaml | 122 ++++++++++++++++++ 2 files changed, 202 insertions(+) create mode 100644 examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml create mode 100644 examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.yaml diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml new file mode 100644 index 0000000000..1e524f25db --- /dev/null +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml @@ -0,0 +1,80 @@ +defaults: ../../vlm_grpo_3B_megatron.yaml + +grpo: + num_prompts_per_step: 8 + val_at_start: false + # Keep the probability-agreement metric observable instead of masking + # sequences before the maintained recipe can enforce its TMPE gate. + seq_logprob_error_threshold: null + +loss_fn: + # Match the Omni Nano GRPO baseline. A reference-policy copy does not fit + # alongside the policy and colocated vLLM engine on one H100 node. + 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 + train_micro_batch_size: 1 + logprob_batch_size: 1 + max_total_sequence_length: 8192 + dynamic_batching: + enabled: false + sequence_packing: + enabled: true + megatron_cfg: + env_vars: + TORCH_CUDA_ARCH_LIST: "9.0" + tensor_model_parallel_size: 8 + expert_model_parallel_size: 8 + expert_tensor_parallel_size: 1 + context_parallel_size: 1 + sequence_parallel: true + # Nemotron-H uses squared ReLU, which is not supported by Megatron's + # bias+activation fusion. + bias_activation_fusion: false + activation_checkpointing: true + gradient_accumulation_fusion: false + generation: + max_new_tokens: 4096 + bad_words: + - + - + - + - + - + - + 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 + # Megatron recomputes raw policy log probabilities. Comparing those with + # vLLM's processed values produces a false generation/training mismatch. + 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 diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.yaml new file mode 100644 index 0000000000..b0113fd1c2 --- /dev/null +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.yaml @@ -0,0 +1,122 @@ +defaults: ./vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml + +# Production-shaped, non-Gym Nano Omni image-GRPO recipe. This is the main-side +# counterpart of Omni's scripts/nanov3_vision_rl.sh workload. Keep behavioral +# changes synchronized with the frozen Omni control and document exceptions. + +grpo: + num_prompts_per_step: 512 + num_generations_per_prompt: 16 + max_num_steps: 1000000 + val_period: 10 + max_val_samples: 256 + val_batch_size: 256 + seed: 42 + overlong_filtering: true + zero_variance_prompt_filtering: false + seq_logprob_error_threshold: null + # The first maintained comparison disables the fork-only PackedTensor dedup + # optimization on both sides. Qualify it separately before enabling it here. + deduplicate_multimodal_data: false + +loss_fn: + reference_policy_kl_penalty: 0.0 + ratio_clip_min: 0.2 + ratio_clip_max: 0.28 + use_on_policy_kl_approximation: true + sequence_level_importance_ratios: true + token_level_loss: false + +checkpointing: + enabled: true + checkpoint_dir: results/vlm_grpo_nemotron_omni_mmpr_megatron + keep_top_k: 4 + save_period: 10 + checkpoint_must_save_by: 00:03:45:00 + +policy: + train_global_batch_size: 2048 + train_micro_batch_size: 1 + logprob_batch_size: 1 + logprob_chunk_size: 1024 + max_total_sequence_length: 8192 + dynamic_batching: + enabled: false + sequence_packing: + enabled: true + train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} + logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} + algorithm: modified_first_fit_decreasing + sequence_length_round: 64 + megatron_cfg: + empty_unused_memory_level: 2 + tensor_model_parallel_size: 8 + expert_model_parallel_size: 16 + expert_tensor_parallel_size: 1 + context_parallel_size: 1 + sequence_parallel: true + bias_activation_fusion: false + activation_checkpointing: true + optimizer: + lr: 3.0e-06 + min_lr: 2.0e-09 + weight_decay: 0.0 + adam_beta1: 0.9 + adam_beta2: 0.99 + optimizer_cpu_offload: false + optimizer_offload_fraction: 0.0 + scheduler: + lr_decay_style: constant + 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 + logprobs_mode: raw_logprobs + enforce_eager: false + max_model_len: ${policy.max_total_sequence_length} + gpu_memory_utilization: 0.75 + enable_prefix_caching: false + vllm_kwargs: + limit_mm_per_prompt: + image: 2 + max_num_batched_tokens: 32768 + max_num_seqs: 512 + mamba_ssm_cache_dtype: float32 + skip_mm_profiling: true + kernel_config: + enable_flashinfer_autotune: false + moe_backend: triton + +data: + train: + dataset_name: mmpr-tiny + download_dir: results/mmpr_tiny_processed + split_validation_size: 0.008 + seed: 42 + default: + prompt_file: null + processor: vlm_hf_data_processor + 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 + gpus_per_node: 8 From b16217eb3f59a7c66aa68123375cf85635d09e55 Mon Sep 17 00:00:00 2001 From: Ali Roshan Ghias Date: Wed, 15 Jul 2026 17:56:43 -0700 Subject: [PATCH 5/9] style(nemotron-omni): apply repository formatting Signed-off-by: Ali Roshan Ghias --- ...-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml | 27 +-------- ...-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.yaml | 57 ------------------- nemo_rl/data/multimodal_utils.py | 11 +--- nemo_rl/data/processors.py | 4 +- nemo_rl/models/megatron/setup.py | 5 +- tests/unit/data/datasets/test_mmpr_tiny.py | 8 +-- .../models/megatron/test_megatron_setup.py | 8 +-- .../megatron/test_nemotron_omni_model.py | 23 +++----- .../models/policy/test_megatron_worker.py | 18 ++++-- 9 files changed, 31 insertions(+), 130 deletions(-) diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml index 1e524f25db..b738cf17a1 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml @@ -1,43 +1,23 @@ defaults: ../../vlm_grpo_3B_megatron.yaml - -grpo: - num_prompts_per_step: 8 - val_at_start: false - # Keep the probability-agreement metric observable instead of masking - # sequences before the maintained recipe can enforce its TMPE gate. - seq_logprob_error_threshold: null - loss_fn: - # Match the Omni Nano GRPO baseline. A reference-policy copy does not fit - # alongside the policy and colocated vLLM engine on one H100 node. 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 - train_micro_batch_size: 1 logprob_batch_size: 1 max_total_sequence_length: 8192 - dynamic_batching: - enabled: false sequence_packing: enabled: true megatron_cfg: env_vars: - TORCH_CUDA_ARCH_LIST: "9.0" + TORCH_CUDA_ARCH_LIST: '9.0' tensor_model_parallel_size: 8 expert_model_parallel_size: 8 - expert_tensor_parallel_size: 1 - context_parallel_size: 1 sequence_parallel: true - # Nemotron-H uses squared ReLU, which is not supported by Megatron's - # bias+activation fusion. bias_activation_fusion: false activation_checkpointing: true - gradient_accumulation_fusion: false generation: max_new_tokens: 4096 bad_words: @@ -54,8 +34,6 @@ policy: cap_max_tokens_to_context: true gpu_memory_utilization: 0.5 enable_prefix_caching: false - # Megatron recomputes raw policy log probabilities. Comparing those with - # vLLM's processed values produces a false generation/training mismatch. logprobs_mode: raw_logprobs vllm_kwargs: limit_mm_per_prompt: @@ -66,15 +44,12 @@ policy: 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 diff --git a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.yaml b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.yaml index b0113fd1c2..844917c1fc 100644 --- a/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.yaml +++ b/examples/configs/recipes/vlm/vlm_grpo-nemotron-omni-30ba3b-mmpr-4n8g-megatron-tp8ep16.v1.yaml @@ -1,72 +1,30 @@ defaults: ./vlm_grpo-nemotron-omni-30ba3b-clevr-1n8g-megatron-tp8ep8.v1.yaml - -# Production-shaped, non-Gym Nano Omni image-GRPO recipe. This is the main-side -# counterpart of Omni's scripts/nanov3_vision_rl.sh workload. Keep behavioral -# changes synchronized with the frozen Omni control and document exceptions. - grpo: num_prompts_per_step: 512 - num_generations_per_prompt: 16 - max_num_steps: 1000000 - val_period: 10 - max_val_samples: 256 - val_batch_size: 256 - seed: 42 overlong_filtering: true zero_variance_prompt_filtering: false - seq_logprob_error_threshold: null - # The first maintained comparison disables the fork-only PackedTensor dedup - # optimization on both sides. Qualify it separately before enabling it here. deduplicate_multimodal_data: false - loss_fn: - reference_policy_kl_penalty: 0.0 - ratio_clip_min: 0.2 ratio_clip_max: 0.28 use_on_policy_kl_approximation: true sequence_level_importance_ratios: true token_level_loss: false - checkpointing: - enabled: true checkpoint_dir: results/vlm_grpo_nemotron_omni_mmpr_megatron keep_top_k: 4 - save_period: 10 checkpoint_must_save_by: 00:03:45:00 - policy: train_global_batch_size: 2048 - train_micro_batch_size: 1 - logprob_batch_size: 1 logprob_chunk_size: 1024 - max_total_sequence_length: 8192 - dynamic_batching: - enabled: false - sequence_packing: - enabled: true - train_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.train_micro_batch_size}} - logprob_mb_tokens: ${mul:${policy.max_total_sequence_length}, ${policy.logprob_batch_size}} - algorithm: modified_first_fit_decreasing - sequence_length_round: 64 megatron_cfg: empty_unused_memory_level: 2 - tensor_model_parallel_size: 8 expert_model_parallel_size: 16 - expert_tensor_parallel_size: 1 - context_parallel_size: 1 - sequence_parallel: true - bias_activation_fusion: false - activation_checkpointing: true optimizer: lr: 3.0e-06 min_lr: 2.0e-09 weight_decay: 0.0 - adam_beta1: 0.9 adam_beta2: 0.99 - optimizer_cpu_offload: false - optimizer_offload_fraction: 0.0 scheduler: - lr_decay_style: constant lr_decay_iters: null lr_warmup_iters: 10 lr_warmup_init: 3.0e-08 @@ -75,22 +33,12 @@ policy: vllm_cfg: tensor_parallel_size: 2 load_format: auto - logprobs_mode: raw_logprobs enforce_eager: false max_model_len: ${policy.max_total_sequence_length} gpu_memory_utilization: 0.75 - enable_prefix_caching: false vllm_kwargs: - limit_mm_per_prompt: - image: 2 max_num_batched_tokens: 32768 max_num_seqs: 512 - mamba_ssm_cache_dtype: float32 - skip_mm_profiling: true - kernel_config: - enable_flashinfer_autotune: false - moe_backend: triton - data: train: dataset_name: mmpr-tiny @@ -99,9 +47,7 @@ data: seed: 42 default: prompt_file: null - processor: vlm_hf_data_processor env_name: mmpr-tiny - env: mmpr-tiny: num_workers: 8 @@ -110,13 +56,10 @@ env: 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 - gpus_per_node: 8 diff --git a/nemo_rl/data/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py index 7d029af578..9bb0f50334 100644 --- a/nemo_rl/data/multimodal_utils.py +++ b/nemo_rl/data/multimodal_utils.py @@ -123,8 +123,7 @@ def as_tensor( # patchification. if ( self.pad_to_max_shape - and - self.dim_to_pack == 0 + 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 ): @@ -191,9 +190,7 @@ 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 - ] + 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" ) @@ -238,9 +235,7 @@ 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 - ] + 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" ) diff --git a/nemo_rl/data/processors.py b/nemo_rl/data/processors.py index eae85a6b78..64a716fde8 100644 --- a/nemo_rl/data/processors.py +++ b/nemo_rl/data/processors.py @@ -638,9 +638,7 @@ def vlm_hf_data_processor( 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 - ) + 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: diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 4141592781..ec31588775 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -339,7 +339,6 @@ def _validate_nemotron_omni_megatron_lm_layout( the only reliable way to distinguish the old LLaVA wrapper from the canonical model-owned-packing implementation. """ - architectures = set(getattr(hf_config, "architectures", None) or ()) if not architectures.intersection(_NEMOTRON_OMNI_ARCHITECTURES): return @@ -348,9 +347,7 @@ def _validate_nemotron_omni_megatron_lm_layout( tensor_metadata = dist_checkpointing.load_tensors_metadata(checkpoint_dir) legacy_keys = [ - str(key) - for key in tensor_metadata - if "llava_model" in str(key).split(".") + str(key) for key in tensor_metadata if "llava_model" in str(key).split(".") ] if not legacy_keys: return diff --git a/tests/unit/data/datasets/test_mmpr_tiny.py b/tests/unit/data/datasets/test_mmpr_tiny.py index a2c61d5dee..74822398d7 100644 --- a/tests/unit/data/datasets/test_mmpr_tiny.py +++ b/tests/unit/data/datasets/test_mmpr_tiny.py @@ -237,9 +237,7 @@ def test_processor_produces_valid_datum_spec(self, tiny_image_path): assert len(result["vllm_images"]) == 1 assert result["task_name"] == "mmpr-tiny" user_message = result["message_log"][0] - assert torch.equal( - user_message["num_frames"].as_tensor(), torch.tensor([1]) - ) + assert torch.equal(user_message["num_frames"].as_tensor(), torch.tensor([1])) def test_historical_tiled_processor_gets_media_metadata(self, tiny_image_path): from nemo_rl.data.interfaces import TaskDataSpec @@ -247,9 +245,7 @@ def test_historical_tiled_processor_gets_media_metadata(self, tiny_image_path): task_data_spec = TaskDataSpec(task_name="mmpr-tiny") task_data_spec.prompt = _TEST_PROMPT_TEMPLATE - processor = _make_stub_nemotron_processor( - include_imgs_sizes=False, num_tiles=3 - ) + processor = _make_stub_nemotron_processor(include_imgs_sizes=False, num_tiles=3) result = vlm_hf_data_processor( datum_dict={ "images": [tiny_image_path], diff --git a/tests/unit/models/megatron/test_megatron_setup.py b/tests/unit/models/megatron/test_megatron_setup.py index 638e128798..2652d635e8 100644 --- a/tests/unit/models/megatron/test_megatron_setup.py +++ b/tests/unit/models/megatron/test_megatron_setup.py @@ -54,9 +54,7 @@ def test_canonical_omni_checkpoint_layout_is_accepted(self): _validate_nemotron_omni_megatron_lm_layout, ) - hf_config = SimpleNamespace( - architectures=["NemotronH_Nano_Omni_Reasoning_V3"] - ) + hf_config = SimpleNamespace(architectures=["NemotronH_Nano_Omni_Reasoning_V3"]) metadata = { "language_model.embedding.word_embeddings.weight": object(), "vision_model.decoder.layers.0.self_attention.linear_qkv.weight": object(), @@ -74,9 +72,7 @@ def test_legacy_llava_checkpoint_layout_is_rejected(self): _validate_nemotron_omni_megatron_lm_layout, ) - hf_config = SimpleNamespace( - architectures=["NemotronH_Super_Omni_Reasoning_V3"] - ) + hf_config = SimpleNamespace(architectures=["NemotronH_Super_Omni_Reasoning_V3"]) metadata = { "model.llava_model.language_model.embedding.word_embeddings.weight": object(), "model.llava_model.vision_model.decoder.layers.0.self_attention.linear_qkv.weight": object(), diff --git a/tests/unit/models/megatron/test_nemotron_omni_model.py b/tests/unit/models/megatron/test_nemotron_omni_model.py index e09b0ae328..b7c0e81a40 100644 --- a/tests/unit/models/megatron/test_nemotron_omni_model.py +++ b/tests/unit/models/megatron/test_nemotron_omni_model.py @@ -22,20 +22,19 @@ import pytest import torch -from megatron.core import dist_checkpointing, parallel_state -from megatron.core.distributed import DistributedDataParallelConfig -from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed - from megatron.bridge.models.nemotron_omni.nemotron_omni_provider import ( NEMOTRON_OMNI_EXPANDED_SEQUENCE_CONTRACT, NemotronOmniModelProvider, ) +from megatron.core import dist_checkpointing, parallel_state +from megatron.core.distributed import DistributedDataParallelConfig +from megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed + from nemo_rl.distributed.model_utils import ( from_parallel_logits_to_logprobs_packed_sequences, ) from nemo_rl.models.megatron.data import process_microbatch - pytestmark = pytest.mark.mcore _IMAGE_TOKEN_ID = 18 @@ -143,9 +142,7 @@ def _expanded_fixture(device: torch.device): generator.manual_seed(2026) images = torch.randn(2, 3, 32, 64, generator=generator, device=device) images[1, :, :, 32:] = 0 - image_sizes = torch.tensor( - [[32, 64], [32, 32]], dtype=torch.int32, device=device - ) + image_sizes = torch.tensor([[32, 64], [32, 32]], dtype=torch.int32, device=device) return input_ids, lengths, images, image_sizes @@ -177,9 +174,9 @@ def _forward(model): inference_only=False, cp_group=parallel_state.get_context_parallel_group(), ) - prediction_mask = torch.arange( - input_ids.shape[1] - 1, device=device - ).unsqueeze(0) < (lengths - 1).unsqueeze(1) + prediction_mask = torch.arange(input_ids.shape[1] - 1, device=device).unsqueeze( + 0 + ) < (lengths - 1).unsqueeze(1) loss = -(logprobs * prediction_mask).sum() / prediction_mask.sum() return loss, output @@ -256,9 +253,7 @@ def _run_training_checkpoint_roundtrip( provider.finalize() restored_model = provider.provide().cuda().eval() restore_template = restored_model.sharded_state_dict(metadata=metadata) - loaded_state = dist_checkpointing.load( - {"model": restore_template}, checkpoint_dir - ) + loaded_state = dist_checkpointing.load({"model": restore_template}, checkpoint_dir) incompatible = restored_model.load_state_dict(loaded_state["model"]) assert not incompatible.missing_keys assert not incompatible.unexpected_keys diff --git a/tests/unit/models/policy/test_megatron_worker.py b/tests/unit/models/policy/test_megatron_worker.py index 694b5d43e7..b51834483f 100644 --- a/tests/unit/models/policy/test_megatron_worker.py +++ b/tests/unit/models/policy/test_megatron_worker.py @@ -83,9 +83,12 @@ def test_refit_size_estimate_preserves_integral_buffer_dtype(): param = torch.zeros(3, dtype=torch.int64) - assert _estimate_refit_tensor_size_in_bytes( - param, export_dtype=torch.bfloat16, tp_size=2, ep_size=4 - ) == 3 * 8 * 2 * 4 + assert ( + _estimate_refit_tensor_size_in_bytes( + param, export_dtype=torch.bfloat16, tp_size=2, ep_size=4 + ) + == 3 * 8 * 2 * 4 + ) def test_refit_size_estimate_casts_floating_weight_to_export_dtype(): @@ -95,9 +98,12 @@ def test_refit_size_estimate_casts_floating_weight_to_export_dtype(): param = torch.zeros(3, dtype=torch.float32) - assert _estimate_refit_tensor_size_in_bytes( - param, export_dtype=torch.bfloat16, tp_size=2, ep_size=4 - ) == 3 * 2 * 2 * 4 + assert ( + _estimate_refit_tensor_size_in_bytes( + param, export_dtype=torch.bfloat16, tp_size=2, ep_size=4 + ) + == 3 * 2 * 2 * 4 + ) def test_qwen3vl_type_fallback_still_delegates_packing(): From 31d577b5e0db8321c61b5afbde47ae8521f71574 Mon Sep 17 00:00:00 2001 From: Ali Roshan Ghias Date: Wed, 15 Jul 2026 17:56:46 -0700 Subject: [PATCH 6/9] docs(nemotron-omni): document Megatron GRPO support Signed-off-by: Ali Roshan Ghias --- docs/guides/nemotron-3-nano-omni.md | 42 ++++++++++++++++++++++++++--- 1 file changed, 39 insertions(+), 3 deletions(-) diff --git a/docs/guides/nemotron-3-nano-omni.md b/docs/guides/nemotron-3-nano-omni.md index 0a35ba3cfe..306d57dcb5 100644 --- a/docs/guides/nemotron-3-nano-omni.md +++ b/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. From cbe13335007d4844398f923365cb2a5ef45c2285 Mon Sep 17 00:00:00 2001 From: Ali Roshan Ghias Date: Thu, 16 Jul 2026 11:31:44 -0700 Subject: [PATCH 7/9] chore: refresh uv lockfile Signed-off-by: Ali Roshan Ghias --- uv.lock | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/uv.lock b/uv.lock index 976bc8f50b..b2eb99c621 100644 --- a/uv.lock +++ b/uv.lock @@ -2599,7 +2599,7 @@ requires-dist = [ { name = "flask-restful", marker = "extra == 'mlm'" }, { name = "flask-restful", marker = "extra == 'training'" }, { name = "hypercorn", marker = "extra == 'dev'" }, - { name = "mamba-ssm", marker = "extra == 'ssm'", specifier = "~=2.2" }, + { name = "mamba-ssm", marker = "extra == 'ssm'", git = "https://github.com/state-spaces/mamba.git?rev=0048fbf2e7b2f214dcbe703ea3dec2b9647595e1" }, { name = "megatron-energon", extras = ["av-decode"], marker = "extra == 'dev'", specifier = "~=7.0" }, { name = "multi-storage-client", marker = "extra == 'dev'", specifier = "~=0.50" }, { name = "numpy" }, @@ -2620,7 +2620,7 @@ requires-dist = [ { name = "tiktoken", marker = "extra == 'training'" }, { name = "torch", specifier = ">=2.6.0" }, { name = "tqdm", marker = "extra == 'dev'" }, - { name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "extra == 'te'", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=4220403e831d29e93868f7793693ea83f6b8b05b" }, + { name = "transformer-engine", extras = ["core-cu13", "pytorch"], marker = "extra == 'te'", git = "https://github.com/NVIDIA/TransformerEngine.git?rev=b9d690e042b1c4e455214e7dab65d6d3512c05d6" }, { name = "transformers", marker = "extra == 'mlm'" }, { name = "transformers", marker = "extra == 'training'" }, { name = "wandb", marker = "extra == 'mlm'" }, @@ -2662,7 +2662,7 @@ linting = [ ] no-pypi-wheels = [ { name = "emerging-optimizers", git = "https://github.com/NVIDIA-NeMo/Emerging-Optimizers.git?rev=v0.2.0" }, - { name = "flash-mla", git = "https://github.com/deepseek-ai/FlashMLA?rev=9edee0c022cd0938148a18e334203b0aab43aa19" }, + { name = "flash-mla", git = "https://github.com/deepseek-ai/FlashMLA?rev=nv_dev" }, ] test = [ { name = "coverage" }, From a2f0c95c9be6dafe8701e3179f5d17ccc398eae2 Mon Sep 17 00:00:00 2001 From: Ali Roshan Ghias Date: Thu, 16 Jul 2026 16:38:46 -0700 Subject: [PATCH 8/9] chore: update Megatron-Bridge pin Signed-off-by: Ali Roshan Ghias --- 3rdparty/Megatron-Bridge-workspace/Megatron-Bridge | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge index ccf4b09e73..c9473b1361 160000 --- a/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge +++ b/3rdparty/Megatron-Bridge-workspace/Megatron-Bridge @@ -1 +1 @@ -Subproject commit ccf4b09e736a317346b277d42cfdccdea4a919c1 +Subproject commit c9473b136167434e687e826ee84b67e5818a5fa8 From df0db31e6db7da15053b924b05c4a7577c6cbf14 Mon Sep 17 00:00:00 2001 From: Ali Roshan Ghias Date: Mon, 20 Jul 2026 04:30:34 -0700 Subject: [PATCH 9/9] refactor: generalize multimodal tensor padding Signed-off-by: Ali Roshan Ghias --- nemo_rl/data/multimodal_utils.py | 57 +++++++++++++++++-------- tests/unit/data/test_multimodal_dict.py | 25 +++++++++++ 2 files changed, 65 insertions(+), 17 deletions(-) diff --git a/nemo_rl/data/multimodal_utils.py b/nemo_rl/data/multimodal_utils.py index 9bb0f50334..e1aefce2b3 100644 --- a/nemo_rl/data/multimodal_utils.py +++ b/nemo_rl/data/multimodal_utils.py @@ -114,24 +114,47 @@ def as_tensor( if len(non_none_tensors) == 0: return None - # 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) + # 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 = [ - F.pad(t, (0, max_width - t.shape[3], 0, max_height - t.shape[2])) - for t in 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) diff --git a/tests/unit/data/test_multimodal_dict.py b/tests/unit/data/test_multimodal_dict.py index ea56ac79c9..a8cf28a3ea 100644 --- a/tests/unit/data/test_multimodal_dict.py +++ b/tests/unit/data/test_multimodal_dict.py @@ -365,3 +365,28 @@ def test_packedtensor_pads_mixed_dynamic_resolution_images(): torch.testing.assert_close(packed[0, :, 2:, :], torch.zeros(3, 2, 4)) torch.testing.assert_close(packed[1, :, :4, :2], second[0]) torch.testing.assert_close(packed[1, :, :, 2:], torch.zeros(3, 4, 2)) + + +@pytest.mark.parametrize( + ("first_shape", "second_shape", "expected_shape"), + [ + ((1, 2, 3), (2, 4, 3), (3, 4, 3)), + ((1, 2, 3, 2, 4), (2, 4, 3, 4, 2), (3, 4, 3, 4, 4)), + ], +) +def test_packedtensor_pad_to_max_shape_supports_audio_and_video( + first_shape, second_shape, expected_shape +): + """Padding is generic across non-packing dimensions and tensor ranks.""" + first = torch.ones(first_shape) + second = 2 * torch.ones(second_shape) + + packed = PackedTensor( + [first, second], dim_to_pack=0, pad_to_max_shape=True + ).as_tensor() + + assert packed.shape == expected_shape + slices = (slice(0, first_shape[0]),) + tuple( + slice(0, size) for size in first_shape[1:] + ) + torch.testing.assert_close(packed[slices], first)