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

[MM] Enable device normalization for Llama Nemotron VL Embed/Rerank (#57928)

Browse Source 
Signed-off-by: Jakub Zakrzewski <jzakrzewski@nvidia.com>
Co-authored-by: OpenAI Codex <codex@openai.com>
This commit is contained in:
Jakub Zakrzewski
2026-09-29 08:33:52 +00:00
committed by GitHub
co-authored by OpenAI Codex
parent 61349e342d
commit 70dc122f81
4 changed files with 132 additions and 20 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/design/mm_processing.md
+14 -12
View File
@@ -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
tests/models/multimodal/pooling/test_llama_nemotron_vl.py
+54
View File
@@ -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],
vllm/model_executor/models/nemotron_vl.py
+20
View File
@@ -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]:
vllm/transformers_utils/processors/nemotron_vl.py
+44 -8
View File
@@ -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
]