feat(multimodal): enable Nemotron-H Omni VLM training

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)

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
 }

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)

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,
 }
 
 

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",
+]

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),

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",
+]

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"]