diff --git a/docs/design/mm_processing.md b/docs/design/mm_processing.md index a267aee38e..20226171a9 100644 --- a/docs/design/mm_processing.md +++ b/docs/design/mm_processing.md @@ -90,6 +90,7 @@ GPU-side fusion is controlled by `multimodal_config.mm_device_do_normalize`: | `llama-nemotron-vl-embed` | `LlamaNemotronVLForEmbedding` | `nvidia/llama-nemotron-embed-vl-1b-v2` | | `llama-nemotron-vl-rerank` | `LlamaNemotronVLForSequenceClassification` | `nvidia/llama-nemotron-rerank-vl-1b-v2` | | `mistral3` | `Mistral3ForConditionalGeneration` | `mistralai/Mistral-Small-3.2-24B-Instruct-2506`, etc. | +| `nemotron-parse` | `NemotronParseForConditionalGeneration` | `nvidia/NVIDIA-Nemotron-Parse-2.0`, etc. | | `qwen2-vl` | `Qwen2VLForConditionalGeneration` | `Qwen/Qwen2-VL-2B-Instruct`, etc. | | `qwen2.5-vl` | `Qwen2_5_VLForConditionalGeneration` | `Qwen/Qwen2.5-VL-3B-Instruct`, etc. | | `qwen3-vl` | `Qwen3VLForConditionalGeneration` | `Qwen/Qwen3-VL-4B-Instruct`, etc. | diff --git a/tests/models/multimodal/processing/test_nemotron_parse.py b/tests/models/multimodal/processing/test_nemotron_parse.py new file mode 100644 index 0000000000..019c06eb15 --- /dev/null +++ b/tests/models/multimodal/processing/test_nemotron_parse.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for Nemotron-Parse's device-side normalisation path.""" + +import numpy as np +import pytest +import torch +from PIL import Image + +from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.platforms import current_platform + +from ...utils import build_model_context + +_MODEL_ID = "nvidia/NVIDIA-Nemotron-Parse-2.0" + + +def _noise_image(width: int, height: int) -> Image.Image: + rng = np.random.default_rng(0) + return Image.fromarray(rng.integers(0, 256, (height, width, 3), dtype=np.uint8)) + + +@pytest.mark.usefixtures("default_vllm_config") +@pytest.mark.parametrize(("width", "height"), [(1240, 1754), (310, 470)]) +def test_mm_device_do_normalize(width: int, height: int): + device = current_platform.device_type + ctx = build_model_context(_MODEL_ID, limit_mm_per_prompt={"image": 1}) + assert ctx.model_config.multimodal_config.mm_device_do_normalize + + images = [_noise_image(width, height)] + prompt = [0] + + ctx.model_config.multimodal_config.mm_device_do_normalize = False + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) + mm_items = processor.info.parse_mm_data({"image": images}) + normalized_values = processor(prompt, mm_items=mm_items)["mm_kwargs"].get_data()[ + "pixel_values" + ] + + ctx.model_config.multimodal_config.mm_device_do_normalize = True + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) + raw_values = processor(prompt, mm_items=mm_items)["mm_kwargs"].get_data()[ + "pixel_values" + ] + assert raw_values.dtype == torch.uint8 + assert raw_values.shape == normalized_values.shape + + input_norm = build_mm_input_norm(ctx.model_config).to(device) + output = input_norm(raw_values.flatten(1).to(device), normalized_values.dtype) + torch.testing.assert_close( + output, normalized_values.flatten(1).to(device), rtol=1.6e-2, atol=1e-5 + ) + + +@pytest.mark.usefixtures("default_vllm_config") +@pytest.mark.parametrize("mm_device_do_normalize", [False, True]) +@pytest.mark.parametrize("layout", ["ndarray_hwc", "tensor_hwc", "tensor_chw"]) +def test_array_image_inputs(mm_device_do_normalize: bool, layout: str): + ctx = build_model_context(_MODEL_ID, limit_mm_per_prompt={"image": 1}) + ctx.model_config.multimodal_config.mm_device_do_normalize = mm_device_do_normalize + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) + + image = _noise_image(310, 470) + array = np.asarray(image) + item = { + "ndarray_hwc": array, + "tensor_hwc": torch.from_numpy(array), + "tensor_chw": torch.from_numpy(array).permute(2, 0, 1), + }[layout] + + def pixel_values(img: object) -> torch.Tensor: + mm_items = processor.info.parse_mm_data({"image": [img]}) + return processor([0], mm_items=mm_items)["mm_kwargs"].get_data()["pixel_values"] + + torch.testing.assert_close(pixel_values(item), pixel_values(image), rtol=0, atol=0) + + +@pytest.mark.usefixtures("default_vllm_config") +def test_explicit_kwargs_without_device_normalize(): + ctx = build_model_context( + _MODEL_ID, + limit_mm_per_prompt={"image": 1}, + mm_processor_kwargs={"do_rescale": False, "do_normalize": False}, + ) + ctx.model_config.multimodal_config.mm_device_do_normalize = False + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) + + mm_items = processor.info.parse_mm_data({"image": [_noise_image(310, 470)]}) + values = processor([0], mm_items=mm_items)["mm_kwargs"].get_data()["pixel_values"] + # The remote processor ignores these kwargs and still normalises on the host. + assert values.is_floating_point() diff --git a/tests/models/registry.py b/tests/models/registry.py index 2bf035d06f..c0568becd2 100644 --- a/tests/models/registry.py +++ b/tests/models/registry.py @@ -1442,7 +1442,9 @@ _MULTIMODAL_EXAMPLE_MODELS = { "CohereLabs/cohere-transcribe-03-2026", trust_remote_code=True ), "NemotronParseForConditionalGeneration": _HfExamplesInfo( - "nvidia/NVIDIA-Nemotron-Parse-v1.2", trust_remote_code=True + "nvidia/NVIDIA-Nemotron-Parse-v1.2", + extras={"2.0": "nvidia/NVIDIA-Nemotron-Parse-2.0"}, + trust_remote_code=True, ), "WhisperForConditionalGeneration": _HfExamplesInfo( "openai/whisper-large-v3-turbo", diff --git a/vllm/model_executor/models/nemotron_parse.py b/vllm/model_executor/models/nemotron_parse.py index 3ba731d2e0..61fcfb3d1f 100644 --- a/vllm/model_executor/models/nemotron_parse.py +++ b/vllm/model_executor/models/nemotron_parse.py @@ -9,11 +9,13 @@ import math from collections.abc import Iterable, Mapping, Sequence -from typing import Annotated, Literal +from typing import Annotated, Any, Literal, cast +import numpy as np import torch import torch.nn as nn from einops import rearrange +from PIL import Image from transformers import ( BartConfig, BatchFeature, @@ -26,6 +28,7 @@ from vllm.config.multimodal import MultiModalDummyOptions from vllm.inputs import MultiModalDataDict from vllm.logger import init_logger from vllm.model_executor.layers.activation import get_act_fn +from vllm.model_executor.layers.fusion.mm_input_norm import build_mm_input_norm from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.quantization.base_config import QuantizationConfig @@ -401,6 +404,56 @@ class NemotronParseDummyInputsBuilder( } +def _to_pil_image(image: object) -> Image.Image: + """Convert a `uint8` array/tensor image item to PIL, since the checkpoint's + image processor calls `.convert("RGB")` on every input.""" + if isinstance(image, Image.Image): + return image + + array = image.numpy() if isinstance(image, torch.Tensor) else np.asarray(image) + if array.dtype != np.uint8 or array.ndim != 3: + raise ValueError( + "Nemotron-Parse expects image arrays/tensors to be 3-D uint8, " + f"got {array.dtype} {array.shape}" + ) + # Same layout heuristic as `ImageProcessorItems.get_image_size`. + if array.shape[-1] not in (1, 3, 4): + array = array.transpose(1, 2, 0) + if array.shape[-1] == 1: + array = array[..., 0] + + return Image.fromarray(np.ascontiguousarray(array)) + + +def _build_uint8_canvas( + image_processor: Any, images: Sequence[Image.Image] +) -> torch.Tensor: + """Reproduce the image processor's pipeline up to, but excluding, its + `ToTensor` rescale, yielding the padded `uint8` canvas as `(N, 3, H, W)`. + + Raises `AttributeError`/`TypeError` if the checkpoint's remote code does + not expose that pipeline. + """ + resize = image_processor._resize_with_aspect_ratio + pad = image_processor._pad_to_size + # v1.x pads with albumentations, 2.0 sets `transform` to None and pads + # with `_pad_to_size`; mirror `preprocess` and pick the same branch. + transform = image_processor.transform + + canvases = [] + for image in images: + array = np.asarray(image.convert("RGB")) + array = resize(array) + array = transform(image=array)["image"] if transform is not None else pad(array) + if array.dtype != np.uint8 or array.ndim != 3 or array.shape[2] != 3: + raise TypeError( + f"expected a uint8 HWC canvas, got {array.dtype} {array.shape}" + ) + canvases.append(torch.from_numpy(np.ascontiguousarray(array)).permute(2, 0, 1)) + + return torch.stack(canvases) + + class NemotronParseMultiModalProcessor( EncDecMultiModalProcessor[NemotronParseProcessingInfo] ): @@ -411,6 +464,35 @@ class NemotronParseMultiModalProcessor( ) -> list[int]: return [0] + def _call_hf_processor( + self, + hf_data: Mapping[str, object], + hf_kwargs: Mapping[str, object], + ) -> BatchFeature: + images = [_to_pil_image(i) for i in cast(list, hf_data.get("images") or [])] + if images: + hf_data = {**hf_data, "images": images} + + # The checkpoint ships its image processor as remote code, and its + # `preprocess` ignores `do_rescale`/`do_normalize`: it always runs + # `ToTensor` and normalises from `self.do_normalize`. Build the uint8 + # canvas here so the device-side norm is not applied on top. + mm_config = self.info.ctx.model_config.get_multimodal_config() + if images and set(hf_data) == {"images"} and mm_config.mm_device_do_normalize: + image_processor = self.info.get_hf_processor().image_processor + try: + pixel_values = _build_uint8_canvas(image_processor, images) + except (AttributeError, TypeError) as exc: + raise RuntimeError( + f"{type(image_processor).__name__} does not expose the " + f"expected pre-rescale pipeline ({exc}), so the uint8 canvas " + "that mm_device_do_normalize requires cannot be built. Pass " + "--no-mm-device-do-normalize to normalise on the host instead." + ) from exc + return BatchFeature({"pixel_values": pixel_values}) + + return super()._call_hf_processor(hf_data, hf_kwargs) + def _get_mm_fields_config( self, hf_inputs: BatchFeature, @@ -542,6 +624,8 @@ class RadioWithNeck(nn.Module): dummy_inputs=NemotronParseDummyInputsBuilder, ) class NemotronParseForConditionalGeneration(nn.Module, SupportsMultiModal): + supports_mm_device_do_normalize = True + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() config = vllm_config.model_config.hf_config @@ -550,6 +634,7 @@ class NemotronParseForConditionalGeneration(nn.Module, SupportsMultiModal): self.vision_config = config.encoder cache_config = vllm_config.cache_config quant_config = vllm_config.quant_config + self.input_norm = build_mm_input_norm(vllm_config.model_config) with self._mark_tower_model(vllm_config, "image"): self.encoder = RadioWithNeck( @@ -622,8 +707,8 @@ class NemotronParseForConditionalGeneration(nn.Module, SupportsMultiModal): assert image_input["type"] == "pixel_values" pixel_values = image_input["data"] dtype = next(self.encoder.parameters()).dtype - pixel_values = pixel_values.to(dtype) - return self.encoder(pixel_values) + flat = self.input_norm(pixel_values.flatten(1), dtype) + return self.encoder(flat.view(pixel_values.shape)) def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings | None: image_input = self._parse_and_validate_image_input(**kwargs)