From 98eadec725ce304f1635c719b6c9770a061686bb Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Wed, 30 Sep 2026 08:03:00 -0700 Subject: [PATCH 1/4] refactor(nemotron): support composite checkpoint conversion --- .../nemotron_h/converting_nemotron_h.py | 36 ++++-- .../models/nemotron_h/modeling_nemotron_h.py | 9 +- .../converting_nemotron_h_omni.py | 91 +++++++++++++++ .../models/nemotron_h_omni/test_conversion.py | 110 ++++++++++++++++++ 4 files changed, 236 insertions(+), 10 deletions(-) create mode 100644 src/prime_rl/trainer/models/nemotron_h_omni/converting_nemotron_h_omni.py create mode 100644 tests/unit/train/models/nemotron_h_omni/test_conversion.py diff --git a/src/prime_rl/trainer/models/nemotron_h/converting_nemotron_h.py b/src/prime_rl/trainer/models/nemotron_h/converting_nemotron_h.py index dead0d9c452..a5fc657e6f3 100644 --- a/src/prime_rl/trainer/models/nemotron_h/converting_nemotron_h.py +++ b/src/prime_rl/trainer/models/nemotron_h/converting_nemotron_h.py @@ -14,16 +14,21 @@ def is_prime_state_dict(state_dict: dict[str, Tensor]) -> bool: ) -def conversion_chain(config: NemotronHConfig) -> list[ConvOp]: +def _conversion_chain( + config: NemotronHConfig, + *, + hf_prefix: str, + prime_prefix: str, +) -> list[ConvOp]: ops: list[ConvOp] = [ - PrefixRename("backbone.", "model."), - Drop("mtp.", is_prefix=True), - Rename("model.embeddings.weight", "model.embed_tokens.weight"), - Rename("model.norm_f.weight", "model.norm.weight"), + PrefixRename(f"{hf_prefix}backbone.", f"{prime_prefix}."), + Drop(f"{hf_prefix}mtp.", is_prefix=True), + Rename(f"{prime_prefix}.embeddings.weight", f"{prime_prefix}.embed_tokens.weight"), + Rename(f"{prime_prefix}.norm_f.weight", f"{prime_prefix}.norm.weight"), ] for layer_idx, layer_type in enumerate(config.layer_types): - prefix = f"model.layers.{layer_idx}" + prefix = f"{prime_prefix}.layers.{layer_idx}" if layer_type == "mamba": ops.append(PrefixRename(f"{prefix}.mixer.", f"{prefix}.mamba.")) elif layer_type == "attention": @@ -52,4 +57,21 @@ def conversion_chain(config: NemotronHConfig) -> list[ConvOp]: return ops -__all__ = ["conversion_chain", "is_hf_state_dict", "is_prime_state_dict"] +def conversion_chain(config: NemotronHConfig) -> list[ConvOp]: + return _conversion_chain(config, hf_prefix="", prime_prefix="model") + + +def composite_language_conversion_chain(config: NemotronHConfig) -> list[ConvOp]: + return _conversion_chain( + config, + hf_prefix="language_model.", + prime_prefix="model.language_model", + ) + + +__all__ = [ + "composite_language_conversion_chain", + "conversion_chain", + "is_hf_state_dict", + "is_prime_state_dict", +] diff --git a/src/prime_rl/trainer/models/nemotron_h/modeling_nemotron_h.py b/src/prime_rl/trainer/models/nemotron_h/modeling_nemotron_h.py index 2dc74b0ae1a..a26311ce93f 100644 --- a/src/prime_rl/trainer/models/nemotron_h/modeling_nemotron_h.py +++ b/src/prime_rl/trainer/models/nemotron_h/modeling_nemotron_h.py @@ -93,7 +93,7 @@ def __init__(self, config: NemotronHConfig, layer_idx: int) -> None: "experts": experts, "shared_expert": shared_expert, "score_before_experts": False, - "load_balance_coeff": config.load_balance_coeff, + "load_balance_coeff": getattr(config, "load_balance_coeff", None), } if config.moe_latent_size is None: self.mlp = MoE(**moe_kwargs) @@ -176,14 +176,17 @@ def __init__(self, config: NemotronHConfig) -> None: def forward( self, - input_ids: torch.LongTensor, + input_ids: torch.LongTensor | None = None, position_ids: torch.LongTensor | None = None, + inputs_embeds: torch.Tensor | None = None, routed_experts: torch.LongTensor | None = None, *, seq_lens: torch.LongTensor, seq_lens_are_pre_shard: bool = False, ) -> BaseModelOutput: - hidden_states = self.embed_tokens(input_ids) + if (input_ids is None) == (inputs_embeds is None): + raise ValueError("Exactly one of input_ids or inputs_embeds must be provided") + hidden_states = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds cu_seqlens, max_seqlen = get_cu_seqlens_from_seq_lens( seq_lens.to(device=hidden_states.device), diff --git a/src/prime_rl/trainer/models/nemotron_h_omni/converting_nemotron_h_omni.py b/src/prime_rl/trainer/models/nemotron_h_omni/converting_nemotron_h_omni.py new file mode 100644 index 00000000000..5a09a9aaba8 --- /dev/null +++ b/src/prime_rl/trainer/models/nemotron_h_omni/converting_nemotron_h_omni.py @@ -0,0 +1,91 @@ +from torch import Tensor + +from prime_rl.trainer.models.conversion_ops import ConvOp, Drop, PrefixRename, Rename, SplitConcat +from prime_rl.trainer.models.nemotron_h.converting_nemotron_h import composite_language_conversion_chain + + +def is_hf_state_dict(state_dict: dict[str, Tensor]) -> bool: + return "language_model.backbone.embeddings.weight" in state_dict + + +def is_prime_state_dict(state_dict: dict[str, Tensor]) -> bool: + return "model.language_model.embed_tokens.weight" in state_dict or any( + name.startswith("model.language_model.layers.") + and any(namespace in name for namespace in (".mamba.", ".self_attn.", ".mlp.router.")) + for name in state_dict + ) + + +def conversion_chain(config) -> list[ConvOp]: + ops: list[ConvOp] = [ + Drop("backbone.", is_prefix=True), + Drop("mtp.", is_prefix=True), + Drop("lm_head.weight"), + *composite_language_conversion_chain(config.llm_config), + Rename("language_model.lm_head.weight", "lm_head.weight"), + Drop("vision_model.radio_model.input_conditioner.", is_prefix=True), + PrefixRename( + "vision_model.radio_model.model.patch_generator.video_embedder.", + "model.vision_model.embeddings.video_patch_projection.", + ), + PrefixRename( + "vision_model.radio_model.model.patch_generator.embedder.", + "model.vision_model.embeddings.patch_projection.", + ), + Rename( + "vision_model.radio_model.model.patch_generator.pos_embed", + "model.vision_model.embeddings.position_embedding", + ), + Rename( + "vision_model.radio_model.model.patch_generator.cls_token.token", + "model.vision_model.embeddings.cls_register_token", + ), + Rename("vision_model.summary_idxs", "model.vision_model.summary_idxs"), + PrefixRename( + "vision_model.radio_model.model.blocks.", + "model.vision_model.encoder.layer.", + ), + ] + + for layer_idx in range(config.vision_config.num_hidden_layers): + prefix = f"model.vision_model.encoder.layer.{layer_idx}" + ops.extend( + [ + Rename(f"{prefix}.attn.proj.weight", f"{prefix}.attention.output.dense.weight"), + Rename(f"{prefix}.attn.proj.bias", f"{prefix}.attention.output.dense.bias"), + SplitConcat( + combined=f"{prefix}.attn.qkv.weight", + parts=[ + (f"{prefix}.attention.attention.query.weight", None), + (f"{prefix}.attention.attention.key.weight", None), + (f"{prefix}.attention.attention.value.weight", None), + ], + dim=0, + ), + SplitConcat( + combined=f"{prefix}.attn.qkv.bias", + parts=[ + (f"{prefix}.attention.attention.query.bias", None), + (f"{prefix}.attention.attention.key.bias", None), + (f"{prefix}.attention.attention.value.bias", None), + ], + dim=0, + ), + ] + ) + + ops.extend( + [ + Rename("mlp1.0.weight", "model.vision_projector.mlp1.norm.weight"), + Rename("mlp1.1.weight", "model.vision_projector.mlp1.linear1.weight"), + Rename("mlp1.3.weight", "model.vision_projector.mlp1.linear2.weight"), + PrefixRename( + "vision_projector.vision_final_layernorm.", + "model.vision_projector.vision_final_layernorm.", + ), + ] + ) + return ops + + +__all__ = ["conversion_chain", "is_hf_state_dict", "is_prime_state_dict"] diff --git a/tests/unit/train/models/nemotron_h_omni/test_conversion.py b/tests/unit/train/models/nemotron_h_omni/test_conversion.py new file mode 100644 index 00000000000..f139ad21f20 --- /dev/null +++ b/tests/unit/train/models/nemotron_h_omni/test_conversion.py @@ -0,0 +1,110 @@ +import torch + +from prime_rl.trainer.models.conversion_ops import apply_hf_to_prime, apply_prime_to_hf +from prime_rl.trainer.models.nemotron_h.configuration_nemotron_h import NemotronHConfig +from prime_rl.trainer.models.nemotron_h.converting_nemotron_h import ( + composite_language_conversion_chain, +) +from prime_rl.trainer.models.nemotron_h.converting_nemotron_h import ( + conversion_chain as standalone_language_conversion_chain, +) +from prime_rl.trainer.models.nemotron_h_omni.converting_nemotron_h_omni import conversion_chain + + +def _config() -> object: + class VisionConfig: + num_hidden_layers = 1 + + class Config: + llm_config = NemotronHConfig( + hidden_size=8, + vocab_size=16, + layers_block_type=["moe"], + n_routed_experts=2, + num_experts_per_tok=1, + moe_intermediate_size=4, + moe_shared_expert_intermediate_size=4, + moe_latent_size=4, + ) + vision_config = VisionConfig() + + return Config() + + +def test_language_conversion_uses_canonical_standalone_and_composite_namespaces(): + config = _config().llm_config + standalone = {"backbone.embeddings.weight": torch.arange(16)} + composite = {"language_model.backbone.embeddings.weight": torch.arange(16)} + + apply_hf_to_prime(standalone, standalone_language_conversion_chain(config)) + apply_hf_to_prime(composite, composite_language_conversion_chain(config)) + + assert "model.embed_tokens.weight" in standalone + assert "model.language_model.embed_tokens.weight" in composite + + +def test_composite_conversion_maps_language_vision_and_projector_weights(): + state_dict = { + "language_model.backbone.embeddings.weight": torch.arange(16), + "language_model.backbone.layers.0.mixer.gate.weight": torch.arange(8), + "language_model.backbone.layers.0.mixer.experts.0.up_proj.weight": torch.ones(4, 4), + "language_model.backbone.layers.0.mixer.experts.1.up_proj.weight": torch.full((4, 4), 2.0), + "language_model.lm_head.weight": torch.arange(16), + "vision_model.radio_model.model.patch_generator.embedder.weight": torch.arange(96).reshape(8, 12), + "vision_model.radio_model.model.blocks.0.attn.qkv.weight": torch.arange(48).reshape(24, 2), + "vision_model.radio_model.model.blocks.0.attn.qkv.bias": torch.arange(24), + "mlp1.0.weight": torch.ones(8), + "vision_projector.vision_final_layernorm.weight": torch.ones(8), + "language_model.mtp.layer.weight": torch.ones(1), + "backbone.embeddings.weight": torch.zeros(16), + } + + apply_hf_to_prime(state_dict, conversion_chain(_config())) + + assert "model.language_model.embed_tokens.weight" in state_dict + assert "model.language_model.layers.0.mlp.router.gate.weight" in state_dict + assert state_dict["model.language_model.layers.0.mlp.experts.up_proj"].shape == (2, 4, 4) + assert "lm_head.weight" in state_dict + assert "model.vision_model.embeddings.patch_projection.weight" in state_dict + assert "model.vision_model.encoder.layer.0.attention.attention.query.weight" in state_dict + assert "model.vision_model.encoder.layer.0.attention.attention.key.weight" in state_dict + assert "model.vision_model.encoder.layer.0.attention.attention.value.weight" in state_dict + assert "model.vision_projector.mlp1.norm.weight" in state_dict + assert "model.vision_projector.vision_final_layernorm.weight" in state_dict + experts = state_dict["model.language_model.layers.0.mlp.experts.up_proj"] + torch.testing.assert_close(experts[0], torch.ones(4, 4)) + torch.testing.assert_close(experts[1], torch.full((4, 4), 2.0)) + qkv = torch.arange(48).reshape(24, 2) + for name, expected in zip( + ("query", "key", "value"), + qkv.chunk(3), + strict=True, + ): + torch.testing.assert_close( + state_dict[f"model.vision_model.encoder.layer.0.attention.attention.{name}.weight"], + expected, + ) + assert not any(name.startswith(("language_model.mtp.", "backbone.")) for name in state_dict) + + +def test_composite_conversion_roundtrip_uses_canonical_checkpoint_names(): + prime_state_dict = { + "model.language_model.embed_tokens.weight": torch.arange(16), + "model.vision_model.encoder.layer.0.attention.attention.query.weight": torch.ones(2, 2), + "model.vision_model.encoder.layer.0.attention.attention.key.weight": torch.full((2, 2), 2.0), + "model.vision_model.encoder.layer.0.attention.attention.value.weight": torch.full((2, 2), 3.0), + "model.vision_projector.mlp1.linear1.weight": torch.ones(4, 8), + "lm_head.weight": torch.arange(16), + } + + apply_prime_to_hf(prime_state_dict, conversion_chain(_config())) + + assert "language_model.backbone.embeddings.weight" in prime_state_dict + assert "vision_model.radio_model.model.blocks.0.attn.qkv.weight" in prime_state_dict + torch.testing.assert_close( + prime_state_dict["vision_model.radio_model.model.blocks.0.attn.qkv.weight"], + torch.cat((torch.ones(2, 2), torch.full((2, 2), 2.0), torch.full((2, 2), 3.0))), + ) + assert "mlp1.1.weight" in prime_state_dict + assert "language_model.lm_head.weight" in prime_state_dict + assert not any(name.startswith("model.") for name in prime_state_dict) From 370157eea01815cbf98ca8ce89228e685f24973b Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Wed, 30 Sep 2026 08:04:23 -0700 Subject: [PATCH 2/4] feat(trainer): add Nemotron-H Omni model --- src/prime_rl/trainer/models/__init__.py | 2 + .../models/nemotron_h_omni/__init__.py | 17 + .../modeling_nemotron_h_omni.py | 515 ++++++++++++++++++ .../models/nemotron_h_omni/test_model.py | 469 ++++++++++++++++ 4 files changed, 1003 insertions(+) create mode 100644 src/prime_rl/trainer/models/nemotron_h_omni/__init__.py create mode 100644 src/prime_rl/trainer/models/nemotron_h_omni/modeling_nemotron_h_omni.py create mode 100644 tests/unit/train/models/nemotron_h_omni/test_model.py diff --git a/src/prime_rl/trainer/models/__init__.py b/src/prime_rl/trainer/models/__init__.py index 2c0355aa18b..d94eb963dec 100644 --- a/src/prime_rl/trainer/models/__init__.py +++ b/src/prime_rl/trainer/models/__init__.py @@ -20,6 +20,7 @@ from prime_rl.trainer.models.llama import LlamaForCausalLM from prime_rl.trainer.models.minimax_m2 import MiniMaxM2Config, MiniMaxM2ForCausalLM from prime_rl.trainer.models.nemotron_h import NemotronHConfig, NemotronHForCausalLM +from prime_rl.trainer.models.nemotron_h_omni import NemotronHOmniForCausalLM from prime_rl.trainer.models.qwen3 import Qwen3ForCausalLM from prime_rl.trainer.models.qwen3_5 import ( Qwen3_5Config, @@ -106,6 +107,7 @@ def supports_custom_impl(model_config: PretrainedConfig) -> bool: _CUSTOM_VLM_MAPPING: dict[str, type] = { "qwen3_5": Qwen3_5ForCausalLM, "qwen3_5_moe": Qwen3_5ForCausalLM, + "nemotron_h_omni": NemotronHOmniForCausalLM, } diff --git a/src/prime_rl/trainer/models/nemotron_h_omni/__init__.py b/src/prime_rl/trainer/models/nemotron_h_omni/__init__.py new file mode 100644 index 00000000000..f7a1bd76c18 --- /dev/null +++ b/src/prime_rl/trainer/models/nemotron_h_omni/__init__.py @@ -0,0 +1,17 @@ +from prime_rl.trainer.models.nemotron_h_omni.converting_nemotron_h_omni import ( + conversion_chain, + is_hf_state_dict, + is_prime_state_dict, +) +from prime_rl.trainer.models.nemotron_h_omni.modeling_nemotron_h_omni import ( + NemotronHOmniForCausalLM, + NemotronHOmniModel, +) + +__all__ = [ + "NemotronHOmniForCausalLM", + "NemotronHOmniModel", + "conversion_chain", + "is_hf_state_dict", + "is_prime_state_dict", +] diff --git a/src/prime_rl/trainer/models/nemotron_h_omni/modeling_nemotron_h_omni.py b/src/prime_rl/trainer/models/nemotron_h_omni/modeling_nemotron_h_omni.py new file mode 100644 index 00000000000..a0bc57f6601 --- /dev/null +++ b/src/prime_rl/trainer/models/nemotron_h_omni/modeling_nemotron_h_omni.py @@ -0,0 +1,515 @@ +import importlib +import re +from functools import lru_cache +from pathlib import Path + +import torch +from torch import Tensor, nn +from transformers.dynamic_module_utils import get_class_from_dynamic_module +from transformers.modeling_outputs import BaseModelOutput + +from prime_rl.trainer.models.layers.lm_head import PrimeLmOutput, VanillaOutputLinear +from prime_rl.trainer.models.layers.moe import MoE +from prime_rl.trainer.models.nemotron_h.modeling_nemotron_h import ( + NemotronHModel, + NemotronHPreTrainedModel, +) +from prime_rl.trainer.models.nemotron_h_omni.converting_nemotron_h_omni import ( + conversion_chain, + is_hf_state_dict, + is_prime_state_dict, +) +from prime_rl.utils.cp import ( + CPContext, + setup_cp_attention_params, + shard_for_cp, + shard_position_ids_for_cp, +) + +_REMOTE_PROJECTOR_CLASS = "NemotronH_Omni_Reasoning_V3VisionProjector" +_EXPECTED_MODEL_TYPE = "nemotron_h_omni" +_EXPECTED_AUTO_MODEL = "modeling_nemotron_h_omni.NemotronH_Omni_Reasoning_V3" +_IMMUTABLE_REVISION = re.compile(r"[0-9a-fA-F]{40}") + + +@lru_cache +def _load_remote_model_module( + class_reference: str, + model_path: str, + revision: str | None, +): + remote_model_class = get_class_from_dynamic_module( + class_reference, + model_path, + revision=revision, + ) + return importlib.import_module(remote_model_class.__module__) + + +def _load_remote_vision_components( + config, + *, + trust_remote_code: bool, +) -> tuple[type[nn.Module], type[nn.Module]]: + if trust_remote_code is not True: + raise ValueError("Nemotron-H Omni requires trust_remote_code=True to load checkpoint vision code") + if getattr(config, "model_type", None) != _EXPECTED_MODEL_TYPE: + raise ValueError(f"Nemotron-H Omni remote vision code requires model_type={_EXPECTED_MODEL_TYPE!r}") + + auto_map = getattr(config, "auto_map", None) or {} + class_reference = auto_map.get("AutoModelForImageTextToText") + if class_reference != _EXPECTED_AUTO_MODEL: + raise ValueError( + "Nemotron-H Omni requires the expected AutoModelForImageTextToText " + f"class reference {_EXPECTED_AUTO_MODEL!r}" + ) + + model_path = str(getattr(config, "_name_or_path", "") or "") + if not model_path: + raise ValueError("Nemotron-H Omni requires a checkpoint model path") + revision = getattr(config, "_commit_hash", None) + local_path = Path(model_path).expanduser() + if not local_path.is_dir() and (not isinstance(revision, str) or _IMMUTABLE_REVISION.fullmatch(revision) is None): + raise ValueError("Nemotron-H Omni remote vision code requires a full immutable commit revision") + module = _load_remote_model_module( + class_reference, + model_path, + revision, + ) + return module.RadioModel, getattr(module, _REMOTE_PROJECTOR_CLASS) + + +def merge_image_embeddings( + input_ids: torch.LongTensor, + inputs_embeds: Tensor, + image_embeds: Tensor, + *, + image_token_id: int, +) -> Tensor: + image_embeds = image_embeds.reshape(-1, image_embeds.shape[-1]) + if image_embeds.shape[-1] != inputs_embeds.shape[-1]: + raise ValueError( + "Nemotron-H Omni projected image embedding hidden size does not match " + f"token embeddings: {image_embeds.shape[-1]} != {inputs_embeds.shape[-1]}" + ) + image_mask = input_ids == image_token_id + image_token_count = int(image_mask.sum().item()) + if image_token_count != image_embeds.shape[0]: + raise ValueError( + "Nemotron-H Omni image token count does not match projected image embeddings: " + f"{image_token_count} tokens != {image_embeds.shape[0]} embeddings" + ) + return inputs_embeds.masked_scatter( + image_mask.unsqueeze(-1).expand_as(inputs_embeds), + image_embeds.to(device=inputs_embeds.device, dtype=inputs_embeds.dtype), + ) + + +def _is_integer_tensor(value: Tensor) -> bool: + return value.dtype in { + torch.int8, + torch.int16, + torch.int32, + torch.int64, + torch.uint8, + } + + +def _normalize_pixel_values( + pixel_values: Tensor | None, + imgs_sizes: torch.LongTensor | None, + *, + num_channels: int, +) -> Tensor | list[Tensor] | None: + if pixel_values is None: + return None + if not torch.is_floating_point(pixel_values): + raise ValueError("Nemotron-H Omni pixel_values must use a floating-point dtype") + if pixel_values.ndim == 4: + return pixel_values + if pixel_values.ndim != 1: + raise ValueError("Nemotron-H Omni pixel_values must be a flat wire tensor or a 4D image tensor") + if not isinstance(num_channels, int) or isinstance(num_channels, bool) or num_channels <= 0: + raise ValueError("Nemotron-H Omni vision channel count must be positive") + if imgs_sizes is None or imgs_sizes.ndim != 2 or imgs_sizes.shape[1] != 2: + raise ValueError("Nemotron-H Omni flat pixel_values require imgs_sizes with shape (images, 2)") + if not _is_integer_tensor(imgs_sizes): + raise ValueError("Nemotron-H Omni imgs_sizes must use an integer dtype") + if imgs_sizes.numel() == 0 or not torch.all(imgs_sizes > 0): + raise ValueError("Nemotron-H Omni imgs_sizes values must be positive") + + image_sizes = imgs_sizes.tolist() + expected_values = sum(num_channels * height * width for height, width in image_sizes) + if pixel_values.numel() != expected_values: + raise ValueError("Nemotron-H Omni flat pixel_values element count does not match imgs_sizes") + + images = [] + offset = 0 + for height, width in image_sizes: + num_values = num_channels * height * width + images.append(pixel_values[offset : offset + num_values].reshape(1, num_channels, height, width)) + offset += num_values + return images + + +def _validate_image_inputs( + input_ids: torch.LongTensor, + pixel_values: Tensor | list[Tensor] | None, + imgs_sizes: torch.LongTensor | None, + num_tokens: torch.LongTensor | None, + num_patches: torch.LongTensor | None, + mm_token_type_ids: torch.LongTensor | None, + *, + image_token_id: int, + num_channels: int, +) -> None: + image_mask = input_ids == image_token_id + if mm_token_type_ids is not None: + if mm_token_type_ids.shape != input_ids.shape: + raise ValueError("Nemotron-H Omni mm_token_type_ids must match input_ids") + if not torch.equal(mm_token_type_ids == 1, image_mask): + raise ValueError("Nemotron-H Omni image markers do not match image tokens") + + metadata = { + "imgs_sizes": imgs_sizes, + "num_tokens": num_tokens, + "num_patches": num_patches, + } + if pixel_values is None: + if any(value is not None for value in metadata.values()): + raise ValueError("Nemotron-H Omni image metadata requires pixel_values") + return + if any(value is not None for value in metadata.values()) and any(value is None for value in metadata.values()): + raise ValueError("Nemotron-H Omni image metadata must be supplied together") + + images = list(pixel_values) if isinstance(pixel_values, list) else list(pixel_values.split(1)) + if not images or any( + image.ndim != 4 + or image.shape[0] != 1 + or image.shape[1] != num_channels + or image.shape[-2] <= 0 + or image.shape[-1] <= 0 + for image in images + ): + raise ValueError( + f"Nemotron-H Omni image tensors must have shape (1, {num_channels}, positive height, positive width)" + ) + if any(not torch.is_floating_point(image) for image in images): + raise ValueError("Nemotron-H Omni pixel_values must use a floating-point dtype") + + num_images = len(images) + if imgs_sizes is not None: + if not _is_integer_tensor(imgs_sizes): + raise ValueError("Nemotron-H Omni imgs_sizes must use an integer dtype") + if imgs_sizes.shape != (num_images, 2): + raise ValueError("Nemotron-H Omni imgs_sizes must have shape (images, 2)") + if not torch.all(imgs_sizes > 0): + raise ValueError("Nemotron-H Omni imgs_sizes values must be positive") + for image, (height, width) in zip(images, imgs_sizes.tolist(), strict=True): + if image.shape[-2:] != (height, width): + raise ValueError("Nemotron-H Omni imgs_sizes do not match pixel_values") + for name, value in (("num_tokens", num_tokens), ("num_patches", num_patches)): + if value is None: + continue + if not _is_integer_tensor(value): + raise ValueError(f"Nemotron-H Omni {name} must use an integer dtype") + if value.shape != (num_images,): + raise ValueError(f"Nemotron-H Omni {name} must have shape (images,)") + if not torch.all(value > 0): + raise ValueError(f"Nemotron-H Omni {name} values must be positive") + if num_patches is not None and not torch.all(num_patches == 1): + raise ValueError("Nemotron-H Omni image inputs require one patch group per image") + if num_tokens is not None and int(num_tokens.sum().item()) != int(image_mask.sum().item()): + raise ValueError("Nemotron-H Omni num_tokens do not match image tokens") + + +def _replace_unit_radio_layer_scales(module: nn.Module, value: float) -> None: + if value != 1.0: + raise ValueError( + "Nemotron-H Omni requires RADIO layerscale_value=1.0 because the " + "checkpoint does not contain LayerScale parameters" + ) + for name, child in tuple(module.named_children()): + if child.__class__.__name__ == "RadioLayerScale": + setattr(module, name, nn.Identity()) + else: + _replace_unit_radio_layer_scales(child, value) + + +class NemotronHOmniModel(nn.Module): + def __init__( + self, + config, + *, + trust_remote_code: bool, + ) -> None: + super().__init__() + self.config = config + vision_model_class, vision_projector_class = _load_remote_vision_components( + config, + trust_remote_code=trust_remote_code, + ) + self.language_model = NemotronHModel(config.llm_config) + self.cp_context = CPContext() + self.vision_model = vision_model_class(config.vision_config) + _replace_unit_radio_layer_scales( + self.vision_model, + config.vision_config.layerscale_value, + ) + self.vision_model.make_preprocessor_external() + self.vision_projector = vision_projector_class(config) + + def get_input_embeddings(self) -> nn.Embedding: + return self.language_model.embed_tokens + + def set_input_embeddings(self, embeddings: nn.Embedding) -> None: + self.language_model.embed_tokens = embeddings + + def _dummy_pixel_values(self, inputs_embeds: Tensor) -> Tensor: + image_size = self.config.patch_size * int(1 / self.config.downsample_ratio) + num_channels = getattr(self.config.vision_config, "num_channels", 3) + vision_parameter = next(self.vision_model.parameters()) + return torch.zeros( + 1, + num_channels, + image_size, + image_size, + device=inputs_embeds.device, + dtype=vision_parameter.dtype, + ) + + def prepare_inputs( + self, + input_ids: torch.LongTensor, + pixel_values: Tensor | None, + image_flags: torch.LongTensor | None, + imgs_sizes: torch.LongTensor | None, + num_tokens: torch.LongTensor | None, + num_patches: torch.LongTensor | None, + mm_token_type_ids: torch.LongTensor | None, + ) -> Tensor: + has_images = pixel_values is not None + num_channels = getattr(self.config.vision_config, "num_channels", 3) + pixel_values = _normalize_pixel_values( + pixel_values, + imgs_sizes, + num_channels=num_channels, + ) + _validate_image_inputs( + input_ids, + pixel_values, + imgs_sizes, + num_tokens, + num_patches, + mm_token_type_ids, + image_token_id=self.config.img_context_token_id, + num_channels=num_channels, + ) + inputs_embeds = self.language_model.embed_tokens(input_ids) + if pixel_values is None: + pixel_values = self._dummy_pixel_values(inputs_embeds) + + if isinstance(pixel_values, list): + image_embeds = torch.cat( + [self.vision_projector(image, self.vision_model).flatten(0, 1) for image in pixel_values], + dim=0, + ) + else: + image_embeds = self.vision_projector(pixel_values, self.vision_model) + if has_images: + if image_flags is not None: + image_embeds = image_embeds[image_flags.squeeze(-1) == 1] + return merge_image_embeddings( + input_ids, + inputs_embeds, + image_embeds, + image_token_id=self.config.img_context_token_id, + ) + return inputs_embeds + image_embeds.sum() * 0.0 + + def forward( + self, + input_ids: torch.LongTensor, + position_ids: torch.LongTensor | None = None, + pixel_values: Tensor | None = None, + image_flags: torch.LongTensor | None = None, + mm_token_type_ids: torch.LongTensor | None = None, + imgs_sizes: torch.LongTensor | None = None, + num_tokens: torch.LongTensor | None = None, + num_patches: torch.LongTensor | None = None, + routed_experts: torch.LongTensor | None = None, + *, + seq_lens: torch.LongTensor, + seq_lens_are_pre_shard: bool = False, + ) -> BaseModelOutput: + has_images = pixel_values is not None + inputs_embeds = self.prepare_inputs( + input_ids, + pixel_values, + image_flags, + imgs_sizes, + num_tokens, + num_patches, + mm_token_type_ids, + ) + + if has_images and self.cp_context.cp_enabled: + if position_ids is None: + position_ids = torch.arange( + inputs_embeds.shape[1], + device=inputs_embeds.device, + ).unsqueeze(0) + rank = self.cp_context.cp_rank + world_size = self.cp_context.cp_world_size + setup_cp_attention_params( + position_ids, + cp_group=self.cp_context.cp_group, + cp_style=self.cp_context.cp_style, + seq_lens=seq_lens, + ) + inputs_embeds = shard_for_cp( + inputs_embeds, + cp_rank=rank, + cp_world_size=world_size, + ) + position_ids = shard_position_ids_for_cp( + position_ids, + cp_rank=rank, + cp_world_size=world_size, + ) + if routed_experts is not None: + routed_experts = shard_for_cp( + routed_experts, + cp_rank=rank, + cp_world_size=world_size, + ) + seq_lens_are_pre_shard = True + + return self.language_model( + inputs_embeds=inputs_embeds, + position_ids=position_ids, + routed_experts=routed_experts, + seq_lens=seq_lens, + seq_lens_are_pre_shard=seq_lens_are_pre_shard, + ) + + +class NemotronHOmniForCausalLM(NemotronHPreTrainedModel): + retain_frozen_vision_weights_after_forward = True + + def __init__( + self, + config, + *, + _prime_trust_remote_code: bool = False, + ) -> None: + super().__init__(config) + self.model = NemotronHOmniModel( + config, + trust_remote_code=_prime_trust_remote_code, + ) + self.lm_head = VanillaOutputLinear( + config.llm_config.hidden_size, + config.llm_config.vocab_size, + ) + self.supports_packed_multimodal_training = True + + @classmethod + def from_config(cls, config, trust_remote_code: bool = False, **kwargs): + return cls._from_config( + config, + _prime_trust_remote_code=trust_remote_code, + **kwargs, + ) + + @classmethod + def from_pretrained( + cls, + pretrained_model_name_or_path, + *model_args, + trust_remote_code: bool = False, + **kwargs, + ): + return super().from_pretrained( + pretrained_model_name_or_path, + *model_args, + trust_remote_code=trust_remote_code, + _prime_trust_remote_code=trust_remote_code, + **kwargs, + ) + + @classmethod + def is_hf_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: + return is_hf_state_dict(state_dict) + + @classmethod + def is_prime_state_dict(cls, state_dict: dict[str, Tensor]) -> bool: + return is_prime_state_dict(state_dict) + + @classmethod + def conversion_chain(cls, config) -> list: + return conversion_chain(config) + + def get_input_embeddings(self) -> nn.Embedding: + return self.model.get_input_embeddings() + + def set_input_embeddings(self, embeddings: nn.Embedding) -> None: + self.model.set_input_embeddings(embeddings) + + def forward( + self, + input_ids: torch.LongTensor, + position_ids: torch.LongTensor | None = None, + labels: torch.LongTensor | None = None, + temperature: Tensor | None = None, + sampling_mask: Tensor | None = None, + routed_experts: torch.LongTensor | None = None, + pixel_values: Tensor | None = None, + image_flags: torch.LongTensor | None = None, + mm_token_type_ids: torch.LongTensor | None = None, + imgs_sizes: torch.LongTensor | None = None, + num_tokens: torch.LongTensor | None = None, + num_patches: torch.LongTensor | None = None, + *, + seq_lens: torch.LongTensor, + seq_lens_are_pre_shard: bool = False, + ) -> PrimeLmOutput: + outputs = self.model( + input_ids=input_ids, + position_ids=position_ids, + pixel_values=pixel_values, + image_flags=image_flags, + mm_token_type_ids=mm_token_type_ids, + imgs_sizes=imgs_sizes, + num_tokens=num_tokens, + num_patches=num_patches, + routed_experts=routed_experts, + seq_lens=seq_lens, + seq_lens_are_pre_shard=seq_lens_are_pre_shard, + ) + return self.lm_head( + outputs.last_hidden_state, + labels, + temperature=temperature, + sampling_mask=sampling_mask, + ) + + @torch.no_grad() + def init_buffers_post_meta(self) -> None: + self.model.vision_model.summary_idxs.copy_( + torch.tensor( + self.config.vision_config.summary_idxs, + device=self.model.vision_model.summary_idxs.device, + ) + ) + for module in self.modules(): + if isinstance(module, MoE): + module.tokens_per_expert.zero_() + module.routing_confidence_sum.zero_() + + +__all__ = [ + "NemotronHOmniForCausalLM", + "NemotronHOmniModel", + "merge_image_embeddings", +] diff --git a/tests/unit/train/models/nemotron_h_omni/test_model.py b/tests/unit/train/models/nemotron_h_omni/test_model.py new file mode 100644 index 00000000000..eb2e16b366a --- /dev/null +++ b/tests/unit/train/models/nemotron_h_omni/test_model.py @@ -0,0 +1,469 @@ +from types import SimpleNamespace + +import pytest +import torch +from torch import nn +from transformers import PretrainedConfig + +from prime_rl.trainer.model import forward as trainer_forward +from prime_rl.trainer.models import get_custom_vlm_cls +from prime_rl.trainer.models.layers.lm_head import inject_prime_lm_head +from prime_rl.trainer.models.nemotron_h.configuration_nemotron_h import NemotronHConfig +from prime_rl.trainer.models.nemotron_h_omni import NemotronHOmniForCausalLM +from prime_rl.trainer.models.nemotron_h_omni.modeling_nemotron_h_omni import ( + _load_remote_vision_components, + merge_image_embeddings, +) +from prime_rl.utils.vlm import VLM_REGISTRY + + +def test_nemotron_h_omni_is_registered_as_packed_vlm(): + config = PretrainedConfig() + config.model_type = "nemotron_h_omni" + + assert get_custom_vlm_cls(config) is NemotronHOmniForCausalLM + assert VLM_REGISTRY[config.model_type].vision_encoder_attr == "model.vision_model" + assert VLM_REGISTRY[config.model_type].language_model_attr == "model.language_model" + assert NemotronHOmniForCausalLM.is_hf_state_dict({"language_model.backbone.embeddings.weight": torch.empty(0)}) + assert NemotronHOmniForCausalLM.is_prime_state_dict({"model.language_model.embed_tokens.weight": torch.empty(0)}) + + +def test_merge_image_embeddings_replaces_only_image_tokens(): + input_ids = torch.tensor([[1, 18, 2, 18]]) + inputs_embeds = torch.zeros(1, 4, 3) + image_embeds = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + + merged = merge_image_embeddings(input_ids, inputs_embeds, image_embeds, image_token_id=18) + + torch.testing.assert_close(merged[0, 1], image_embeds[0]) + torch.testing.assert_close(merged[0, 3], image_embeds[1]) + torch.testing.assert_close(merged[0, [0, 2]], torch.zeros(2, 3)) + + +@pytest.mark.parametrize("num_image_embeddings", [1, 3]) +def test_merge_image_embeddings_rejects_token_count_mismatch( + num_image_embeddings, +): + with pytest.raises(ValueError, match="image token count"): + merge_image_embeddings( + torch.tensor([[18, 18]]), + torch.zeros(1, 2, 3), + torch.ones(num_image_embeddings, 3), + image_token_id=18, + ) + + +def test_merge_image_embeddings_rejects_hidden_size_mismatch(): + with pytest.raises(ValueError, match="hidden size"): + merge_image_embeddings( + torch.tensor([[18]]), + torch.zeros(1, 1, 3), + torch.ones(1, 4), + image_token_id=18, + ) + + +def test_nemotron_h_model_requires_exactly_one_embedding_input(): + from prime_rl.trainer.models.nemotron_h.modeling_nemotron_h import NemotronHModel + + model = NemotronHModel(_language_config()) + seq_lens = torch.tensor([2]) + + with pytest.raises(ValueError, match="[Ee]xactly one"): + model(seq_lens=seq_lens) + with pytest.raises(ValueError, match="[Ee]xactly one"): + model( + input_ids=torch.tensor([[1, 2]]), + inputs_embeds=torch.zeros(1, 2, 4), + seq_lens=seq_lens, + ) + + +def test_nemotron_h_model_accepts_remote_config_without_load_balance_coeff(): + from prime_rl.trainer.models.nemotron_h.modeling_nemotron_h import NemotronHModel + + config = _language_config() + config.layers_block_type = ["moe"] + del config.load_balance_coeff + + model = NemotronHModel(config) + + assert model.layers[0].mlp.load_balance_coeff is None + + +class RadioLayerScale(nn.Module): + def __init__(self, config): + super().__init__() + self.lambda1 = nn.Parameter(torch.full((4,), config.layerscale_value)) + + +class _FakeVisionModel(nn.Module): + def __init__(self, config): + super().__init__() + self.config = config + self.anchor = nn.Parameter(torch.ones(1)) + self.layer_scale = RadioLayerScale(config) + self.register_buffer( + "summary_idxs", + torch.tensor(config.summary_idxs), + ) + self.preprocessor_external = False + + def make_preprocessor_external(self): + self.preprocessor_external = True + + +class _FakeVisionProjector(nn.Module): + def __init__(self, config): + super().__init__() + + def forward(self, pixel_values, vision_model): + if isinstance(pixel_values, list): + return torch.cat([self.forward(item, vision_model) for item in pixel_values], dim=0) + if pixel_values.ndim == 4: + return pixel_values.new_ones(pixel_values.shape[0], 1, 4) * vision_model.anchor + return pixel_values.unsqueeze(0) * vision_model.anchor + + +def _language_config(): + return NemotronHConfig( + hidden_size=4, + vocab_size=32, + layers_block_type=[], + n_routed_experts=2, + num_experts_per_tok=1, + moe_intermediate_size=4, + moe_shared_expert_intermediate_size=4, + moe_latent_size=4, + ) + + +def _composite_config(): + config = PretrainedConfig() + config.model_type = "nemotron_h_omni" + config.llm_config = _language_config() + config.vision_config = SimpleNamespace( + layerscale_value=1.0, + summary_idxs=[0], + num_channels=3, + ) + config.img_context_token_id = 18 + config.force_image_size = 2 + config.patch_size = 2 + config.downsample_ratio = 0.5 + config.tie_word_embeddings = False + return config + + +def _remote_composite_config(model_path: str, revision: str | None): + config = _composite_config() + config._name_or_path = model_path + config._commit_hash = revision + config.auto_map = {"AutoModelForImageTextToText": "modeling_nemotron_h_omni.NemotronH_Omni_Reasoning_V3"} + return config + + +def test_remote_vision_loader_requires_explicit_trust(monkeypatch): + import prime_rl.trainer.models.nemotron_h_omni.modeling_nemotron_h_omni as modeling + + calls = [] + monkeypatch.setattr(modeling, "_load_remote_model_module", lambda *args: calls.append(args)) + + with pytest.raises(ValueError, match="trust_remote_code=True"): + _load_remote_vision_components( + _remote_composite_config("nvidia/model", "a" * 40), + trust_remote_code=False, + ) + + assert calls == [] + + +def test_remote_vision_loader_rejects_unpinned_remote_source(monkeypatch): + import prime_rl.trainer.models.nemotron_h_omni.modeling_nemotron_h_omni as modeling + + calls = [] + monkeypatch.setattr(modeling, "_load_remote_model_module", lambda *args: calls.append(args)) + + with pytest.raises(ValueError, match="immutable commit"): + _load_remote_vision_components( + _remote_composite_config("nvidia/model", None), + trust_remote_code=True, + ) + + assert calls == [] + + +def test_remote_vision_loader_accepts_trusted_local_checkpoint(monkeypatch, tmp_path): + import prime_rl.trainer.models.nemotron_h_omni.modeling_nemotron_h_omni as modeling + + model_path = tmp_path / "model" + model_path.mkdir() + module = SimpleNamespace( + RadioModel=_FakeVisionModel, + NemotronH_Omni_Reasoning_V3VisionProjector=_FakeVisionProjector, + ) + calls = [] + monkeypatch.setattr( + modeling, + "_load_remote_model_module", + lambda *args: calls.append(args) or module, + ) + + components = _load_remote_vision_components( + _remote_composite_config(str(model_path), None), + trust_remote_code=True, + ) + + assert components == (_FakeVisionModel, _FakeVisionProjector) + assert calls == [ + ( + "modeling_nemotron_h_omni.NemotronH_Omni_Reasoning_V3", + str(model_path), + None, + ) + ] + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("model_type", "nemotron_h", "model_type"), + ( + "auto_map", + {"AutoModelForImageTextToText": "modeling_other.UnexpectedModel"}, + "AutoModelForImageTextToText", + ), + ], +) +def test_remote_vision_loader_rejects_unexpected_checkpoint_contract(monkeypatch, field, value, message): + import prime_rl.trainer.models.nemotron_h_omni.modeling_nemotron_h_omni as modeling + + config = _remote_composite_config("nvidia/model", "a" * 40) + setattr(config, field, value) + calls = [] + monkeypatch.setattr(modeling, "_load_remote_model_module", lambda *args: calls.append(args)) + + with pytest.raises(ValueError, match=message): + _load_remote_vision_components(config, trust_remote_code=True) + + assert calls == [] + + +def test_custom_loading_paths_propagate_explicit_trust(monkeypatch): + from prime_rl.trainer.models.nemotron_h.modeling_nemotron_h import NemotronHPreTrainedModel + + config_calls = [] + pretrained_calls = [] + + def fake_from_config(cls, config, **kwargs): + config_calls.append((config, kwargs)) + return "from-config" + + def fake_from_pretrained(cls, *args, **kwargs): + pretrained_calls.append((args, kwargs)) + return "from-pretrained" + + monkeypatch.setattr(NemotronHOmniForCausalLM, "_from_config", classmethod(fake_from_config)) + monkeypatch.setattr(NemotronHPreTrainedModel, "from_pretrained", classmethod(fake_from_pretrained)) + + config = _remote_composite_config("nvidia/model", "a" * 40) + assert NemotronHOmniForCausalLM.from_config(config, trust_remote_code=True) == "from-config" + assert ( + NemotronHOmniForCausalLM.from_pretrained("nvidia/model", config=config, trust_remote_code=True) + == "from-pretrained" + ) + assert config_calls == [(config, {"_prime_trust_remote_code": True})] + assert pretrained_calls == [ + ( + ("nvidia/model",), + { + "config": config, + "trust_remote_code": True, + "_prime_trust_remote_code": True, + }, + ) + ] + + +@pytest.fixture +def composite_model(monkeypatch): + import prime_rl.trainer.models.nemotron_h_omni.modeling_nemotron_h_omni as modeling + + monkeypatch.setattr( + modeling, + "_load_remote_vision_components", + lambda config, *, trust_remote_code: (_FakeVisionModel, _FakeVisionProjector), + ) + return NemotronHOmniForCausalLM(_composite_config(), _prime_trust_remote_code=True) + + +def test_composite_text_only_forward(composite_model): + input_ids = torch.tensor([[1, 2, 3]]) + seq_lens = torch.tensor([3]) + + output = composite_model(input_ids=input_ids, seq_lens=seq_lens) + + output["logits"].sum().backward() + + assert composite_model.supports_packed_multimodal_training + assert composite_model.model.vision_model.preprocessor_external + assert composite_model.model.vision_model.anchor.grad is not None + torch.testing.assert_close( + composite_model.model.vision_model.anchor.grad, + torch.zeros(1), + ) + assert output["logits"].shape == (1, 3, 32) + + +def test_composite_minimal_image_forward_without_weights(composite_model): + input_ids = torch.tensor([[1, 18, 3]]) + inject_prime_lm_head(composite_model) + output = composite_model( + input_ids=input_ids, + pixel_values=torch.full((1, 3, 2, 2), 0.5), + mm_token_type_ids=torch.tensor([[0, 1, 0]]), + seq_lens=torch.tensor([3]), + ) + + assert output["logits"].shape == (1, input_ids.shape[1], 32) + + +def test_trainer_forward_accepts_renderer_nemotron_image_fields(composite_model): + input_ids = torch.tensor([[1, 18, 3]]) + position_ids = torch.arange(input_ids.shape[1]).unsqueeze(0) + inject_prime_lm_head(composite_model) + + output = trainer_forward( + composite_model, + input_ids, + position_ids, + seq_lens=torch.tensor([3]), + mm_kwargs={ + "pixel_values": torch.zeros(1, 3, 2, 2), + "imgs_sizes": torch.tensor([[2, 2]]), + "num_tokens": torch.tensor([1]), + "num_patches": torch.tensor([1]), + }, + mm_token_type_ids=torch.tensor([[0, 1, 0]]), + ) + + assert output["logits"].shape == (1, input_ids.shape[1], 32) + + +def test_composite_rejects_renderer_image_markers_that_disagree_with_tokens(composite_model): + inject_prime_lm_head(composite_model) + + with pytest.raises(ValueError, match="image markers"): + composite_model( + input_ids=torch.tensor([[1, 18, 3]]), + pixel_values=torch.zeros(1, 3, 2, 2), + imgs_sizes=torch.tensor([[2, 2]]), + num_tokens=torch.tensor([1]), + num_patches=torch.tensor([1]), + mm_token_type_ids=torch.tensor([[0, 0, 0]]), + seq_lens=torch.tensor([3]), + ) + + +def test_composite_reconstructs_dynamic_resolution_images_from_flat_wire(composite_model): + inject_prime_lm_head(composite_model) + pixels = torch.cat( + ( + torch.zeros(1, 3, 2, 3).reshape(-1), + torch.ones(1, 3, 3, 2).reshape(-1), + ) + ) + + output = composite_model( + input_ids=torch.tensor([[18, 2, 18]]), + pixel_values=pixels, + imgs_sizes=torch.tensor([[2, 3], [3, 2]]), + num_tokens=torch.tensor([1, 1]), + num_patches=torch.tensor([1, 1]), + mm_token_type_ids=torch.tensor([[1, 0, 1]]), + seq_lens=torch.tensor([3]), + ) + + assert output["logits"].shape == (1, 3, 32) + + +@pytest.mark.parametrize( + ("pixel_values", "imgs_sizes", "num_tokens", "message"), + [ + (torch.arange(11, dtype=torch.float32), torch.tensor([[2, 2]]), torch.tensor([1]), "flat pixel_values"), + (torch.arange(12, dtype=torch.float32), torch.tensor([[0, 2]]), torch.tensor([1]), "positive"), + (torch.arange(12, dtype=torch.int64), torch.tensor([[2, 2]]), torch.tensor([1]), "floating-point"), + (torch.arange(12, dtype=torch.float32), torch.tensor([[2, 2]]), torch.tensor([0]), "positive"), + ], +) +def test_composite_rejects_malformed_flat_image_wire( + composite_model, + pixel_values, + imgs_sizes, + num_tokens, + message, +): + with pytest.raises(ValueError, match=message): + composite_model( + input_ids=torch.tensor([[1, 18, 3]]), + pixel_values=pixel_values, + imgs_sizes=imgs_sizes, + num_tokens=num_tokens, + num_patches=torch.tensor([1]), + mm_token_type_ids=torch.tensor([[0, 1, 0]]), + seq_lens=torch.tensor([3]), + ) + + +def test_composite_merges_images_before_context_parallel_sharding(composite_model, monkeypatch): + import prime_rl.trainer.models.nemotron_h_omni.modeling_nemotron_h_omni as modeling + + events = [] + original_merge = modeling.merge_image_embeddings + + def record_merge(*args, **kwargs): + events.append("merge") + return original_merge(*args, **kwargs) + + def record_shard(tensor, **kwargs): + events.append("shard") + return tensor + + monkeypatch.setattr(modeling, "merge_image_embeddings", record_merge) + monkeypatch.setattr(modeling, "shard_for_cp", record_shard) + monkeypatch.setattr(modeling, "shard_position_ids_for_cp", lambda position_ids, **kwargs: position_ids) + monkeypatch.setattr(modeling, "setup_cp_attention_params", lambda *args, **kwargs: None) + composite_model.model.cp_context = modeling.CPContext( + cp_group=object(), + cp_rank=0, + cp_world_size=2, + cp_style="ulysses", + ) + + output = composite_model.model( + input_ids=torch.tensor([[1, 18, 3]]), + position_ids=torch.arange(3).unsqueeze(0), + pixel_values=torch.zeros(1, 3, 2, 2), + imgs_sizes=torch.tensor([[2, 2]]), + num_tokens=torch.tensor([1]), + num_patches=torch.tensor([1]), + mm_token_type_ids=torch.tensor([[0, 1, 0]]), + routed_experts=torch.zeros(1, 3, 0, 0, dtype=torch.long), + seq_lens=torch.tensor([3]), + ) + + assert output.last_hidden_state.shape == (1, 3, 4) + assert events == ["merge", "shard", "shard"] + + +def test_composite_handles_checkpoint_missing_radio_state_after_to_empty( + composite_model, +): + assert isinstance(composite_model.model.vision_model.layer_scale, nn.Identity) + composite_model.to_empty(device="cpu") + composite_model.init_buffers_post_meta() + + torch.testing.assert_close( + composite_model.model.vision_model.summary_idxs, + torch.tensor([0]), + ) From b514c4894b4d9da532e863135d7d887d1ec94f53 Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Wed, 30 Sep 2026 08:43:40 -0700 Subject: [PATCH 3/4] fix(nemotron): register composite architecture for text training --- src/prime_rl/utils/vlm.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/src/prime_rl/utils/vlm.py b/src/prime_rl/utils/vlm.py index c40c93cd9a8..583e62b2999 100644 --- a/src/prime_rl/utils/vlm.py +++ b/src/prime_rl/utils/vlm.py @@ -30,6 +30,10 @@ class VLMModelInfo: "qwen3_5_moe": VLMModelInfo(vision_encoder_attr="model.visual", language_model_attr="model.language_model"), "qwen3_vl_moe": VLMModelInfo(vision_encoder_attr="model.visual", language_model_attr="model.language_model"), "gemma4": VLMModelInfo(vision_encoder_attr="model.vision_tower", language_model_attr="model.language_model"), + "nemotron_h_omni": VLMModelInfo( + vision_encoder_attr="model.vision_model", + language_model_attr="model.language_model", + ), } # Text-only default From 308e35713992c80e26d3565b17f0307535963a9e Mon Sep 17 00:00:00 2001 From: Biswa Panda Date: Wed, 30 Sep 2026 08:05:50 -0700 Subject: [PATCH 4/4] feat(multimodal): enable Nemotron-H Omni training --- src/prime_rl/multimodal/nemotron_h_omni.py | 36 ++++++++++++++++++ src/prime_rl/multimodal/registry.py | 3 +- src/prime_rl/trainer/model.py | 9 ++++- tests/unit/train/test_model_forward.py | 44 ++++++++++++++++++++++ tests/unit/train/test_multimodal.py | 43 +++++++++++++++++++++ 5 files changed, 133 insertions(+), 2 deletions(-) create mode 100644 src/prime_rl/multimodal/nemotron_h_omni.py diff --git a/src/prime_rl/multimodal/nemotron_h_omni.py b/src/prime_rl/multimodal/nemotron_h_omni.py new file mode 100644 index 00000000000..5213cc7ee3f --- /dev/null +++ b/src/prime_rl/multimodal/nemotron_h_omni.py @@ -0,0 +1,36 @@ +from __future__ import annotations + +from typing import Any + +from PIL.Image import Image + +from prime_rl.multimodal.base import ForwardPolicy, MaterializedMM, required_tensors + + +class NemotronHOmniAdapter: + model_types = frozenset({"nemotron_h_omni"}) + forward_policy = ForwardPolicy( + requires_mm_token_type_ids=True, + defer_context_parallelism=True, + ) + + def materialize( + self, + image_processor: Any, + images: list[Image], + placeholder_lengths: list[int], + ) -> MaterializedMM: + values = image_processor(images=images, return_tensors="pt") + pixel_values = values.get("pixel_values") + if isinstance(pixel_values, list): + raise ValueError("Nemotron images in one batch must have the same processed shape") + kwargs = required_tensors( + values, + ("pixel_values", "imgs_sizes", "num_tokens", "num_patches"), + ) + lengths = [int(length) for length in kwargs["num_tokens"].reshape(-1)] + if lengths != placeholder_lengths: + raise ValueError( + f"Nemotron image placeholder lengths differ from vLLM: expected {placeholder_lengths}, got {lengths}" + ) + return MaterializedMM(kwargs=kwargs, forward_policy=self.forward_policy) diff --git a/src/prime_rl/multimodal/registry.py b/src/prime_rl/multimodal/registry.py index d86bc96bd2e..85bec7dabf7 100644 --- a/src/prime_rl/multimodal/registry.py +++ b/src/prime_rl/multimodal/registry.py @@ -2,9 +2,10 @@ from prime_rl.multimodal.base import MultimodalAdapter from prime_rl.multimodal.kimi_k25 import KimiK25Adapter +from prime_rl.multimodal.nemotron_h_omni import NemotronHOmniAdapter from prime_rl.multimodal.qwen_vl import QwenVLAdapter -_ADAPTERS = (QwenVLAdapter(), KimiK25Adapter()) +_ADAPTERS = (QwenVLAdapter(), KimiK25Adapter(), NemotronHOmniAdapter()) _BY_MODEL_TYPE: dict[str, MultimodalAdapter] = { model_type: adapter for adapter in _ADAPTERS for model_type in adapter.model_types } diff --git a/src/prime_rl/trainer/model.py b/src/prime_rl/trainer/model.py index b2115593a8e..f3a10a839b7 100644 --- a/src/prime_rl/trainer/model.py +++ b/src/prime_rl/trainer/model.py @@ -537,7 +537,14 @@ def shard_placement_fn(parameter: nn.Parameter) -> Shard | None: if vision_encoder is None: raise ValueError(f"VLM model {config.name} has no recognized vision encoder") - fully_shard(vision_encoder, mesh=hsdp_mesh, **fsdp_config) + retain_frozen_vision_weights = config.vlm.freeze_vision_encoder and getattr( + model, "retain_frozen_vision_weights_after_forward", False + ) + vision_fsdp_config = { + **fsdp_config, + "reshard_after_forward": config.reshard_after_forward and not retain_frozen_vision_weights, + } + fully_shard(vision_encoder, mesh=hsdp_mesh, **vision_fsdp_config) get_logger().info(f"Applied FSDP to vision encoder (frozen={config.vlm.freeze_vision_encoder})") language_model = get_language_model(model, override=config.vlm.language_model_attr if is_vlm_training else None) diff --git a/tests/unit/train/test_model_forward.py b/tests/unit/train/test_model_forward.py index 853eef3ebc9..c1fe018229c 100644 --- a/tests/unit/train/test_model_forward.py +++ b/tests/unit/train/test_model_forward.py @@ -1,8 +1,10 @@ from types import SimpleNamespace +import pytest import torch import torch.nn as nn +import prime_rl.trainer.model as trainer_model from prime_rl.multimodal.qwen_vl import QwenVLAdapter from prime_rl.trainer.model import forward @@ -85,3 +87,45 @@ def test_forward_keeps_position_ids_for_non_mrope_vlm(): assert model.kwargs is not None torch.testing.assert_close(model.kwargs["position_ids"], position_ids) + + +@pytest.mark.parametrize( + "freeze_vision_encoder,retain_frozen_vision_weights_after_forward,expected", + [(True, True, False), (True, False, True), (False, True, True)], +) +def test_setup_fsdp_keeps_frozen_vision_encoder_unsharded_after_forward( + monkeypatch: pytest.MonkeyPatch, + freeze_vision_encoder: bool, + retain_frozen_vision_weights_after_forward: bool, + expected: bool, +): + vision_encoder = object() + language_model = SimpleNamespace(layers=[], embed_tokens=object(), norm=object()) + model = SimpleNamespace( + config=SimpleNamespace(tie_word_embeddings=False), + lm_head=object(), + retain_frozen_vision_weights_after_forward=retain_frozen_vision_weights_after_forward, + ) + config = SimpleNamespace( + reduce_dtype="bfloat16", + fsdp_cpu_offload=False, + fusions=SimpleNamespace(shard_fused_on_dim1=False), + reshard_after_forward=True, + vlm=SimpleNamespace( + vision_encoder_attr="vision", + language_model_attr="language", + freeze_vision_encoder=freeze_vision_encoder, + ), + moe_router_dtype="bfloat16", + ) + parallel_dims = SimpleNamespace(ep_enabled=False, get_mesh=lambda _: object()) + shard_calls = [] + + monkeypatch.setattr(trainer_model, "get_vision_encoder", lambda *_args, **_kwargs: vision_encoder) + monkeypatch.setattr(trainer_model, "get_language_model", lambda *_args, **_kwargs: language_model) + monkeypatch.setattr(trainer_model, "fully_shard", lambda target, **kwargs: shard_calls.append((target, kwargs))) + + trainer_model.setup_fsdp(model, config, parallel_dims) + + vision_kwargs = next(kwargs for target, kwargs in shard_calls if target is vision_encoder) + assert vision_kwargs["reshard_after_forward"] is expected diff --git a/tests/unit/train/test_multimodal.py b/tests/unit/train/test_multimodal.py index a7dd5694ee6..7d918985159 100644 --- a/tests/unit/train/test_multimodal.py +++ b/tests/unit/train/test_multimodal.py @@ -5,6 +5,7 @@ from prime_rl.multimodal import get_multimodal_adapter from prime_rl.multimodal.kimi_k25 import KimiK25Adapter +from prime_rl.multimodal.nemotron_h_omni import NemotronHOmniAdapter from prime_rl.multimodal.qwen_vl import QwenVLAdapter from prime_rl.trainer.multimodal import materialize_mm_refs from prime_rl.transports.batch import MMImageRef, MMRefs @@ -54,3 +55,45 @@ def preprocess(self, media, *, return_tensors): assert set(materialized.kwargs) == {"pixel_values", "grid_thws"} assert materialized.forward_policy == KimiK25Adapter.forward_policy + + +def test_nemotron_adapter_materializes_and_validates_expansion(): + class ImageProcessor: + def __call__(self, *, images, return_tensors): + assert len(images) == 1 and images[0].mode == "RGB" and return_tensors == "pt" + return { + "pixel_values": torch.ones(1, 3, 16, 16), + "imgs_sizes": torch.tensor([[16, 16]]), + "num_tokens": torch.tensor([4]), + "num_patches": torch.tensor([1]), + } + + processor = SimpleNamespace(image_processor=ImageProcessor()) + adapter = get_multimodal_adapter("nemotron_h_omni") + materialized = materialize_mm_refs(_refs(4), processor, adapter) + + assert set(materialized.kwargs) == {"pixel_values", "imgs_sizes", "num_tokens", "num_patches"} + assert materialized.forward_policy == NemotronHOmniAdapter.forward_policy + with pytest.raises(ValueError, match="placeholder lengths differ"): + materialize_mm_refs(_refs(3), processor, adapter) + + +def test_nemotron_adapter_rejects_variable_image_shapes(): + class ImageProcessor: + def __call__(self, *, images, return_tensors): + return { + "pixel_values": [torch.ones(3, 16, 16), torch.ones(3, 32, 16)], + "imgs_sizes": [(16, 16), (32, 16)], + "num_tokens": [4, 8], + "num_patches": [1, 1], + } + + processor = SimpleNamespace(image_processor=ImageProcessor()) + refs = MMRefs( + images=[ + MMImageRef(url=_IMAGE_URL, offset=1, length=4), + MMImageRef(url=_IMAGE_URL, offset=5, length=8), + ] + ) + with pytest.raises(ValueError, match="same processed shape"): + materialize_mm_refs(refs, processor, get_multimodal_adapter("nemotron_h_omni"))