From 70dc122f81346032ebb658b3d00346f31ff1714d Mon Sep 17 00:00:00 2001 From: Jakub Zakrzewski Date: Tue, 29 Sep 2026 10:33:52 +0200 Subject: [PATCH] [MM] Enable device normalization for Llama Nemotron VL Embed/Rerank (#57928) Signed-off-by: Jakub Zakrzewski Co-authored-by: OpenAI Codex --- docs/design/mm_processing.md | 26 ++++----- .../pooling/test_llama_nemotron_vl.py | 54 +++++++++++++++++++ vllm/model_executor/models/nemotron_vl.py | 20 +++++++ .../processors/nemotron_vl.py | 52 +++++++++++++++--- 4 files changed, 132 insertions(+), 20 deletions(-) diff --git a/docs/design/mm_processing.md b/docs/design/mm_processing.md index f53f7d2dd1..3a37a2d8a4 100644 --- a/docs/design/mm_processing.md +++ b/docs/design/mm_processing.md @@ -78,18 +78,20 @@ GPU-side fusion is controlled by `multimodal_config.mm_device_do_normalize`: - The flag is **enabled by default** for all models that support it. - Currently, it’s on by default for these architectures: -| name | Architecture | Example HF Models | -|--------------------|------------------------------------------|-------------------------------------------------------| -| `intern-s1-pro` | `InternS1ProForConditionalGeneration` | `internlm/Intern-S1-Pro` | -| `intern-s2-mobius` | `InternS2MobiusForConditionalGeneration` | `internlm/Intern-S2-Mobius` | -| `mistral3` | `Mistral3ForConditionalGeneration` | `mistralai/Mistral-Small-3.2-24B-Instruct-2506`, 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. | -| `qwen3-vl-moe` | `Qwen3VLMoeForConditionalGeneration` | `Qwen/Qwen3-VL-30B-A3B-Instruct`, etc. | -| `qwen3.5` | `Qwen3_5ForConditionalGeneration` | `Qwen/Qwen3.5-0.8B`, etc. | -| `qwen3.5-moe` | `Qwen3_5MoeForConditionalGeneration` | `Qwen/Qwen3.5-35B-A3B`, etc. | -| `qwen4-exp` | `Qwen4ExpForConditionalGeneration` | `Qwen/Qwen3.8-Flash-Next` | +| name | Architecture | Example HF Models | +| -------------------------- | ------------------------------------------ | ----------------------------------------------------- | +| `intern-s1-pro` | `InternS1ProForConditionalGeneration` | `internlm/Intern-S1-Pro` | +| `intern-s2-mobius` | `InternS2MobiusForConditionalGeneration` | `internlm/Intern-S2-Mobius` | +| `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. | +| `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. | +| `qwen3-vl-moe` | `Qwen3VLMoeForConditionalGeneration` | `Qwen/Qwen3-VL-30B-A3B-Instruct`, etc. | +| `qwen3.5` | `Qwen3_5ForConditionalGeneration` | `Qwen/Qwen3.5-0.8B`, etc. | +| `qwen3.5-moe` | `Qwen3_5MoeForConditionalGeneration` | `Qwen/Qwen3.5-35B-A3B`, etc. | +| `qwen4-exp` | `Qwen4ExpForConditionalGeneration` | `Qwen/Qwen3.8-Flash-Next` | #### Key Properties and Gains diff --git a/tests/models/multimodal/pooling/test_llama_nemotron_vl.py b/tests/models/multimodal/pooling/test_llama_nemotron_vl.py index 52802939ce..2a0ba96722 100644 --- a/tests/models/multimodal/pooling/test_llama_nemotron_vl.py +++ b/tests/models/multimodal/pooling/test_llama_nemotron_vl.py @@ -23,7 +23,13 @@ from vllm.entrypoints.chat_utils import ( ChatCompletionContentPartTextParam, ) from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam +from vllm.model_executor.layers.fusion.mm_input_norm import FusedMMInputNorm from vllm.platforms import current_platform +from vllm.transformers_utils.processors.nemotron_vl import ( + SIGLIP_MEAN, + SIGLIP_STD, + LlamaNemotronVLEmbedImageProcessor, +) from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner from ....utils import ROCM_ENGINE_KWARGS @@ -52,6 +58,54 @@ HF_IMAGE_PROMPTS = IMAGE_ASSETS.prompts( MODELS = ["nvidia/llama-nemotron-embed-vl-1b-v2"] +@pytest.mark.usefixtures("default_vllm_config") +def test_device_normalization_matches_cpu_transform() -> None: + device = current_platform.device_type + image = Image.new("RGB", (48, 32), color=(12, 128, 240)) + processor_kwargs = dict( + image_size=32, + min_dynamic_patch=1, + max_dynamic_patch=1, + dynamic_image_size=True, + use_thumbnail=False, + ) + + cpu_processor = LlamaNemotronVLEmbedImageProcessor(**processor_kwargs) + device_processor = LlamaNemotronVLEmbedImageProcessor( + **processor_kwargs, do_rescale=False, do_normalize=False + ) + + expected = cpu_processor._images_to_pixel_values_lst([image])[0] + raw_pixels = device_processor._images_to_pixel_values_lst([image])[0] + input_norm = FusedMMInputNorm( + image_mean=list(SIGLIP_MEAN), + image_std=list(SIGLIP_STD), + rescale_factor=1 / 255, + ).to(device) + actual = input_norm( + raw_pixels.flatten(start_dim=1).to(device), expected.dtype + ).view_as(expected) + + assert raw_pixels.dtype == torch.uint8 + torch.testing.assert_close(actual, expected.to(device)) + + +@pytest.mark.parametrize(("do_rescale", "do_normalize"), [(True, False), (False, True)]) +def test_device_normalization_rejects_partial_cpu_processing( + do_rescale: bool, do_normalize: bool +) -> None: + with pytest.raises(ValueError, match="both enabled or both disabled"): + LlamaNemotronVLEmbedImageProcessor( + image_size=32, + min_dynamic_patch=1, + max_dynamic_patch=1, + dynamic_image_size=True, + use_thumbnail=False, + do_rescale=do_rescale, + do_normalize=do_normalize, + ) + + def _run_test( hf_runner: type[HfRunner], vllm_runner: type[VllmRunner], diff --git a/vllm/model_executor/models/nemotron_vl.py b/vllm/model_executor/models/nemotron_vl.py index 0ec4589e51..b5948048da 100644 --- a/vllm/model_executor/models/nemotron_vl.py +++ b/vllm/model_executor/models/nemotron_vl.py @@ -8,6 +8,10 @@ import torch.nn as nn from transformers import AutoModel, PreTrainedConfig from vllm.config import VllmConfig +from vllm.model_executor.layers.fusion.mm_input_norm import ( + FusedMMInputNorm, + IdentityInputNorm, +) from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.layers.pooler import DispatchPooler from vllm.model_executor.layers.quantization import QuantizationConfig @@ -26,6 +30,8 @@ from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.sequence import IntermediateTensors from vllm.transformers_utils.processor import cached_image_processor_from_config from vllm.transformers_utils.processors.nemotron_vl import ( + SIGLIP_MEAN, + SIGLIP_STD, LlamaNemotronNanoVLImageProcessor, LlamaNemotronNanoVLProcessor, LlamaNemotronVLEmbedImageProcessor, @@ -474,6 +480,7 @@ class LlamaNemotronVLForEmbedding(LlamaNemotronVLChatModel, VllmModelForPooling) """ is_pooling_model = True + supports_mm_device_do_normalize = True # Weight mapping from checkpoint format to vLLM format # Different from parent class due to different vision model structure @@ -498,6 +505,15 @@ class LlamaNemotronVLForEmbedding(LlamaNemotronVLChatModel, VllmModelForPooling) # Override: get img_context_token_id from config (parent sets None) self.img_context_token_id = getattr(config, "img_context_token_id", None) + if self.model_config.get_multimodal_config().mm_device_do_normalize: + self.input_norm = FusedMMInputNorm( + image_mean=list(SIGLIP_MEAN), + image_std=list(SIGLIP_STD), + rescale_factor=1 / 255, + ) + else: + self.input_norm = IdentityInputNorm() + # Initialize pooler for embedding output pooler_config = vllm_config.model_config.pooler_config assert pooler_config is not None @@ -537,6 +553,10 @@ class LlamaNemotronVLForEmbedding(LlamaNemotronVLChatModel, VllmModelForPooling) def _call_vision_model(self, pixel_values: torch.Tensor) -> torch.Tensor: """Override to handle SigLIP interface.""" + original_shape = pixel_values.shape + pixel_values = self.input_norm( + pixel_values.flatten(start_dim=1), self.vision_model.dtype + ).view(original_shape) return self.vision_model(pixel_values) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: diff --git a/vllm/transformers_utils/processors/nemotron_vl.py b/vllm/transformers_utils/processors/nemotron_vl.py index 9504f5ee2b..fefb05cff4 100644 --- a/vllm/transformers_utils/processors/nemotron_vl.py +++ b/vllm/transformers_utils/processors/nemotron_vl.py @@ -244,20 +244,53 @@ SIGLIP_MEAN = (0.5, 0.5, 0.5) SIGLIP_STD = (0.5, 0.5, 0.5) -def build_siglip_transform(input_size: int): +def build_siglip_transform( + input_size: int, + *, + normalize_on_cpu: bool = True, +): """Build transform for SigLIP vision encoder with normalization. Extends the base transform from nemotron_vl with SigLIP-specific normalization. """ - return T.Compose( - [ - build_transform(input_size=input_size), - T.Normalize(mean=SIGLIP_MEAN, std=SIGLIP_STD), - ] - ) + transforms = [ + T.Lambda(lambda img: convert_image_mode(img, "RGB")), + T.Resize((input_size, input_size), interpolation=T.InterpolationMode.BICUBIC), + T.ToTensor() if normalize_on_cpu else T.PILToTensor(), + ] + if normalize_on_cpu: + transforms.append(T.Normalize(mean=SIGLIP_MEAN, std=SIGLIP_STD)) + + return T.Compose(transforms) class LlamaNemotronVLEmbedImageProcessor(InternVLImageProcessor): + def __init__( + self, + image_size: int, + min_dynamic_patch: int, + max_dynamic_patch: int, + dynamic_image_size: bool, + use_thumbnail: bool, + *, + do_rescale: bool = True, + do_normalize: bool = True, + ) -> None: + if do_rescale != do_normalize: + raise ValueError( + "LlamaNemotronVLEmbedImageProcessor requires do_rescale and " + "do_normalize to be both enabled or both disabled." + ) + + super().__init__( + image_size=image_size, + min_dynamic_patch=min_dynamic_patch, + max_dynamic_patch=max_dynamic_patch, + dynamic_image_size=dynamic_image_size, + use_thumbnail=use_thumbnail, + ) + self.normalize_on_cpu = do_normalize + def _images_to_pixel_values_lst( self, images: list[Image.Image], @@ -279,7 +312,10 @@ class LlamaNemotronVLEmbedImageProcessor(InternVLImageProcessor): min_num=min_num, max_num=max_num, use_thumbnail=self.use_thumbnail, - transform=build_siglip_transform(self.image_size), + transform=build_siglip_transform( + self.image_size, + normalize_on_cpu=self.normalize_on_cpu, + ), ) for image in images ]