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, 10:15 UTC
[Model] Add support for nvidia/llama-nemotron-rerank-vl-1b-v2 (#35735)
Signed-off-by: Jakub Zakrzewski <jzakrzewski@nvidia.com>
Jakub Zakrzewski committed
Mar 3, 2026 at 01:32 UTC
c8b678e53e37b24aa457502e2e47a650b27fd0ec
9 files changed
+503
-149
docs/models/pooling_models.md
+67
-1
@@ -498,7 +498,9 @@ 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
501
+### Llama Nemotron Multimodal
502
+
503
+#### Embedding Model
504
505
Llama Nemotron VL Embedding models combine the bidirectional Llama embedding backbone
506
(from `nvidia/llama-nemotron-embed-1b-v2`) with SigLIP as the vision encoder to produce
@@ -559,6 +561,70 @@ curl -s http://localhost:8000/v1/embeddings -H "Content-Type: application/json"
561
}'
562
```
563
564
+#### Reranker Model
565
+
566
+Llama Nemotron VL reranker models combine the same bidirectional Llama + SigLIP
567
+backbone with a sequence-classification head for cross-encoder scoring and reranking.
568
+
569
+| Architecture | Backbone | Example HF Models |
570
+|---|---|---|
571
+| `LlamaNemotronVLForSequenceClassification` | Bidirectional Llama + SigLIP | `nvidia/llama-nemotron-rerank-vl-1b-v2` |
572
+
573
+Start the server:
574
+
575
+```shell
576
+vllm serve nvidia/llama-nemotron-rerank-vl-1b-v2 \
577
+ --runner pooling \
578
+ --trust-remote-code \
579
+ --chat-template examples/pooling/score/template/nemotron-vl-rerank.jinja
580
+```
581
+
582
+!!! note
583
+ The chat template bundled with this checkpoint's tokenizer is not suitable
584
+ for the Score/Rerank APIs. Use the provided override template when serving:
585
+ `examples/pooling/score/template/nemotron-vl-rerank.jinja`.
586
+
587
+Score a text query against an image document:
588
+
589
+```shell
590
+curl -s http://localhost:8000/score -H "Content-Type: application/json" -d '{
591
+ "model": "nvidia/llama-nemotron-rerank-vl-1b-v2",
592
+ "data_1": "Find diagrams about autonomous robots",
593
+ "data_2": [
594
+ {
595
+ "content": [
596
+ {"type": "image_url", "image_url": {"url": "data:image/png;base64,<BASE64>"}},
597
+ {"type": "text", "text": "Robotics workflow diagram."}
598
+ ]
599
+ }
600
+ ]
601
+}'
602
+```
603
+
604
+Rerank image documents by a text query:
605
+
606
+```shell
607
+curl -s http://localhost:8000/rerank -H "Content-Type: application/json" -d '{
608
+ "model": "nvidia/llama-nemotron-rerank-vl-1b-v2",
609
+ "query": "Find diagrams about autonomous robots",
610
+ "documents": [
611
+ {
612
+ "content": [
613
+ {"type": "image_url", "image_url": {"url": "data:image/png;base64,<BASE64_1>"}},
614
+ {"type": "text", "text": "Robotics workflow diagram."}
615
+ ]
616
+ },
617
+ {
618
+ "content": [
619
+ {"type": "image_url", "image_url": {"url": "data:image/png;base64,<BASE64_2>"}},
620
+ {"type": "text", "text": "General skyline photo."}
621
+ ]
622
+ }
623
+ ],
624
+ "top_n": 2
625
+}'
626
+```
627
+
628
### BAAI/bge-m3
629
630
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
@@ -842,6 +842,7 @@ These models primarily support the [`LLM.score`](./pooling_models.md#llmscore) A
842
| Architecture | Models | Inputs | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) |
843
|--------------|--------|--------|-------------------|----------------------|---------------------------|
844
| `JinaVLForSequenceClassification` | JinaVL-based | T + I<sup>E+</sup> | `jinaai/jina-reranker-m0`, etc. | ✅︎ | ✅︎ |
845
+| `LlamaNemotronVLForSequenceClassification` | Llama Nemotron Reranker + SigLIP | T + I<sup>E+</sup> | `nvidia/llama-nemotron-rerank-vl-1b-v2` | | |
846
| `Qwen3VLForSequenceClassification` | Qwen3-VL-Reranker | T + I<sup>E+</sup> + V<sup>E+</sup> | `Qwen/Qwen3-VL-Reranker-2B`(see note), etc. | ✅︎ | ✅︎ |
847
848
<sup>C</sup> Automatically converted into a classification model via `--convert classify`. ([details](./pooling_models.md#model-conversion))
examples/pooling/score/template/nemotron-vl-rerank.jinja
new
+15
@@ -0,0 +1,15 @@
1
+{%- set query_msg = (messages | selectattr('role', 'equalto', 'query') | list | first) -%}
2
+{%- set doc_msg = (messages | selectattr('role', 'equalto', 'document') | list | first) -%}
3
+
4
+{%- set q = query_msg['content'] -%}
5
+{%- set d = doc_msg['content'] -%}
6
+
7
+{# If the doc contains <image> anywhere, hoist a single <image> to the front #}
8
+{%- set has_image = ("<image>" in d) -%}
9
+{%- set d_clean = d | replace("<image>", "") -%}
10
+{%- set q_clean = q | replace("<image>", "") -%}
11
+
12
+{%- if has_image -%}<image>{{ " " }}{%- endif -%}
13
+question:{{ q_clean }}{{ " " }}
14
+{{ " " }}
15
+{{ " " }}passage:{{ d_clean }}
\ No newline at end of file
tests/models/multimodal/pooling/test_llama_nemotron_vl.py
new
+355
@@ -0,0 +1,355 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""
4
+Tests for the LlamaNemotronVL model family:
5
+ - nvidia/llama-nemotron-embed-vl-1b-v2 (LlamaNemotronVLForCausalLM / embed)
6
+ - nvidia/llama-nemotron-rerank-vl-1b-v2
7
+ (LlamaNemotronVLForSequenceClassification / rerank)
8
+
9
+Both variants share a SigLIP vision encoder with a bidirectional LLaMA backbone.
10
+"""
11
+
12
+import base64
13
+from io import BytesIO
14
+from pathlib import Path
15
+
16
+import pytest
17
+import torch
18
+from transformers import AutoModel, AutoModelForSequenceClassification, AutoProcessor
19
+
20
+from vllm.entrypoints.chat_utils import (
21
+ ChatCompletionContentPartImageParam,
22
+ ChatCompletionContentPartTextParam,
23
+)
24
+from vllm.entrypoints.pooling.score.utils import ScoreMultiModalParam
25
+
26
+from ....conftest import IMAGE_ASSETS, HfRunner, PromptImageInput, VllmRunner
27
+from ...utils import check_embeddings_close
28
+
29
+# Prefixes used by the model API
30
+QUERY_PREFIX = "query: "
31
+PASSAGE_PREFIX = "passage: "
32
+
33
+# Text prompts for text-only embedding
34
+HF_TEXT_PROMPTS = [
35
+ # T -> X (text embedding queries)
36
+ f"{QUERY_PREFIX}The label of the object is stop sign",
37
+ f"{QUERY_PREFIX}cherry blossom",
38
+]
39
+
40
+# Image prompts using the model's expected format
41
+HF_IMAGE_PROMPTS = IMAGE_ASSETS.prompts(
42
+ {
43
+ # I -> X (image embedding as passage/document)
44
+ "stop_sign": f"{PASSAGE_PREFIX}<image>",
45
+ "cherry_blossom": f"{PASSAGE_PREFIX}<image>",
46
+ }
47
+)
48
+
49
+MODELS = ["nvidia/llama-nemotron-embed-vl-1b-v2"]
50
+
51
+
52
+def _run_test(
53
+ hf_runner: type[HfRunner],
54
+ vllm_runner: type[VllmRunner],
55
+ input_texts: list[str],
56
+ input_images: PromptImageInput,
57
+ model: str,
58
+ *,
59
+ dtype: str,
60
+) -> None:
61
+ """Run embedding comparison test between HF and vLLM.
62
+
63
+ NOTE: Run vLLM first to avoid CUDA initialization issues with multiprocessing.
64
+ """
65
+ # Run vLLM inference first
66
+ with vllm_runner(
67
+ model,
68
+ runner="pooling",
69
+ dtype=dtype,
70
+ max_model_len=2048,
71
+ enforce_eager=True,
72
+ trust_remote_code=True,
73
+ ) as vllm_model:
74
+ vllm_outputs = vllm_model.embed(input_texts, images=input_images)
75
+
76
+ # Run HF inference using the model's encode_queries/encode_documents API
77
+ with hf_runner(model, dtype=dtype, auto_cls=AutoModel) as hf_model:
78
+ hf_outputs = []
79
+ for text, image in zip(input_texts, input_images):
80
+ with torch.inference_mode():
81
+ if text.startswith(QUERY_PREFIX):
82
+ # Strip prefix and use encode_queries for query texts
83
+ query_text = text[len(QUERY_PREFIX) :]
84
+ embedding = hf_model.model.encode_queries([query_text])
85
+ elif text.startswith(PASSAGE_PREFIX):
86
+ # Strip prefix and use encode_documents for passages/images
87
+ passage_text = text[len(PASSAGE_PREFIX) :]
88
+ if image is not None:
89
+ # Image document - pass image to encode_documents
90
+ embedding = hf_model.model.encode_documents(
91
+ images=[image],
92
+ texts=[passage_text],
93
+ )
94
+ else:
95
+ # Text-only document
96
+ embedding = hf_model.model.encode_documents(
97
+ texts=[passage_text]
98
+ )
99
+ else:
100
+ raise ValueError(
101
+ f"Text must start with '{QUERY_PREFIX}' or '{PASSAGE_PREFIX}'"
102
+ )
103
+
104
+ hf_outputs.append(embedding[0].tolist())
105
+
106
+ check_embeddings_close(
107
+ embeddings_0_lst=hf_outputs,
108
+ embeddings_1_lst=vllm_outputs,
109
+ name_0="hf",
110
+ name_1="vllm",
111
+ )
112
+
113
+
114
+@pytest.mark.parametrize("model", MODELS)
115
+@pytest.mark.parametrize("dtype", ["half"])
116
+def test_models_text(
117
+ hf_runner,
118
+ vllm_runner,
119
+ image_assets,
120
+ model: str,
121
+ dtype: str,
122
+) -> None:
123
+ """Test text-only embedding."""
124
+ input_texts_images = [(text, None) for text in HF_TEXT_PROMPTS]
125
+ input_texts = [text for text, _ in input_texts_images]
126
+ input_images = [image for _, image in input_texts_images]
127
+
128
+ _run_test(
129
+ hf_runner,
130
+ vllm_runner,
131
+ input_texts,
132
+ input_images, # type: ignore
133
+ model,
134
+ dtype=dtype,
135
+ )
136
+
137
+
138
+@pytest.mark.parametrize("model", MODELS)
139
+@pytest.mark.parametrize("dtype", ["half"])
140
+def test_models_image(
141
+ hf_runner,
142
+ vllm_runner,
143
+ image_assets,
144
+ model: str,
145
+ dtype: str,
146
+) -> None:
147
+ """Test image embedding."""
148
+ input_texts_images = [
149
+ (text, asset.pil_image) for text, asset in zip(HF_IMAGE_PROMPTS, image_assets)
150
+ ]
151
+ input_texts = [text for text, _ in input_texts_images]
152
+ input_images = [image for _, image in input_texts_images]
153
+
154
+ _run_test(
155
+ hf_runner,
156
+ vllm_runner,
157
+ input_texts,
158
+ input_images,
159
+ model,
160
+ dtype=dtype,
161
+ )
162
+
163
+
164
+# ---------------------------------------------------------------------------
165
+# Reranker tests — nvidia/llama-nemotron-rerank-vl-1b-v2
166
+# ---------------------------------------------------------------------------
167
+
168
+RERANKER_MODELS = ["nvidia/llama-nemotron-rerank-vl-1b-v2"]
169
+
170
+# The tokenizer's built-in chat template is not suitable for the Score/Rerank
171
+# APIs (it's inherited from the base LLM). We must use the provided override.
172
+_RERANKER_SCORE_TEMPLATE = (
173
+ Path(__file__).parents[4]
174
+ / "examples/pooling/score/template/nemotron-vl-rerank.jinja"
175
+).read_text()
176
+
177
+RERANKER_TEXT_QUERY = "How is AI improving the intelligence and capabilities of robots?"
178
+RERANKER_TEXT_DOCS = [
179
+ "AI enables robots to perceive, plan, and act autonomously.",
180
+ (
181
+ "A biological foundation model designed to analyze DNA, RNA, "
182
+ "and protein sequences."
183
+ ),
184
+]
185
+
186
+RERANKER_IMAGE_QUERY = "photo of a red stop sign on a street"
187
+
188
+
189
+def _pil_to_data_uri(image) -> str:
190
+ buf = BytesIO()
191
+ image.save(buf, format="PNG")
192
+ b64 = base64.b64encode(buf.getvalue()).decode()
193
+ return f"data:image/png;base64,{b64}"
194
+
195
+
196
+def _run_hf_reranker(
197
+ hf_runner: type[HfRunner],
198
+ model: str,
199
+ dtype: str,
200
+ query: str,
201
+ docs: list,
202
+) -> list[float]:
203
+ """Run HF reranker inference; docs is a list of (doc_text, doc_image|None)."""
204
+ with hf_runner(
205
+ model,
206
+ dtype=dtype,
207
+ trust_remote_code=True,
208
+ auto_cls=AutoModelForSequenceClassification,
209
+ ) as hf_model:
210
+ processor = AutoProcessor.from_pretrained(
211
+ model,
212
+ trust_remote_code=True,
213
+ max_input_tiles=6,
214
+ use_thumbnail=True,
215
+ rerank_max_length=2048,
216
+ )
217
+ examples = [
218
+ {
219
+ "question": query,
220
+ "doc_text": doc_text if doc_text is not None else "",
221
+ "doc_image": doc_image if doc_image is not None else "",
222
+ }
223
+ for doc_text, doc_image in docs
224
+ ]
225
+ batch_dict = processor.process_queries_documents_crossencoder(examples)
226
+ batch_dict = {
227
+ k: v.to(hf_model.model.device) if isinstance(v, torch.Tensor) else v
228
+ for k, v in batch_dict.items()
229
+ }
230
+ with torch.inference_mode():
231
+ logits = hf_model.model(**batch_dict, return_dict=True).logits
232
+ # vLLM applies sigmoid activation to the raw logits before returning
233
+ # scores; apply the same here so both sides are comparable.
234
+ scores = torch.sigmoid(logits.squeeze(-1).float())
235
+ return scores.detach().cpu().tolist()
236
+
237
+
238
+def _run_vllm_reranker(
239
+ vllm_runner: type[VllmRunner],
240
+ model: str,
241
+ dtype: str,
242
+ query: str,
243
+ docs: list,
244
+) -> list[float]:
245
+ """Run vLLM reranker inference; docs is a list of (doc_text, doc_image|None)."""
246
+ with vllm_runner(
247
+ model,
248
+ runner="pooling",
249
+ dtype=dtype,
250
+ max_model_len=2048,
251
+ enforce_eager=True,
252
+ trust_remote_code=True,
253
+ ) as vllm_model:
254
+ has_images = any(img is not None for _, img in docs)
255
+
256
+ if not has_images:
257
+ # Text-only path: use the simple string score API.
258
+ queries = [query] * len(docs)
259
+ doc_texts = [doc_text for doc_text, _ in docs]
260
+ outputs = vllm_model.score(
261
+ queries,
262
+ doc_texts,
263
+ chat_template=_RERANKER_SCORE_TEMPLATE,
264
+ )
265
+ else:
266
+ # Multimodal path: build ScoreMultiModalParam for each pair.
267
+ query_params = [
268
+ ScoreMultiModalParam(
269
+ content=[
270
+ ChatCompletionContentPartTextParam(
271
+ type="text",
272
+ text=query,
273
+ )
274
+ ]
275
+ )
276
+ ] * len(docs)
277
+
278
+ doc_params = []
279
+ for doc_text, doc_image in docs:
280
+ content: list = []
281
+ if doc_image is not None:
282
+ content.append(
283
+ ChatCompletionContentPartImageParam(
284
+ type="image_url",
285
+ image_url={"url": _pil_to_data_uri(doc_image)},
286
+ )
287
+ )
288
+ if doc_text:
289
+ content.append(
290
+ ChatCompletionContentPartTextParam(
291
+ type="text",
292
+ text=doc_text,
293
+ )
294
+ )
295
+ doc_params.append(ScoreMultiModalParam(content=content))
296
+
297
+ raw_outputs = vllm_model.llm.score(
298
+ query_params,
299
+ doc_params,
300
+ chat_template=_RERANKER_SCORE_TEMPLATE,
301
+ )
302
+ outputs = [o.outputs.score for o in raw_outputs]
303
+
304
+ return outputs
305
+
306
+
307
+def _run_reranker_test(
308
+ hf_runner: type[HfRunner],
309
+ vllm_runner: type[VllmRunner],
310
+ model: str,
311
+ dtype: str,
312
+ query: str,
313
+ docs: list,
314
+) -> None:
315
+ """Compare HF and vLLM reranker scores.
316
+
317
+ NOTE: Run vLLM first to avoid CUDA initialization issues with multiprocessing.
318
+ """
319
+ vllm_scores = _run_vllm_reranker(vllm_runner, model, dtype, query, docs)
320
+ hf_scores = _run_hf_reranker(hf_runner, model, dtype, query, docs)
321
+
322
+ assert len(hf_scores) == len(vllm_scores), (
323
+ f"Output length mismatch: HF={len(hf_scores)}, vLLM={len(vllm_scores)}"
324
+ )
325
+ for i, (hf_score, vllm_score) in enumerate(zip(hf_scores, vllm_scores)):
326
+ assert hf_score == pytest.approx(vllm_score, rel=0.02), (
327
+ f"Score mismatch at index {i}: HF={hf_score:.4f}, vLLM={vllm_score:.4f}"
328
+ )
329
+
330
+
331
+@pytest.mark.parametrize("model", RERANKER_MODELS)
332
+@pytest.mark.parametrize("dtype", ["half"])
333
+def test_reranker_text(
334
+ hf_runner,
335
+ vllm_runner,
336
+ model: str,
337
+ dtype: str,
338
+) -> None:
339
+ """Test reranking with text-only query and text documents."""
340
+ docs = [(text, None) for text in RERANKER_TEXT_DOCS]
341
+ _run_reranker_test(hf_runner, vllm_runner, model, dtype, RERANKER_TEXT_QUERY, docs)
342
+
343
+
344
+@pytest.mark.parametrize("model", RERANKER_MODELS)
345
+@pytest.mark.parametrize("dtype", ["half"])
346
+def test_reranker_image_doc(
347
+ hf_runner,
348
+ vllm_runner,
349
+ image_assets,
350
+ model: str,
351
+ dtype: str,
352
+) -> None:
353
+ """Test reranking with text query against image documents."""
354
+ docs = [(None, asset.pil_image) for asset in image_assets]
355
+ _run_reranker_test(hf_runner, vllm_runner, model, dtype, RERANKER_IMAGE_QUERY, docs)
tests/models/multimodal/pooling/test_llama_nemotron_vl_embed.py
deleted
-148
@@ -1,148 +0,0 @@
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
@@ -653,6 +653,9 @@ _SEQUENCE_CLASSIFICATION_EXAMPLE_MODELS = {
653
"LlamaBidirectionalForSequenceClassification": _HfExamplesInfo(
654
"nvidia/llama-nemotron-rerank-1b-v2", trust_remote_code=True
655
),
656
+ "LlamaNemotronVLForSequenceClassification": _HfExamplesInfo(
657
+ "nvidia/llama-nemotron-rerank-vl-1b-v2", trust_remote_code=True
658
+ ),
659
"ModernBertForSequenceClassification": _HfExamplesInfo(
660
"Alibaba-NLP/gte-reranker-modernbert-base"
661
),
vllm/model_executor/models/config.py
+1
@@ -664,6 +664,7 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
664
"LlamaBidirectionalForSequenceClassification": LlamaBidirectionalConfig,
665
"LlamaBidirectionalModel": LlamaBidirectionalConfig,
666
"LlamaNemotronVLModel": LlamaNemotronVLConfig,
667
+ "LlamaNemotronVLForSequenceClassification": LlamaNemotronVLConfig,
668
"NomicBertModel": NomicBertModelConfig,
669
"Qwen2ForProcessRewardModel": Qwen2ForProcessRewardModelConfig,
670
"Qwen2ForRewardModel": Qwen2ForRewardModelConfig,
vllm/model_executor/models/nemotron_vl.py
+57
@@ -7,6 +7,7 @@
7
# Copyright (c) 2023 OpenGVLab
8
# Licensed under The MIT License [see LICENSE for details]
9
# --------------------------------------------------------
10
+import math
11
from abc import ABC
12
from collections.abc import Iterable
13
@@ -18,6 +19,7 @@ from transformers import AutoModel, PretrainedConfig
19
from transformers.image_processing_utils_fast import BaseImageProcessorFast
20
21
from vllm.config import VllmConfig
22
+from vllm.model_executor.layers.linear import ReplicatedLinear
23
from vllm.model_executor.layers.pooler import DispatchPooler
24
from vllm.model_executor.layers.quantization import QuantizationConfig
25
from vllm.model_executor.layers.quantization.awq import AWQConfig
@@ -42,6 +44,7 @@ from vllm.transformers_utils.repo_utils import get_hf_file_to_dict
44
45
from .interfaces import (
46
MultiModalEmbeddings,
47
+ SupportsCrossEncoding,
48
SupportsLoRA,
49
SupportsMultiModal,
50
SupportsPP,
@@ -883,3 +886,57 @@ class LlamaNemotronVLForEmbedding(LlamaNemotronVLChatModel, VllmModelForPooling)
886
"""Override to use different weight mapping for SigLIP."""
887
loader = AutoWeightsLoader(self)
888
return loader.load_weights(weights, mapper=self.weight_mapper)
889
+
890
+
891
+class LlamaNemotronVLForSequenceClassification(
892
+ LlamaNemotronVLForEmbedding, SupportsCrossEncoding
893
+):
894
+ """LlamaNemotronVL model variant for sequence classification / reranking."""
895
+
896
+ # Reranker checkpoint places base model weights under `model.*`,
897
+ # while `score.*` remains at the top level.
898
+ weight_mapper = WeightsMapper(orig_to_new_prefix={"model.": ""}) | (
899
+ LlamaNemotronVLForEmbedding.weight_mapper
900
+ )
901
+
902
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
903
+ super().__init__(vllm_config=vllm_config, prefix=prefix)
904
+
905
+ text_config = vllm_config.model_config.hf_config.get_text_config()
906
+ model_config = vllm_config.model_config
907
+ quant_config = vllm_config.quant_config
908
+
909
+ self.score = ReplicatedLinear(
910
+ model_config.get_hidden_size(),
911
+ text_config.num_labels,
912
+ bias=False,
913
+ params_dtype=model_config.head_dtype,
914
+ quant_config=quant_config,
915
+ return_bias=False,
916
+ prefix=maybe_prefix(prefix, "score"),
917
+ )
918
+
919
+ pooler_config = model_config.pooler_config
920
+ assert pooler_config is not None
921
+ self.pooler = DispatchPooler.for_seq_cls(pooler_config, classifier=self.score)
922
+
923
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
924
+ loaded_weights = super().load_weights(weights)
925
+
926
+ # reranker checkpoint omits the inner LM seq-cls head
927
+ # (`language_model.score.*`). It is unused by this outer model, but
928
+ # the default loader expects all parameters to be initialized.
929
+ for name, param in self.named_parameters():
930
+ if not name.startswith("language_model.score.") or name in loaded_weights:
931
+ continue
932
+
933
+ if name.endswith(".weight"):
934
+ torch.nn.init.kaiming_uniform_(param, a=math.sqrt(5))
935
+ elif name.endswith(".bias"):
936
+ torch.nn.init.zeros_(param)
937
+ else:
938
+ torch.nn.init.normal_(param, mean=0.0, std=0.02)
939
+
940
+ loaded_weights.add(name)
941
+
942
+ return loaded_weights
vllm/model_executor/models/registry.py
+4
@@ -284,6 +284,10 @@ _CROSS_ENCODER_MODELS = {
284
"llama",
285
"LlamaBidirectionalForSequenceClassification",
286
),
287
+ "LlamaNemotronVLForSequenceClassification": (
288
+ "nemotron_vl",
289
+ "LlamaNemotronVLForSequenceClassification",
290
+ ),
291
"ModernBertForSequenceClassification": (
292
"modernbert",
293
"ModernBertForSequenceClassification",