[MM] Enable device normalization for Llama Nemotron VL Embed/Rerank (#57928)
Signed-off-by: Jakub Zakrzewski <jzakrzewski@nvidia.com> Co-authored-by: OpenAI Codex <codex@openai.com>
Jakub Zakrzewski committed
Sep 29, 2026 at 10:33 UTC
70dc122f81346032ebb658b3d00346f31ff1714d
4 files changed
+132
-20
docs/design/mm_processing.md
+14
-12
@@ -78,18 +78,20 @@ GPU-side fusion is controlled by `multimodal_config.mm_device_do_normalize`:
78
- The flag is **enabled by default** for all models that support it.
79
- Currently, it’s on by default for these architectures:
80
81
-| name | Architecture | Example HF Models |
82
-|--------------------|------------------------------------------|-------------------------------------------------------|
83
-| `intern-s1-pro` | `InternS1ProForConditionalGeneration` | `internlm/Intern-S1-Pro` |
84
-| `intern-s2-mobius` | `InternS2MobiusForConditionalGeneration` | `internlm/Intern-S2-Mobius` |
85
-| `mistral3` | `Mistral3ForConditionalGeneration` | `mistralai/Mistral-Small-3.2-24B-Instruct-2506`, etc. |
86
-| `qwen2-vl` | `Qwen2VLForConditionalGeneration` | `Qwen/Qwen2-VL-2B-Instruct`, etc. |
87
-| `qwen2.5-vl` | `Qwen2_5_VLForConditionalGeneration` | `Qwen/Qwen2.5-VL-3B-Instruct`, etc. |
88
-| `qwen3-vl` | `Qwen3VLForConditionalGeneration` | `Qwen/Qwen3-VL-4B-Instruct`, etc. |
89
-| `qwen3-vl-moe` | `Qwen3VLMoeForConditionalGeneration` | `Qwen/Qwen3-VL-30B-A3B-Instruct`, etc. |
90
-| `qwen3.5` | `Qwen3_5ForConditionalGeneration` | `Qwen/Qwen3.5-0.8B`, etc. |
91
-| `qwen3.5-moe` | `Qwen3_5MoeForConditionalGeneration` | `Qwen/Qwen3.5-35B-A3B`, etc. |
92
-| `qwen4-exp` | `Qwen4ExpForConditionalGeneration` | `Qwen/Qwen3.8-Flash-Next` |
81
+| name | Architecture | Example HF Models |
82
+| -------------------------- | ------------------------------------------ | ----------------------------------------------------- |
83
+| `intern-s1-pro` | `InternS1ProForConditionalGeneration` | `internlm/Intern-S1-Pro` |
84
+| `intern-s2-mobius` | `InternS2MobiusForConditionalGeneration` | `internlm/Intern-S2-Mobius` |
85
+| `llama-nemotron-vl-embed` | `LlamaNemotronVLForEmbedding` | `nvidia/llama-nemotron-embed-vl-1b-v2` |
86
+| `llama-nemotron-vl-rerank` | `LlamaNemotronVLForSequenceClassification` | `nvidia/llama-nemotron-rerank-vl-1b-v2` |
87
+| `mistral3` | `Mistral3ForConditionalGeneration` | `mistralai/Mistral-Small-3.2-24B-Instruct-2506`, etc. |
88
+| `qwen2-vl` | `Qwen2VLForConditionalGeneration` | `Qwen/Qwen2-VL-2B-Instruct`, etc. |
89
+| `qwen2.5-vl` | `Qwen2_5_VLForConditionalGeneration` | `Qwen/Qwen2.5-VL-3B-Instruct`, etc. |
90
+| `qwen3-vl` | `Qwen3VLForConditionalGeneration` | `Qwen/Qwen3-VL-4B-Instruct`, etc. |
91
+| `qwen3-vl-moe` | `Qwen3VLMoeForConditionalGeneration` | `Qwen/Qwen3-VL-30B-A3B-Instruct`, etc. |
92
+| `qwen3.5` | `Qwen3_5ForConditionalGeneration` | `Qwen/Qwen3.5-0.8B`, etc. |
93
+| `qwen3.5-moe` | `Qwen3_5MoeForConditionalGeneration` | `Qwen/Qwen3.5-35B-A3B`, etc. |
94
+| `qwen4-exp` | `Qwen4ExpForConditionalGeneration` | `Qwen/Qwen3.8-Flash-Next` |
95
96
#### Key Properties and Gains
97
tests/models/multimodal/pooling/test_llama_nemotron_vl.py
+54
@@ -23,7 +23,13 @@ from vllm.entrypoints.chat_utils import (
23
ChatCompletionContentPartTextParam,
24
)
25
from vllm.entrypoints.pooling.scoring.typing import ScoreMultiModalParam
26
+from vllm.model_executor.layers.fusion.mm_input_norm import FusedMMInputNorm
27
from vllm.platforms import current_platform
28
+from vllm.transformers_utils.processors.nemotron_vl import (
29
+ SIGLIP_MEAN,
30
+ SIGLIP_STD,
31
+ LlamaNemotronVLEmbedImageProcessor,
32
+)
33
34
from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner
35
from ....utils import ROCM_ENGINE_KWARGS
@@ -52,6 +58,54 @@ HF_IMAGE_PROMPTS = IMAGE_ASSETS.prompts(
58
MODELS = ["nvidia/llama-nemotron-embed-vl-1b-v2"]
59
60
61
+@pytest.mark.usefixtures("default_vllm_config")
62
+def test_device_normalization_matches_cpu_transform() -> None:
63
+ device = current_platform.device_type
64
+ image = Image.new("RGB", (48, 32), color=(12, 128, 240))
65
+ processor_kwargs = dict(
66
+ image_size=32,
67
+ min_dynamic_patch=1,
68
+ max_dynamic_patch=1,
69
+ dynamic_image_size=True,
70
+ use_thumbnail=False,
71
+ )
72
+
73
+ cpu_processor = LlamaNemotronVLEmbedImageProcessor(**processor_kwargs)
74
+ device_processor = LlamaNemotronVLEmbedImageProcessor(
75
+ **processor_kwargs, do_rescale=False, do_normalize=False
76
+ )
77
+
78
+ expected = cpu_processor._images_to_pixel_values_lst([image])[0]
79
+ raw_pixels = device_processor._images_to_pixel_values_lst([image])[0]
80
+ input_norm = FusedMMInputNorm(
81
+ image_mean=list(SIGLIP_MEAN),
82
+ image_std=list(SIGLIP_STD),
83
+ rescale_factor=1 / 255,
84
+ ).to(device)
85
+ actual = input_norm(
86
+ raw_pixels.flatten(start_dim=1).to(device), expected.dtype
87
+ ).view_as(expected)
88
+
89
+ assert raw_pixels.dtype == torch.uint8
90
+ torch.testing.assert_close(actual, expected.to(device))
91
+
92
+
93
+@pytest.mark.parametrize(("do_rescale", "do_normalize"), [(True, False), (False, True)])
94
+def test_device_normalization_rejects_partial_cpu_processing(
95
+ do_rescale: bool, do_normalize: bool
96
+) -> None:
97
+ with pytest.raises(ValueError, match="both enabled or both disabled"):
98
+ LlamaNemotronVLEmbedImageProcessor(
99
+ image_size=32,
100
+ min_dynamic_patch=1,
101
+ max_dynamic_patch=1,
102
+ dynamic_image_size=True,
103
+ use_thumbnail=False,
104
+ do_rescale=do_rescale,
105
+ do_normalize=do_normalize,
106
+ )
107
+
108
+
109
def _run_test(
110
hf_runner: type[HfRunner],
111
vllm_runner: type[VllmRunner],
vllm/model_executor/models/nemotron_vl.py
+20
@@ -8,6 +8,10 @@ import torch.nn as nn
8
from transformers import AutoModel, PreTrainedConfig
9
10
from vllm.config import VllmConfig
11
+from vllm.model_executor.layers.fusion.mm_input_norm import (
12
+ FusedMMInputNorm,
13
+ IdentityInputNorm,
14
+)
15
from vllm.model_executor.layers.linear import ReplicatedLinear
16
from vllm.model_executor.layers.pooler import DispatchPooler
17
from vllm.model_executor.layers.quantization import QuantizationConfig
@@ -26,6 +30,8 @@ from vllm.multimodal import MULTIMODAL_REGISTRY
30
from vllm.sequence import IntermediateTensors
31
from vllm.transformers_utils.processor import cached_image_processor_from_config
32
from vllm.transformers_utils.processors.nemotron_vl import (
33
+ SIGLIP_MEAN,
34
+ SIGLIP_STD,
35
LlamaNemotronNanoVLImageProcessor,
36
LlamaNemotronNanoVLProcessor,
37
LlamaNemotronVLEmbedImageProcessor,
@@ -474,6 +480,7 @@ class LlamaNemotronVLForEmbedding(LlamaNemotronVLChatModel, VllmModelForPooling)
480
"""
481
482
is_pooling_model = True
483
+ supports_mm_device_do_normalize = True
484
485
# Weight mapping from checkpoint format to vLLM format
486
# Different from parent class due to different vision model structure
@@ -498,6 +505,15 @@ class LlamaNemotronVLForEmbedding(LlamaNemotronVLChatModel, VllmModelForPooling)
505
# Override: get img_context_token_id from config (parent sets None)
506
self.img_context_token_id = getattr(config, "img_context_token_id", None)
507
508
+ if self.model_config.get_multimodal_config().mm_device_do_normalize:
509
+ self.input_norm = FusedMMInputNorm(
510
+ image_mean=list(SIGLIP_MEAN),
511
+ image_std=list(SIGLIP_STD),
512
+ rescale_factor=1 / 255,
513
+ )
514
+ else:
515
+ self.input_norm = IdentityInputNorm()
516
+
517
# Initialize pooler for embedding output
518
pooler_config = vllm_config.model_config.pooler_config
519
assert pooler_config is not None
@@ -537,6 +553,10 @@ class LlamaNemotronVLForEmbedding(LlamaNemotronVLChatModel, VllmModelForPooling)
553
554
def _call_vision_model(self, pixel_values: torch.Tensor) -> torch.Tensor:
555
"""Override to handle SigLIP interface."""
556
+ original_shape = pixel_values.shape
557
+ pixel_values = self.input_norm(
558
+ pixel_values.flatten(start_dim=1), self.vision_model.dtype
559
+ ).view(original_shape)
560
return self.vision_model(pixel_values)
561
562
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
vllm/transformers_utils/processors/nemotron_vl.py
+44
-8
@@ -244,20 +244,53 @@ SIGLIP_MEAN = (0.5, 0.5, 0.5)
244
SIGLIP_STD = (0.5, 0.5, 0.5)
245
246
247
-def build_siglip_transform(input_size: int):
247
+def build_siglip_transform(
248
+ input_size: int,
249
+ *,
250
+ normalize_on_cpu: bool = True,
251
+):
252
"""Build transform for SigLIP vision encoder with normalization.
253
254
Extends the base transform from nemotron_vl with SigLIP-specific normalization.
255
"""
252
- return T.Compose(
253
- [
254
- build_transform(input_size=input_size),
255
- T.Normalize(mean=SIGLIP_MEAN, std=SIGLIP_STD),
256
- ]
257
- )
256
+ transforms = [
257
+ T.Lambda(lambda img: convert_image_mode(img, "RGB")),
258
+ T.Resize((input_size, input_size), interpolation=T.InterpolationMode.BICUBIC),
259
+ T.ToTensor() if normalize_on_cpu else T.PILToTensor(),
260
+ ]
261
+ if normalize_on_cpu:
262
+ transforms.append(T.Normalize(mean=SIGLIP_MEAN, std=SIGLIP_STD))
263
+
264
+ return T.Compose(transforms)
265
266
267
class LlamaNemotronVLEmbedImageProcessor(InternVLImageProcessor):
268
+ def __init__(
269
+ self,
270
+ image_size: int,
271
+ min_dynamic_patch: int,
272
+ max_dynamic_patch: int,
273
+ dynamic_image_size: bool,
274
+ use_thumbnail: bool,
275
+ *,
276
+ do_rescale: bool = True,
277
+ do_normalize: bool = True,
278
+ ) -> None:
279
+ if do_rescale != do_normalize:
280
+ raise ValueError(
281
+ "LlamaNemotronVLEmbedImageProcessor requires do_rescale and "
282
+ "do_normalize to be both enabled or both disabled."
283
+ )
284
+
285
+ super().__init__(
286
+ image_size=image_size,
287
+ min_dynamic_patch=min_dynamic_patch,
288
+ max_dynamic_patch=max_dynamic_patch,
289
+ dynamic_image_size=dynamic_image_size,
290
+ use_thumbnail=use_thumbnail,
291
+ )
292
+ self.normalize_on_cpu = do_normalize
293
+
294
def _images_to_pixel_values_lst(
295
self,
296
images: list[Image.Image],
@@ -279,7 +312,10 @@ class LlamaNemotronVLEmbedImageProcessor(InternVLImageProcessor):
312
min_num=min_num,
313
max_num=max_num,
314
use_thumbnail=self.use_thumbnail,
282
- transform=build_siglip_transform(self.image_size),
315
+ transform=build_siglip_transform(
316
+ self.image_size,
317
+ normalize_on_cpu=self.normalize_on_cpu,
318
+ ),
319
)
320
for image in images
321
]