Add NVIDIA Nemotron Parse 2.0 support

README.md

@@ -29,7 +29,7 @@ multi-component export for pipelines.
 |---|---|
 | **Text Generation** | Llama 2/3/4, Mistral, Qwen 2/2.5/3/3.5/3.6, Phi-3/3.5, Gemma 1/2/3/4, Granite, GPT-2, OPT, OLMo, SmolLM3, and many more |
 | **Mixture of Experts** | PhiMoE, GPTOSS, Mixtral, OLMoE, DeepSeek-V2/V3, Qwen2-MoE, Qwen3-MoE, Qwen3-Next, GLM-4-MoE, Arctic, DBRX, Jamba |
-| **Multimodal** | Gemma 3/4, Phi-4MM (vision + audio + LoRA), LLaVA, InternVL2, Qwen2.5-VL, Qwen3-VL, Qwen3.5/3.6-VL, Pixtral |
+| **Multimodal** | Gemma 3/4, Phi-4MM (vision + audio + LoRA), Nemotron Parse, LLaVA, InternVL2, Qwen2.5-VL, Qwen3-VL, Qwen3.5/3.6-VL, Pixtral |
 | **Encoder-only** | BERT, RoBERTa, ALBERT, DeBERTa, DistilBERT, ELECTRA, XLNet |
 | **Encoder-Decoder** | BART, T5/mT5, Marian, M2M-100, Pegasus, BigBird-Pegasus |
 | **Speech-to-Text** | Whisper, FastConformer-RNNT, FunASR, Qwen3-ASR, SenseVoice |

scripts/generate_golden.py

@@ -511,6 +511,86 @@ def _generate_vision_language(case: TestCase, json_path: Path, device: str) -> N
         )
 
 
+def _generate_image_to_text(case: TestCase, json_path: Path, device: str) -> None:
+    """Generate Nemotron Parse-style image-to-text reference data."""
+    import torch
+    from PIL import Image
+    from transformers import AutoModel, AutoProcessor
+
+    from mobius._testing.golden import save_generation_json, save_golden_ref
+
+    dtype_map = {
+        "float32": torch.float32,
+        "float16": torch.float16,
+        "bfloat16": torch.bfloat16,
+    }
+    model = AutoModel.from_pretrained(
+        case.model_id,
+        revision=case.revision,
+        torch_dtype=dtype_map[case.dtype],
+        trust_remote_code=case.trust_remote_code,
+    ).to(device)
+    model.eval()
+    processor = AutoProcessor.from_pretrained(
+        case.model_id,
+        revision=case.revision,
+        trust_remote_code=case.trust_remote_code,
+    )
+    images = [
+        Image.open(Path("testdata") / image_path).convert("RGB") for image_path in case.images
+    ]
+    processed = processor(
+        images=images,
+        text=case.decoder_prompt,
+        return_tensors="pt",
+        add_special_tokens=False,
+    ).to(device)
+    decoder_input_ids = processed["input_ids"]
+
+    with torch.no_grad():
+        outputs = model(
+            pixel_values=processed["pixel_values"],
+            decoder_input_ids=decoder_input_ids,
+        )
+    last_logits = outputs.logits[0, -1].float().cpu().numpy()
+    golden = _extract_logits_golden(last_logits)
+    input_ids = decoder_input_ids.cpu().numpy()
+
+    generated_ids = None
+    if "L5" in case.level:
+        with torch.no_grad():
+            generated = model.generate(
+                pixel_values=processed["pixel_values"],
+                decoder_input_ids=decoder_input_ids,
+                max_new_tokens=case.generation_params.get("max_new_tokens", 20),
+                do_sample=False,
+            )
+        generated_ids = generated[0, input_ids.shape[1] :].cpu().numpy()
+
+    save_golden_ref(
+        json_path,
+        top1_id=golden["top1_id"],
+        top2_id=golden["top2_id"],
+        top10_ids=golden["top10_ids"],
+        top10_logits=golden["top10_logits"],
+        logits_summary=golden["logits_summary"],
+        input_ids=input_ids,
+    )
+
+    if generated_ids is not None:
+        gen_path = json_path.with_name(json_path.stem + "_generation.json")
+        save_generation_json(
+            gen_path,
+            model_id=case.model_id,
+            prompt=case.decoder_prompt,
+            generated_tokens=generated_ids.tolist(),
+            generated_text=processor.decode(
+                generated_ids.tolist(),
+                skip_special_tokens=False,
+            ),
+        )
+
+
 def _generate_speech_to_text(case: TestCase, json_path: Path, device: str) -> None:
     """Generate golden data for a speech-to-text (Whisper) model."""
     import librosa
@@ -1646,6 +1726,7 @@ def _hook(_module, _args, kwargs, output):
     "feature-extraction": _generate_encoder,
     "seq2seq": _generate_seq2seq,
     "image-text-to-text": _generate_vision_language,
+    "image-to-text": _generate_image_to_text,
     "image-classification": _generate_image_classification,
     "speech-to-text": _generate_speech_to_text,
     "speech-language": _generate_speech_language,

src/mobius/_configs/__init__.py

@@ -45,6 +45,7 @@
     MMSConfig,
     NanoChatConfig,
     NemotronHConfig,
+    NemotronParseConfig,
     Qwen35MtpConfig,
     Sam2Config,
     SegformerConfig,
@@ -107,6 +108,7 @@
     "MllamaConfig",
     "MMSConfig",
     "NanoChatConfig",
+    "NemotronParseConfig",
     "NemotronHConfig",
     "QuantizationConfig",
     "Qwen35MtpConfig",

src/mobius/_configs/_base.py

@@ -1267,6 +1267,101 @@ class VisionLanguageConfig(CausalLMConfig):
     """
 
 
+@dataclasses.dataclass
+class NemotronParseConfig(ArchitectureConfig):
+    """Configuration for NVIDIA Nemotron Parse image-to-text models."""
+
+    image_height: int = 2048
+    image_width: int = 1664
+    vision_max_grid_size: int = 128
+    num_summary_tokens: int = 3
+    decoder_start_token_id: int = 2
+    scale_embedding: bool = True
+    add_final_layer_norm: bool = True
+
+    @classmethod
+    def from_transformers(cls, config, parent_config=None) -> NemotronParseConfig:
+        del parent_config
+        import types
+
+        def _namespace(value):
+            if isinstance(value, dict):
+                return types.SimpleNamespace(
+                    **{key: _namespace(item) for key, item in value.items()}
+                )
+            if isinstance(value, list):
+                return [_namespace(item) for item in value]
+            return value
+
+        decoder = _namespace(getattr(config, "decoder", None))
+        if decoder is None:
+            raise ValueError("Nemotron Parse config is missing its decoder sub-config")
+        base = ArchitectureConfig.from_transformers(decoder, parent_config=config)
+        fields = _shallow_fields(base)
+        num_attention_heads = int(
+            getattr(decoder, "decoder_attention_heads", None)
+            or getattr(decoder, "num_attention_heads", fields["num_attention_heads"])
+        )
+        hidden_size = int(fields["hidden_size"])
+        if hidden_size % num_attention_heads:
+            raise ValueError(
+                "Nemotron Parse decoder hidden size must be divisible by its attention heads"
+            )
+        fields.update(
+            num_attention_heads=num_attention_heads,
+            num_key_value_heads=num_attention_heads,
+            head_dim=hidden_size // num_attention_heads,
+        )
+
+        raw_image_size = getattr(config, "image_size", (2048, 1664))
+        if isinstance(raw_image_size, int):
+            image_height = image_width = raw_image_size
+        else:
+            image_height, image_width = (int(raw_image_size[0]), int(raw_image_size[1]))
+
+        encoder = _namespace(getattr(config, "encoder", None))
+        patch_size = int(getattr(encoder, "patch_size", 16))
+        max_resolution = int(
+            getattr(encoder, "max_resolution", max(image_height, image_width))
+        )
+        fields.update(
+            model_type="nemotron_parse",
+            bos_token_id=getattr(config, "bos_token_id", fields.get("bos_token_id")),
+            eos_token_id=getattr(config, "eos_token_id", fields.get("eos_token_id")),
+            pad_token_id=getattr(config, "pad_token_id", fields["pad_token_id"]),
+            tie_word_embeddings=getattr(config, "tie_word_embeddings", True),
+            max_position_embeddings=(
+                getattr(config, "max_sequence_length", None)
+                or fields["max_position_embeddings"]
+            ),
+            vision=VisionConfig(
+                hidden_size=1280,
+                intermediate_size=5120,
+                num_hidden_layers=32,
+                num_attention_heads=16,
+                image_size=max_resolution,
+                patch_size=patch_size,
+                norm_eps=1e-6,
+                model_type="radio_v2.5-h",
+                in_channels=3,
+            ),
+        )
+        resolved_dtype = _resolve_dtype(config)
+        if resolved_dtype is not None:
+            fields["dtype"] = resolved_dtype
+
+        return cls(
+            **fields,
+            image_height=image_height,
+            image_width=image_width,
+            vision_max_grid_size=max_resolution // patch_size,
+            num_summary_tokens=3,
+            decoder_start_token_id=int(getattr(config, "decoder_start_token_id", 2)),
+            scale_embedding=bool(getattr(decoder, "scale_embedding", True)),
+            add_final_layer_norm=bool(getattr(decoder, "add_final_layer_norm", True)),
+        )
+
+
 # ---------------------------------------------------------------------------
 # Model-family subclasses — add model-specific fields
 # ---------------------------------------------------------------------------

src/mobius/_configs/_base_test.py

@@ -0,0 +1,42 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+"""Tests for architecture-specific configuration extraction."""
+
+from __future__ import annotations
+
+import types
+
+from mobius._configs import NemotronParseConfig
+
+
+def test_nemotron_parse_maps_raw_mbart_decoder_attention_heads():
+    """The Hub's non-trusted config exposes MBART's decoder-specific aliases."""
+    config = types.SimpleNamespace(
+        model_type="nemotron_parse",
+        decoder={
+            "model_type": "nemotron_parse_text",
+            "d_model": 1024,
+            "decoder_attention_heads": 16,
+            "decoder_ffn_dim": 4096,
+            "decoder_layers": 10,
+            "num_hidden_layers": 12,
+            "vocab_size": 72256,
+            "pad_token_id": 1,
+        },
+        encoder={"patch_size": 16, "max_resolution": 2048},
+        image_size=[2048, 1664],
+        max_sequence_length=9000,
+        bos_token_id=0,
+        eos_token_id=2,
+        pad_token_id=1,
+        tie_word_embeddings=True,
+        decoder_start_token_id=2,
+    )
+
+    extracted = NemotronParseConfig.from_transformers(config)
+
+    assert extracted.num_attention_heads == 16
+    assert extracted.num_key_value_heads == 16
+    assert extracted.head_dim == 64
+    assert extracted.num_decoder_layers == 10

src/mobius/_registry.py

@@ -28,6 +28,7 @@
     Gemma4AssistantConfig,
     Gemma4Config,
     MMSConfig,
+    NemotronParseConfig,
     WhisperConfig,
 )
 from mobius.models import (
@@ -73,6 +74,7 @@
     MoECausalLMModel,
     NanoChatCausalLMModel,
     NemotronCausalLMModel,
+    NemotronParseForConditionalGeneration,
     OLMo2CausalLMModel,
     OLMoCausalLMModel,
     Phi3CausalLMModel,
@@ -564,6 +566,11 @@ def _detect_fallback_registration(hf_config) -> ModelRegistration | None:
     "bamba": ModelRegistration(BambaCausalLMModel),
     "jamba": ModelRegistration(JambaCausalLMModel),
     "nemotron_h": ModelRegistration(NemotronHCausalLMModel),
+    "nemotron_parse": ModelRegistration(
+        NemotronParseForConditionalGeneration,
+        task="vision-encoder-decoder",
+        config_class=NemotronParseConfig,
+    ),
     # --- Hybrid linear-attention ---
     "longcat_flash": ModelRegistration(LongcatFlashCausalLMModel),
     # --- Multimodal ---
@@ -863,6 +870,7 @@ def _create_default_registry() -> ModelRegistry:
     "csm": "sesame/csm-1b",
     "evolla": "westlake-repl/Evolla-10B-hf",
     "nemotron_h": "nvidia/NVIDIA-Nemotron-3-Nano-4B-BF16",
+    "nemotron_parse": "nvidia/NVIDIA-Nemotron-Parse-2.0",
     "open-llama": "openlm-research/open_llama_3b",
     "persimmon": "adept/persimmon-8b-base",
     "shieldgemma2": "google/shieldgemma-2b",

src/mobius/components/__init__.py

@@ -56,6 +56,7 @@
     "PostNormDecoderLayer",
     "QuantizedEmbedding",
     "QuantizedLinear",
+    "RadioVisionModel",
     "RMSNorm",
     "SelectiveScan",
     "SiLU",
@@ -251,6 +252,7 @@
 from mobius.components._qwen25_vl_vision import (
     Qwen25VLVisionRotaryEmbedding as Qwen25VLVisionRotaryEmbedding,
 )
+from mobius.components._radio_vision import RadioVisionModel
 from mobius.components._rms_norm import (
     GatedRMSNorm,
     OffsetRMSNorm,