A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 16:30 UTC
[Model] Add nvidia/llama-nemotron-embed-vl-1b-v2 multimodal embedding model (#35297)
Signed-off-by: Jakub Zakrzewski <jzakrzewski@nvidia.com>
Jakub Zakrzewski committed
Feb 26, 2026 at 15:17 UTC
111d8690699927af686fa6750cfbbc692a1f8740
8 files changed
+545
-31
docs/models/pooling_models.md
+61
@@ -498,6 +498,67 @@ curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{
498
- Multi-vector retrieval: [examples/pooling/token_embed/colqwen3_token_embed_online.py](../../examples/pooling/token_embed/colqwen3_token_embed_online.py)
499
- Reranking (text + multi-modal): [examples/pooling/score/colqwen3_rerank_online.py](../../examples/pooling/score/colqwen3_rerank_online.py)
500
501
+### Llama Nemotron Multimodal Embedding Models
502
+
503
+Llama Nemotron VL Embedding models combine the bidirectional Llama embedding backbone
504
+(from `nvidia/llama-nemotron-embed-1b-v2`) with SigLIP as the vision encoder to produce
505
+single-vector embeddings from text and/or images.
506
+
507
+| Architecture | Backbone | Example HF Models |
508
+|---|---|---|
509
+| `LlamaNemotronVLModel` | Bidirectional Llama + SigLIP | `nvidia/llama-nemotron-embed-vl-1b-v2` |
510
+
511
+Start the server:
512
+
513
+```shell
514
+vllm serve nvidia/llama-nemotron-embed-vl-1b-v2 \
515
+ --trust-remote-code \
516
+ --chat-template examples/pooling/embed/template/nemotron_embed_vl.jinja
517
+```
518
+
519
+!!! note
520
+ The chat template bundled with this model's tokenizer is not suitable for
521
+ the embeddings API. Use the provided override template above when serving
522
+ with the `messages`-based (chat-style) embeddings endpoint.
523
+
524
+ The override template uses the message `role` to automatically prepend the
525
+ appropriate prefix: set `role` to `"query"` for queries (prepends `query: `)
526
+ or `"document"` for passages (prepends `passage: `). Any other role omits
527
+ the prefix.
528
+
529
+Embed text queries:
530
+
531
+```shell
532
+curl -s http://localhost:8000/v1/embeddings -H "Content-Type: application/json" -d '{
533
+ "model": "nvidia/llama-nemotron-embed-vl-1b-v2",
534
+ "messages": [
535
+ {
536
+ "role": "query",
537
+ "content": [
538
+ {"type": "text", "text": "What is machine learning?"}
539
+ ]
540
+ }
541
+ ]
542
+}'
543
+```
544
+
545
+Embed images via the chat-style `messages` field:
546
+
547
+```shell
548
+curl -s http://localhost:8000/v1/embeddings -H "Content-Type: application/json" -d '{
549
+ "model": "nvidia/llama-nemotron-embed-vl-1b-v2",
550
+ "messages": [
551
+ {
552
+ "role": "document",
553
+ "content": [
554
+ {"type": "image_url", "image_url": {"url": "data:image/png;base64,<BASE64>"}},
555
+ {"type": "text", "text": "Describe the image."}
556
+ ]
557
+ }
558
+ ]
559
+}'
560
+```
561
+
562
### BAAI/bge-m3
563
564
The `BAAI/bge-m3` model comes with extra weights for sparse and colbert embeddings but unfortunately in its `config.json`
docs/models/supported_models.md
+1
@@ -821,6 +821,7 @@ The following table lists those that are tested in vLLM.
821
|--------------|--------|--------|-------------------|----------------------|---------------------------|
822
| `CLIPModel` | CLIP | T / I | `openai/clip-vit-base-patch32`, `openai/clip-vit-large-patch14`, etc. | | |
823
| `ColModernVBertForRetrieval` | ColModernVBERT | T / I | `ModernVBERT/colmodernvbert-merged` | | |
824
+| `LlamaNemotronVLModel` | Llama Nemotron Embedding + SigLIP | T + I | `nvidia/llama-nemotron-embed-vl-1b-v2` | | |
825
| `LlavaNextForConditionalGeneration`<sup>C</sup> | LLaVA-NeXT-based | T / I | `royokong/e5-v` | | ✅︎ |
826
| `Phi3VForCausalLM`<sup>C</sup> | Phi-3-Vision-based | T + I | `TIGER-Lab/VLM2Vec-Full` | | ✅︎ |
827
| `Qwen3VLForConditionalGeneration`<sup>C</sup> | Qwen3-VL | T + I + V | `Qwen/Qwen3-VL-Embedding-2B`, etc. | ✅︎ | ✅︎ |
examples/pooling/embed/template/nemotron_embed_vl.jinja
new
+20
@@ -0,0 +1,20 @@
1
+{%- if messages | length > 1 -%}
2
+ {{ raise_exception('Embedding models should only embed one message at a time') }}
3
+{%- endif -%}
4
+
5
+{% set vars = namespace(prefix='', images=[], texts=[]) %}
6
+{%- for message in messages -%}
7
+ {%- if message['role'] == 'query' -%}
8
+ {%- set vars.prefix = 'query: ' %}
9
+ {%- elif message['role'] == 'document' -%}
10
+ {%- set vars.prefix = 'passage: ' %}
11
+ {%- endif -%}
12
+ {%- for content in message['content'] -%}
13
+ {%- if content['type'] == 'text' -%}
14
+ {%- set vars.texts = vars.texts + [content['text']] %}
15
+ {%- elif content['type'] == 'image' -%}
16
+ {%- set vars.images = vars.images + ['<image> '] %}
17
+ {%- endif -%}
18
+ {%- endfor -%}
19
+{%- endfor -%}
20
+{{- bos_token }}{{ vars.prefix }}{{ (vars.images + vars.texts) | join('') }}
tests/models/multimodal/pooling/test_llama_nemotron_vl_embed.py
new
+148
@@ -0,0 +1,148 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Tests for LlamaNemotronVL embedding model (nvidia/llama-nemotron-embed-vl-1b-v2).
5
+
6
+This model uses SigLIP vision encoder with bidirectional LLaMA for embeddings.
7
+"""
8
+
9
+import pytest
10
+import torch
11
+from transformers import AutoModel
12
+
13
+from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner
14
+from ...utils import check_embeddings_close
15
+
16
+# Prefixes used by the model API
17
+QUERY_PREFIX = "query: "
18
+PASSAGE_PREFIX = "passage: "
19
+
20
+# Text prompts for text-only embedding
21
+HF_TEXT_PROMPTS = [
22
+ # T -> X (text embedding queries)
23
+ f"{QUERY_PREFIX}The label of the object is stop sign",
24
+ f"{QUERY_PREFIX}cherry blossom",
25
+]
26
+
27
+# Image prompts using the model's expected format
28
+HF_IMAGE_PROMPTS = IMAGE_ASSETS.prompts(
29
+ {
30
+ # I -> X (image embedding as passage/document)
31
+ "stop_sign": f"{PASSAGE_PREFIX}<image>",
32
+ "cherry_blossom": f"{PASSAGE_PREFIX}<image>",
33
+ }
34
+)
35
+
36
+MODELS = ["nvidia/llama-nemotron-embed-vl-1b-v2"]
37
+
38
+
39
+def _run_test(
40
+ hf_runner: type[HfRunner],
41
+ vllm_runner: type[VllmRunner],
42
+ input_texts: list[str],
43
+ input_images: PromptImageInput,
44
+ model: str,
45
+ *,
46
+ dtype: str,
47
+) -> None:
48
+ """Run embedding comparison test between HF and vLLM.
49
+
50
+ NOTE: Run vLLM first to avoid CUDA initialization issues with multiprocessing.
51
+ """
52
+ # Run vLLM inference first
53
+ with vllm_runner(
54
+ model,
55
+ runner="pooling",
56
+ dtype=dtype,
57
+ max_model_len=2048,
58
+ enforce_eager=True,
59
+ trust_remote_code=True,
60
+ ) as vllm_model:
61
+ vllm_outputs = vllm_model.embed(input_texts, images=input_images)
62
+
63
+ # Run HF inference using the model's encode_queries/encode_documents API
64
+ with hf_runner(model, dtype=dtype, auto_cls=AutoModel) as hf_model:
65
+ hf_outputs = []
66
+ for text, image in zip(input_texts, input_images):
67
+ with torch.inference_mode():
68
+ if text.startswith(QUERY_PREFIX):
69
+ # Strip prefix and use encode_queries for query texts
70
+ query_text = text[len(QUERY_PREFIX) :]
71
+ embedding = hf_model.model.encode_queries([query_text])
72
+ elif text.startswith(PASSAGE_PREFIX):
73
+ # Strip prefix and use encode_documents for passages/images
74
+ passage_text = text[len(PASSAGE_PREFIX) :]
75
+ if image is not None:
76
+ # Image document - pass image to encode_documents
77
+ embedding = hf_model.model.encode_documents(
78
+ images=[image],
79
+ texts=[passage_text],
80
+ )
81
+ else:
82
+ # Text-only document
83
+ embedding = hf_model.model.encode_documents(
84
+ texts=[passage_text]
85
+ )
86
+ else:
87
+ raise ValueError(
88
+ f"Text must start with '{QUERY_PREFIX}' or '{PASSAGE_PREFIX}'"
89
+ )
90
+
91
+ hf_outputs.append(embedding[0].tolist())
92
+
93
+ check_embeddings_close(
94
+ embeddings_0_lst=hf_outputs,
95
+ embeddings_1_lst=vllm_outputs,
96
+ name_0="hf",
97
+ name_1="vllm",
98
+ )
99
+
100
+
101
+@pytest.mark.parametrize("model", MODELS)
102
+@pytest.mark.parametrize("dtype", ["bfloat16"])
103
+def test_models_text(
104
+ hf_runner,
105
+ vllm_runner,
106
+ image_assets,
107
+ model: str,
108
+ dtype: str,
109
+) -> None:
110
+ """Test text-only embedding."""
111
+ input_texts_images = [(text, None) for text in HF_TEXT_PROMPTS]
112
+ input_texts = [text for text, _ in input_texts_images]
113
+ input_images = [image for _, image in input_texts_images]
114
+
115
+ _run_test(
116
+ hf_runner,
117
+ vllm_runner,
118
+ input_texts,
119
+ input_images, # type: ignore
120
+ model,
121
+ dtype=dtype,
122
+ )
123
+
124
+
125
+@pytest.mark.parametrize("model", MODELS)
126
+@pytest.mark.parametrize("dtype", ["bfloat16"])
127
+def test_models_image(
128
+ hf_runner,
129
+ vllm_runner,
130
+ image_assets,
131
+ model: str,
132
+ dtype: str,
133
+) -> None:
134
+ """Test image embedding."""
135
+ input_texts_images = [
136
+ (text, asset.pil_image) for text, asset in zip(HF_IMAGE_PROMPTS, image_assets)
137
+ ]
138
+ input_texts = [text for text, _ in input_texts_images]
139
+ input_images = [image for _, image in input_texts_images]
140
+
141
+ _run_test(
142
+ hf_runner,
143
+ vllm_runner,
144
+ input_texts,
145
+ input_images,
146
+ model,
147
+ dtype=dtype,
148
+ )
tests/models/registry.py
+3
@@ -598,6 +598,9 @@ _EMBEDDING_EXAMPLE_MODELS = {
598
"ColModernVBertForRetrieval": _HfExamplesInfo(
599
"ModernVBERT/colmodernvbert-merged",
600
),
601
+ "LlamaNemotronVLModel": _HfExamplesInfo(
602
+ "nvidia/llama-nemotron-embed-vl-1b-v2", trust_remote_code=True
603
+ ),
604
"LlavaNextForConditionalGeneration": _HfExamplesInfo("royokong/e5-v"),
605
"Phi3VForCausalLM": _HfExamplesInfo(
606
"TIGER-Lab/VLM2Vec-Full", trust_remote_code=True
vllm/model_executor/models/config.py
+37
@@ -112,6 +112,42 @@ class LlamaBidirectionalConfig(VerifyAndUpdateConfig):
112
model_config.pooler_config.seq_pooling_type = pooling_type
113
114
115
+class LlamaNemotronVLConfig(VerifyAndUpdateConfig):
116
+ """Config handler for LlamaNemotronVL embedding models."""
117
+
118
+ @staticmethod
119
+ def verify_and_update_model_config(model_config: "ModelConfig") -> None:
120
+ from vllm.config.pooler import SequencePoolingType
121
+
122
+ hf_config = model_config.hf_config
123
+
124
+ # Set bidirectional attention on the language model config
125
+ hf_config.is_causal = False
126
+ if hasattr(hf_config, "llm_config"):
127
+ hf_config.llm_config.is_causal = False
128
+
129
+ if hasattr(hf_config, "vision_config"):
130
+ hf_config.patch_size = hf_config.vision_config.patch_size
131
+
132
+ # Set up pooling type
133
+ pooling_type_map: dict[str, SequencePoolingType] = {
134
+ "avg": "MEAN",
135
+ "cls": "CLS",
136
+ "last": "LAST",
137
+ }
138
+
139
+ # Get pooling type from config (check both top-level and llm_config)
140
+ pooling = getattr(hf_config, "pooling", None)
141
+ if pooling is None and hasattr(hf_config, "llm_config"):
142
+ pooling = getattr(hf_config.llm_config, "pooling", "avg")
143
+
144
+ pooling_type = pooling_type_map.get(pooling)
145
+ if pooling_type is None:
146
+ raise ValueError(f"pool_type {pooling!r} not supported")
147
+
148
+ model_config.pooler_config.seq_pooling_type = pooling_type
149
+
150
+
151
class NomicBertModelConfig(VerifyAndUpdateConfig):
152
@staticmethod
153
def verify_and_update_model_config(model_config: "ModelConfig") -> None:
@@ -619,6 +655,7 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
655
"Gemma3TextModel": Gemma3TextModelConfig,
656
"LlamaBidirectionalForSequenceClassification": LlamaBidirectionalConfig,
657
"LlamaBidirectionalModel": LlamaBidirectionalConfig,
658
+ "LlamaNemotronVLModel": LlamaNemotronVLConfig,
659
"NomicBertModel": NomicBertModelConfig,
660
"Qwen2ForProcessRewardModel": Qwen2ForProcessRewardModelConfig,
661
"Qwen2ForRewardModel": Qwen2ForRewardModelConfig,
vllm/model_executor/models/nemotron_vl.py
+271
-31
@@ -18,6 +18,7 @@ from transformers import AutoModel, PretrainedConfig
18
from transformers.image_processing_utils_fast import BaseImageProcessorFast
19
20
from vllm.config import VllmConfig
21
+from vllm.model_executor.layers.pooler import DispatchPooler
22
from vllm.model_executor.layers.quantization import QuantizationConfig
23
from vllm.model_executor.layers.quantization.awq import AWQConfig
24
from vllm.model_executor.models.internvl import (
@@ -30,12 +31,14 @@ from vllm.model_executor.models.internvl import (
31
InternVLProcessor,
32
)
33
from vllm.model_executor.models.module_mapping import MultiModelKeys
34
+from vllm.model_executor.models.siglip import SiglipVisionModel
35
from vllm.multimodal import MULTIMODAL_REGISTRY
36
from vllm.multimodal.image import convert_image_mode
37
from vllm.multimodal.processing import PromptUpdateDetails
38
from vllm.sequence import IntermediateTensors
39
from vllm.tokenizers import TokenizerLike
40
from vllm.transformers_utils.processor import cached_image_processor_from_config
41
+from vllm.transformers_utils.repo_utils import get_hf_file_to_dict
42
43
from .interfaces import (
44
MultiModalEmbeddings,
@@ -43,11 +46,13 @@ from .interfaces import (
46
SupportsMultiModal,
47
SupportsPP,
48
)
46
-from .utils import AutoWeightsLoader, init_vllm_registered_model, maybe_prefix
47
-
48
-IMG_START = "<img>"
49
-IMG_END = "</img>"
50
-IMG_CONTEXT = "<image>"
49
+from .interfaces_base import VllmModelForPooling
50
+from .utils import (
51
+ AutoWeightsLoader,
52
+ WeightsMapper,
53
+ init_vllm_registered_model,
54
+ maybe_prefix,
55
+)
56
57
58
def build_transform(input_size: int):
@@ -183,10 +188,12 @@ def image_to_pixel_values_nemotron_vl(
188
min_num: int,
189
max_num: int,
190
use_thumbnail: bool,
191
+ transform: T.Compose | None = None,
192
) -> torch.Tensor:
193
target_ratios = get_nemotron_vl_target_ratios(min_num, max_num)
194
189
- transform = build_transform(input_size=input_size)
195
+ if transform is None:
196
+ transform = build_transform(input_size=input_size)
197
198
images = dynamic_preprocess_nemotron_vl(
199
image,
@@ -200,11 +207,15 @@ def image_to_pixel_values_nemotron_vl(
207
208
209
class NemotronVLProcessor(InternVLProcessor):
210
+ IMG_START = "<img>"
211
+ IMG_END = "</img>"
212
+ IMG_CONTEXT = "<image>"
213
+
214
def __init__(
215
self,
216
config: PretrainedConfig,
217
tokenizer: TokenizerLike,
207
- image_processor: BaseImageProcessorFast,
218
+ image_processor: BaseImageProcessorFast | None = None,
219
*,
220
min_dynamic_patch: int | None = None,
221
max_dynamic_patch: int | None = None,
@@ -236,11 +247,18 @@ class NemotronVLProcessor(InternVLProcessor):
247
self.min_dynamic_patch = min_dynamic_patch
248
self.max_dynamic_patch = max_dynamic_patch
249
self.dynamic_image_size = dynamic_image_size
239
- self.use_thumbnail: bool = self.image_processor.use_thumbnail
250
+
251
+ if image_processor is not None:
252
+ self.use_thumbnail = image_processor.use_thumbnail
253
+ else:
254
+ self.use_thumbnail = getattr(config, "use_thumbnail", True)
255
256
@property
257
def image_token_id(self) -> int:
243
- return self.tokenizer.get_vocab()[IMG_CONTEXT]
258
+ return self.tokenizer.get_vocab()[self.IMG_CONTEXT]
259
+
260
+ def _get_transform(self) -> T.Compose:
261
+ return build_transform(input_size=self.image_size)
262
263
def get_num_image_tokens(
264
self,
@@ -283,10 +301,26 @@ class NemotronVLProcessor(InternVLProcessor):
301
min_num=min_num,
302
max_num=max_num,
303
use_thumbnail=self.use_thumbnail,
304
+ transform=self._get_transform(),
305
)
306
for image in images
307
]
308
309
+ def _replace_image_tokens(
310
+ self,
311
+ text: list[str],
312
+ pixel_values_lst: list[torch.Tensor],
313
+ ) -> list[str]:
314
+ """Replace <image> placeholders with image tokens."""
315
+ for pixel_values in pixel_values_lst:
316
+ num_patches = pixel_values.shape[0]
317
+ feature_size = num_patches * self.num_image_token
318
+ image_repl = self.get_image_repl(feature_size, num_patches)
319
+ # Use temporary placeholder to avoid replacing tokens we just inserted
320
+ NVL_IMAGE_CONTEXT = image_repl.full.replace("<image>", "<NVL_IMG_CONTEXT>")
321
+ text = [t.replace("<image>", NVL_IMAGE_CONTEXT, 1) for t in text]
322
+ return [t.replace("<NVL_IMG_CONTEXT>", self.IMG_CONTEXT) for t in text]
323
+
324
def _preprocess_image(
325
self,
326
text: list[str],
@@ -311,15 +345,7 @@ class NemotronVLProcessor(InternVLProcessor):
345
),
346
}
347
314
- for pixel_values in pixel_values_lst:
315
- num_patches = pixel_values.shape[0]
316
- feature_size = num_patches * self.num_image_token
317
- image_repl = self.get_image_repl(feature_size, num_patches)
318
- NVL_IMAGE_CONTEXT = image_repl.full.replace(
319
- "<image>", "<NVL_IMG_CONTEXT>"
320
- )
321
- text = [t.replace("<image>", NVL_IMAGE_CONTEXT, 1) for t in text]
322
- text = [t.replace("<NVL_IMG_CONTEXT>", IMG_CONTEXT) for t in text]
348
+ text = self._replace_image_tokens(text, pixel_values_lst)
349
return text, image_inputs
350
351
def get_image_repl(
@@ -327,10 +353,10 @@ class NemotronVLProcessor(InternVLProcessor):
353
feature_size: int,
354
num_patches: int | None,
355
) -> PromptUpdateDetails[str]:
330
- repl_features = IMG_CONTEXT * feature_size
331
- repl_full = IMG_START + repl_features + IMG_END
356
+ repl_features = self.IMG_CONTEXT * feature_size
357
+ repl_full = self.IMG_START + repl_features + self.IMG_END
358
333
- return PromptUpdateDetails.select_text(repl_full, IMG_CONTEXT)
359
+ return PromptUpdateDetails.select_text(repl_full, self.IMG_CONTEXT)
360
361
362
class NemotronVLProcessingInfo(BaseInternVLProcessingInfo):
@@ -396,7 +422,7 @@ class LlamaNemotronVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, Suppor
422
with self._mark_language_model(vllm_config):
423
self.language_model = init_vllm_registered_model(
424
vllm_config=vllm_config,
399
- hf_config=config.text_config,
425
+ hf_config=config.get_text_config(),
426
prefix=maybe_prefix(prefix, "language_model"),
427
)
428
@@ -413,7 +439,7 @@ class LlamaNemotronVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, Suppor
439
# the awq models from OpenGVLab missing `modules_to_not_convert`
440
# patch the quant_config to add `modules_to_not_convert` back
441
if isinstance(quant_config, AWQConfig):
416
- text_config = config.text_config
442
+ text_config = config.get_text_config()
443
llm_quant_config = getattr(text_config, "quantization_config", None)
444
if (not quant_config.modules_to_not_convert) and (
445
llm_quant_config is not None
@@ -429,10 +455,17 @@ class LlamaNemotronVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, Suppor
455
):
456
return AutoModel.from_config(config.vision_config, trust_remote_code=True)
457
432
- def _init_mlp1(self, config: PretrainedConfig) -> nn.Module:
433
- vit_hidden_size = config.vit_hidden_size
434
- vision_projection_hidden_size = config.projector_hidden_size
435
- llm_hidden_size = config.text_config.hidden_size
458
+ def _init_mlp1(
459
+ self,
460
+ config: PretrainedConfig,
461
+ vit_hidden_size: int | None = None,
462
+ vision_projection_hidden_size: int | None = None,
463
+ ) -> nn.Module:
464
+ if vit_hidden_size is None:
465
+ vit_hidden_size = config.vit_hidden_size
466
+ if vision_projection_hidden_size is None:
467
+ vision_projection_hidden_size = config.projector_hidden_size
468
+ llm_hidden_size = config.get_text_config().hidden_size
469
470
return nn.Sequential(
471
nn.LayerNorm(
@@ -465,10 +498,18 @@ class LlamaNemotronVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, Suppor
498
x = x.permute(0, 2, 1, 3).contiguous()
499
return x
500
501
+ def _call_vision_model(self, pixel_values: torch.Tensor) -> torch.Tensor:
502
+ """Call vision model and return embeddings.
503
+
504
+ Override this method in subclasses to handle different vision model
505
+ interfaces (e.g., SigLIP vs C-RADIO).
506
+ """
507
+ vit_embeds = self.vision_model(x=pixel_values).features
508
+ return vit_embeds.to(dtype=torch.bfloat16)
509
+
510
def extract_feature(self, pixel_values: torch.Tensor) -> torch.Tensor:
511
# https://huggingface.co/nvidia/Llama-3.1-Nemotron-Nano-VL-8B-V1/blob/main/modeling.py#L177
470
- vit_embeds = self.vision_model(x=pixel_values).features
471
- vit_embeds = vit_embeds.to(dtype=torch.bfloat16)
512
+ vit_embeds = self._call_vision_model(pixel_values)
513
514
h = w = int(vit_embeds.shape[1] ** 0.5)
515
vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)
@@ -523,15 +564,16 @@ class LlamaNemotronVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, Suppor
564
image_embeds = self.extract_feature(image_input["pixel_values_flat"])
565
566
num_patches = image_input["num_patches"]
567
+ hidden_size = self.config.get_text_config().hidden_size
568
569
# Only one image in the current batch
570
if len(num_patches) == 1:
529
- return (image_embeds.view(-1, self.config.text_config.hidden_size),)
571
+ return (image_embeds.view(-1, hidden_size),)
572
573
# NOTE: Image embeddings are split into separate tensors for each image
574
# by the size of each embedding.
575
feature_size = image_embeds.shape[1]
534
- image_embeds = image_embeds.view(-1, self.config.text_config.hidden_size)
576
+ image_embeds = image_embeds.view(-1, hidden_size)
577
image_feature_sizes = [
578
num_patches * feature_size for num_patches in num_patches
579
]
@@ -643,3 +685,201 @@ class LlamaNemotronVLChatModel(nn.Module, SupportsMultiModal, SupportsPP, Suppor
685
connector="mlp1",
686
tower_model="vision_model",
687
)
688
+
689
+
690
+# --------------------------------------------------------
691
+# LlamaNemotronVL Embedding Model (nvidia/llama-nemotron-embed-vl-1b-v2)
692
+# Extends LlamaNemotronVLChatModel for embedding/pooling tasks:
693
+# - SigLIP vision encoder (instead of C-RADIO)
694
+# - Bidirectional (non-causal) LLaMA language model
695
+# - Pooler output instead of generative logits
696
+# --------------------------------------------------------
697
+
698
+# SigLIP normalization constants
699
+SIGLIP_MEAN = (0.5, 0.5, 0.5)
700
+SIGLIP_STD = (0.5, 0.5, 0.5)
701
+
702
+
703
+def build_siglip_transform(input_size: int):
704
+ """Build transform for SigLIP vision encoder with normalization.
705
+
706
+ Extends the base transform from nemotron_vl with SigLIP-specific normalization.
707
+ """
708
+ base_transform = build_transform(input_size=input_size)
709
+ return T.Compose(
710
+ [
711
+ base_transform,
712
+ T.Normalize(mean=SIGLIP_MEAN, std=SIGLIP_STD),
713
+ ]
714
+ )
715
+
716
+
717
+class LlamaNemotronVLEmbedProcessor(NemotronVLProcessor):
718
+ """
719
+ Processor for LlamaNemotronVL embedding model.
720
+
721
+ Inherits from NemotronVLProcessor and specializes it for embedding tasks:
722
+ - Uses SigLIP transform with normalization instead of base transform
723
+ - Uses different image context token (<IMG_CONTEXT> vs <image>)
724
+ """
725
+
726
+ IMG_CONTEXT = "<IMG_CONTEXT>"
727
+
728
+ def __init__(
729
+ self,
730
+ config: PretrainedConfig,
731
+ tokenizer: TokenizerLike,
732
+ processor_config: dict,
733
+ *,
734
+ min_dynamic_patch: int | None = None,
735
+ max_dynamic_patch: int | None = None,
736
+ dynamic_image_size: bool | None = None,
737
+ ) -> None:
738
+ if min_dynamic_patch is None:
739
+ min_dynamic_patch = processor_config.get(
740
+ "min_input_tiles",
741
+ getattr(config, "min_dynamic_patch", 1),
742
+ )
743
+ if max_dynamic_patch is None:
744
+ max_dynamic_patch = processor_config.get(
745
+ "max_input_tiles",
746
+ getattr(config, "max_dynamic_patch", 1),
747
+ )
748
+ if dynamic_image_size is None:
749
+ dynamic_image_size = processor_config.get(
750
+ "dynamic_image_size",
751
+ getattr(config, "dynamic_image_size", True),
752
+ )
753
+ super().__init__(
754
+ config=config,
755
+ tokenizer=tokenizer,
756
+ image_processor=None,
757
+ min_dynamic_patch=min_dynamic_patch,
758
+ max_dynamic_patch=max_dynamic_patch,
759
+ dynamic_image_size=dynamic_image_size,
760
+ )
761
+
762
+ def _get_transform(self) -> T.Compose:
763
+ """Override to add SigLIP normalization."""
764
+ return build_siglip_transform(input_size=self.image_size)
765
+
766
+ def _replace_image_tokens(
767
+ self,
768
+ text: list[str],
769
+ pixel_values_lst: list[torch.Tensor],
770
+ ) -> list[str]:
771
+ """Override with simpler token replacement for embedding model.
772
+
773
+ No temporary placeholder needed because IMG_CONTEXT is <IMG_CONTEXT>,
774
+ not <image>, so there's no collision risk.
775
+ """
776
+ for pixel_values in pixel_values_lst:
777
+ num_patches = pixel_values.shape[0]
778
+ feature_size = num_patches * self.num_image_token
779
+ image_repl = self.get_image_repl(feature_size, num_patches)
780
+ text = [t.replace("<image>", image_repl.full, 1) for t in text]
781
+ return text
782
+
783
+
784
+class LlamaNemotronVLEmbedProcessingInfo(NemotronVLProcessingInfo):
785
+ """Processing info for LlamaNemotronVL embedding model."""
786
+
787
+ def get_hf_processor(self, **kwargs: object) -> LlamaNemotronVLEmbedProcessor:
788
+ """Override to create embedding-specific processor without image_processor."""
789
+ model_config = self.ctx.model_config
790
+ processor_config = {}
791
+ if model_config.model is not None:
792
+ processor_config = (
793
+ get_hf_file_to_dict(
794
+ "processor_config.json",
795
+ model_config.model,
796
+ model_config.revision,
797
+ )
798
+ or {}
799
+ )
800
+
801
+ return self.ctx.init_processor(
802
+ LlamaNemotronVLEmbedProcessor,
803
+ config=self.get_hf_config(),
804
+ tokenizer=self.get_tokenizer(),
805
+ processor_config=processor_config,
806
+ **kwargs,
807
+ )
808
+
809
+
810
+@MULTIMODAL_REGISTRY.register_processor(
811
+ BaseInternVLMultiModalProcessor[LlamaNemotronVLEmbedProcessingInfo],
812
+ info=LlamaNemotronVLEmbedProcessingInfo,
813
+ dummy_inputs=BaseInternVLDummyInputsBuilder[LlamaNemotronVLEmbedProcessingInfo],
814
+)
815
+class LlamaNemotronVLForEmbedding(LlamaNemotronVLChatModel, VllmModelForPooling):
816
+ """
817
+ LlamaNemotronVL model for embeddings.
818
+
819
+ Inherits from LlamaNemotronVLChatModel and specializes it for embedding tasks:
820
+ - Uses SigLIP vision encoder instead of C-RADIO
821
+ - Uses bidirectional LLaMA (via llm_config) instead of causal LLaMA
822
+ - Adds pooler for embedding output instead of generating logits
823
+ """
824
+
825
+ is_pooling_model = True
826
+
827
+ # Weight mapping from checkpoint format to vLLM format
828
+ # Different from parent class due to different vision model structure
829
+ weight_mapper = WeightsMapper(
830
+ orig_to_new_prefix={
831
+ # Language model mapping
832
+ "language_model.layers.": "language_model.model.layers.",
833
+ "language_model.embed_tokens.": "language_model.model.embed_tokens.",
834
+ "language_model.norm.": "language_model.model.norm.",
835
+ # Vision model mapping (SiglipVisionModel has nested vision_model)
836
+ "vision_model.encoder.": "vision_model.vision_model.encoder.",
837
+ "vision_model.embeddings.": "vision_model.vision_model.embeddings.",
838
+ "vision_model.post_layernorm.": "vision_model.vision_model.post_layernorm.",
839
+ }
840
+ )
841
+
842
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
843
+ super().__init__(vllm_config=vllm_config, prefix=prefix)
844
+
845
+ config = vllm_config.model_config.hf_config
846
+
847
+ # Override: get img_context_token_id from config (parent sets None)
848
+ self.img_context_token_id = getattr(config, "img_context_token_id", None)
849
+
850
+ # Initialize pooler for embedding output
851
+ pooler_config = vllm_config.model_config.pooler_config
852
+ assert pooler_config is not None
853
+ self.pooler = DispatchPooler.for_embedding(pooler_config)
854
+
855
+ def _init_vision_model(
856
+ self,
857
+ config: PretrainedConfig,
858
+ quant_config,
859
+ *,
860
+ prefix: str,
861
+ ) -> nn.Module:
862
+ """Override to use SigLIP instead of C-RADIO."""
863
+ return SiglipVisionModel(
864
+ config.vision_config,
865
+ quant_config=quant_config,
866
+ prefix=prefix,
867
+ use_head=False,
868
+ )
869
+
870
+ def _init_mlp1(self, config: PretrainedConfig) -> nn.Module:
871
+ """Override to use different MLP structure for embedding model."""
872
+ return super()._init_mlp1(
873
+ config,
874
+ vit_hidden_size=config.vision_config.hidden_size,
875
+ vision_projection_hidden_size=config.get_text_config().hidden_size,
876
+ )
877
+
878
+ def _call_vision_model(self, pixel_values: torch.Tensor) -> torch.Tensor:
879
+ """Override to handle SigLIP interface."""
880
+ return self.vision_model(pixel_values)
881
+
882
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
883
+ """Override to use different weight mapping for SigLIP."""
884
+ loader = AutoWeightsLoader(self)
885
+ return loader.load_weights(weights, mapper=self.weight_mapper)
vllm/model_executor/models/registry.py
+4
@@ -260,6 +260,10 @@ _EMBEDDING_MODELS = {
260
"OpsColQwen3Model": ("colqwen3", "ColQwen3Model"),
261
"Qwen3VLNemotronEmbedModel": ("colqwen3", "ColQwen3Model"),
262
"SiglipModel": ("siglip", "SiglipEmbeddingModel"),
263
+ "LlamaNemotronVLModel": (
264
+ "nemotron_vl",
265
+ "LlamaNemotronVLForEmbedding",
266
+ ),
267
# Technically Terratorch models work on images, both in
268
# input and output. I am adding it here because it piggy-backs on embedding
269
# models for the time being.