iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Perf][Nemotron-Parse] Do image rescale/normalize on the device (#60500)

Browse Source 
Signed-off-by: DiegoCao <DiegoCao@users.noreply.github.com>
Co-authored-by: DiegoCao <DiegoCao@users.noreply.github.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Hangrui Cao
2026-10-14 18:33:22 +00:00
committed by GitHub
co-authored by DiegoCao Claude Opus 5.5
parent 0b6d07cbe2
commit d1d8e6c096
4 changed files with 184 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/design/mm_processing.md
+1
View File
@@ -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. |
tests/models/multimodal/processing/test_nemotron_parse.py
+92
View File
@@ -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()
tests/models/registry.py
+3 -1
View File
@@ -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",
vllm/model_executor/models/nemotron_parse.py
+88 -3
View File
@@ -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)