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 14, 05:30 UTC
[Doc] Update V1 status of various pooling models (#23189)
Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk>
Cyrus Leung committed
Aug 20, 2025 at 10:33 UTC
64ab3c7253afb8cc2008777153812109bf92d7c8
4 files changed
+28
-24
docs/models/supported_models.md
+13
-13
@@ -363,7 +363,7 @@ th {
363
| `GraniteMoeForCausalLM` | Granite 3.0 MoE, PowerMoE | `ibm-granite/granite-3.0-1b-a400m-base`, `ibm-granite/granite-3.0-3b-a800m-instruct`, `ibm/PowerMoE-3b`, etc. | ✅︎ | ✅︎ | ✅︎ |
364
| `GraniteMoeHybridForCausalLM` | Granite 4.0 MoE Hybrid | `ibm-granite/granite-4.0-tiny-preview`, etc. | ✅︎ | ✅︎ | ✅︎ |
365
| `GraniteMoeSharedForCausalLM` | Granite MoE Shared | `ibm-research/moe-7b-1b-active-shared-experts` (test model) | ✅︎ | ✅︎ | ✅︎ |
366
-| `GritLM` | GritLM | `parasail-ai/GritLM-7B-vllm`. | ✅︎ | ✅︎ | |
366
+| `GritLM` | GritLM | `parasail-ai/GritLM-7B-vllm`. | ✅︎ | ✅︎ | ✅︎ |
367
| `Grok1ModelForCausalLM` | Grok1 | `hpcai-tech/grok-1`. | ✅︎ | ✅︎ | ✅︎ |
368
| `HunYuanDenseV1ForCausalLM` | Hunyuan-7B-Instruct-0124 | `tencent/Hunyuan-7B-Instruct-0124` | ✅︎ | | ✅︎ |
369
| `HunYuanMoEV1ForCausalLM` | Hunyuan-80B-A13B | `tencent/Hunyuan-A13B-Instruct`, `tencent/Hunyuan-A13B-Pretrain`, `tencent/Hunyuan-A13B-Instruct-FP8`, etc. | ✅︎ | | ✅︎ |
@@ -436,17 +436,17 @@ These models primarily support the [`LLM.embed`](./pooling_models.md#llmembed) A
436
437
| Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) | [V1](gh-issue:8779) |
438
|--------------|--------|-------------------|----------------------|---------------------------|---------------------|
439
-| `BertModel`<sup>C</sup> | BERT-based | `BAAI/bge-base-en-v1.5`, `Snowflake/snowflake-arctic-embed-xs`, etc. | | | |
440
-| `Gemma2Model`<sup>C</sup> | Gemma 2-based | `BAAI/bge-multilingual-gemma2`, etc. | ✅︎ | | ✅︎ |
441
-| `GritLM` | GritLM | `parasail-ai/GritLM-7B-vllm`. | ✅︎ | ✅︎ | |
442
-| `GteModel`<sup>C</sup> | Arctic-Embed-2.0-M | `Snowflake/snowflake-arctic-embed-m-v2.0`. | | | |
443
-| `GteNewModel`<sup>C</sup> | mGTE-TRM (see note) | `Alibaba-NLP/gte-multilingual-base`, etc. | | | |
444
-| `ModernBertModel`<sup>C</sup> | ModernBERT-based | `Alibaba-NLP/gte-modernbert-base`, etc. | | | |
445
-| `NomicBertModel`<sup>C</sup> | Nomic BERT | `nomic-ai/nomic-embed-text-v1`, `nomic-ai/nomic-embed-text-v2-moe`, `Snowflake/snowflake-arctic-embed-m-long`, etc. | | | |
439
+| `BertModel`<sup>C</sup> | BERT-based | `BAAI/bge-base-en-v1.5`, `Snowflake/snowflake-arctic-embed-xs`, etc. | | | ✅︎ |
440
+| `Gemma2Model`<sup>C</sup> | Gemma 2-based | `BAAI/bge-multilingual-gemma2`, etc. | ✅︎ | ✅︎ | ✅︎ |
441
+| `GritLM` | GritLM | `parasail-ai/GritLM-7B-vllm`. | ✅︎ | ✅︎ | ✅︎ |
442
+| `GteModel`<sup>C</sup> | Arctic-Embed-2.0-M | `Snowflake/snowflake-arctic-embed-m-v2.0`. | | | ✅︎ |
443
+| `GteNewModel`<sup>C</sup> | mGTE-TRM (see note) | `Alibaba-NLP/gte-multilingual-base`, etc. | | | ✅︎ |
444
+| `ModernBertModel`<sup>C</sup> | ModernBERT-based | `Alibaba-NLP/gte-modernbert-base`, etc. | | | ✅︎ |
445
+| `NomicBertModel`<sup>C</sup> | Nomic BERT | `nomic-ai/nomic-embed-text-v1`, `nomic-ai/nomic-embed-text-v2-moe`, `Snowflake/snowflake-arctic-embed-m-long`, etc. | | | ✅︎ |
446
| `LlamaModel`<sup>C</sup>, `LlamaForCausalLM`<sup>C</sup>, `MistralModel`<sup>C</sup>, etc. | Llama-based | `intfloat/e5-mistral-7b-instruct`, etc. | ✅︎ | ✅︎ | ✅︎ |
447
| `Qwen2Model`<sup>C</sup>, `Qwen2ForCausalLM`<sup>C</sup> | Qwen2-based | `ssmits/Qwen2-7B-Instruct-embed-base` (see note), `Alibaba-NLP/gte-Qwen2-7B-instruct` (see note), etc. | ✅︎ | ✅︎ | ✅︎ |
448
| `Qwen3Model`<sup>C</sup>, `Qwen3ForCausalLM`<sup>C</sup> | Qwen3-based | `Qwen/Qwen3-Embedding-0.6B`, etc. | ✅︎ | ✅︎ | ✅︎ |
449
-| `RobertaModel`, `RobertaForMaskedLM` | RoBERTa-based | `sentence-transformers/all-roberta-large-v1`, etc. | | | |
449
+| `RobertaModel`, `RobertaForMaskedLM` | RoBERTa-based | `sentence-transformers/all-roberta-large-v1`, etc. | | | ✅︎ |
450
| `*Model`<sup>C</sup>, `*ForCausalLM`<sup>C</sup>, etc. | Generative models | N/A | \* | \* | \* |
451
452
<sup>C</sup> Automatically converted into an embedding model via `--convert embed`. ([details](./pooling_models.md#model-conversion))
@@ -476,7 +476,7 @@ These models primarily support the [`LLM.classify`](./pooling_models.md#llmclass
476
477
| Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) | [V1](gh-issue:8779) |
478
|--------------|--------|-------------------|----------------------|---------------------------|---------------------|
479
-| `JambaForSequenceClassification` | Jamba | `ai21labs/Jamba-tiny-reward-dev`, etc. | ✅︎ | ✅︎ | |
479
+| `JambaForSequenceClassification` | Jamba | `ai21labs/Jamba-tiny-reward-dev`, etc. | ✅︎ | ✅︎ | ✅︎ |
480
| `GPT2ForSequenceClassification` | GPT2 | `nie3e/sentiment-polish-gpt2-small` | | | ✅︎ |
481
| `*Model`<sup>C</sup>, `*ForCausalLM`<sup>C</sup>, etc. | Generative models | N/A | \* | \* | \* |
482
@@ -493,12 +493,12 @@ These models primarily support the [`LLM.score`](./pooling_models.md#llmscore) A
493
494
| Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) | [V1](gh-issue:8779) |
495
|--------------|--------|-------------------|----------------------|---------------------------|---------------------|
496
-| `BertForSequenceClassification` | BERT-based | `cross-encoder/ms-marco-MiniLM-L-6-v2`, etc. | | | |
496
+| `BertForSequenceClassification` | BERT-based | `cross-encoder/ms-marco-MiniLM-L-6-v2`, etc. | | | ✅︎ |
497
| `GemmaForSequenceClassification` | Gemma-based | `BAAI/bge-reranker-v2-gemma` (see note), etc. | ✅︎ | ✅︎ | ✅︎ |
498
| `Qwen2ForSequenceClassification` | Qwen2-based | `mixedbread-ai/mxbai-rerank-base-v2` (see note), etc. | ✅︎ | ✅︎ | ✅︎ |
499
| `Qwen3ForSequenceClassification` | Qwen3-based | `tomaarsen/Qwen3-Reranker-0.6B-seq-cls`, `Qwen/Qwen3-Reranker-0.6B` (see note), etc. | ✅︎ | ✅︎ | ✅︎ |
500
-| `RobertaForSequenceClassification` | RoBERTa-based | `cross-encoder/quora-roberta-base`, etc. | | | |
501
-| `XLMRobertaForSequenceClassification` | XLM-RoBERTa-based | `BAAI/bge-reranker-v2-m3`, etc. | | | |
500
+| `RobertaForSequenceClassification` | RoBERTa-based | `cross-encoder/quora-roberta-base`, etc. | | | ✅︎ |
501
+| `XLMRobertaForSequenceClassification` | XLM-RoBERTa-based | `BAAI/bge-reranker-v2-m3`, etc. | | | ✅︎ |
502
| `*Model`<sup>C</sup>, `*ForCausalLM`<sup>C</sup>, etc. | Generative models | N/A | \* | \* | \* |
503
504
<sup>C</sup> Automatically converted into a classification model via `--convert classify`. ([details](./pooling_models.md#model-conversion))
tests/models/language/pooling/test_gritlm.py
+5
-4
@@ -14,6 +14,7 @@ from ....utils import RemoteOpenAIServer
14
15
MODEL_NAME = "parasail-ai/GritLM-7B-vllm"
16
MAX_MODEL_LEN = 4000
17
+ATOL = 0.002
18
19
20
def _arr(arr):
@@ -97,16 +98,16 @@ def get_test_data():
98
99
def validate_embed_output(q_rep: list[list[float]], d_rep: list[list[float]]):
100
cosine_sim_q0_d0 = 1 - cosine(q_rep[0], d_rep[0])
100
- assert cosine_sim_q0_d0 == pytest.approx(0.609, abs=0.001)
101
+ assert cosine_sim_q0_d0 == pytest.approx(0.609, abs=ATOL)
102
103
cosine_sim_q0_d1 = 1 - cosine(q_rep[0], d_rep[1])
103
- assert cosine_sim_q0_d1 == pytest.approx(0.101, abs=0.001)
104
+ assert cosine_sim_q0_d1 == pytest.approx(0.101, abs=ATOL)
105
106
cosine_sim_q1_d0 = 1 - cosine(q_rep[1], d_rep[0])
106
- assert cosine_sim_q1_d0 == pytest.approx(0.120, abs=0.001)
107
+ assert cosine_sim_q1_d0 == pytest.approx(0.120, abs=ATOL)
108
109
cosine_sim_q1_d1 = 1 - cosine(q_rep[1], d_rep[1])
109
- assert cosine_sim_q1_d1 == pytest.approx(0.534, abs=0.001)
110
+ assert cosine_sim_q1_d1 == pytest.approx(0.534, abs=ATOL)
111
112
113
def test_gritlm_offline_embedding(vllm_runner):
vllm/model_executor/models/gritlm.py
+3
-3
@@ -20,7 +20,7 @@ from vllm.sequence import PoolerOutput
20
from vllm.tasks import PoolingTask
21
from vllm.transformers_utils.tokenizer import cached_tokenizer_from_config
22
23
-from .interfaces import SupportsV0Only
23
+from .interfaces import default_pooling_type
24
25
logger = init_logger(__name__)
26
@@ -215,7 +215,8 @@ class GritLMPooler(Pooler):
215
return build_output(pooled_data)
216
217
218
-class GritLM(LlamaForCausalLM, SupportsV0Only):
218
+@default_pooling_type("MEAN")
219
+class GritLM(LlamaForCausalLM):
220
"""This class implements the embedding model for parasail-ai/GritLM-7B-vllm.
221
222
The class inherits from LlamaForCausalLM and provides a custom pooling
@@ -241,7 +242,6 @@ class GritLM(LlamaForCausalLM, SupportsV0Only):
242
prefix: str = "",
243
**kwargs,
244
) -> None:
244
- # Use full attention for pooling (this is why V1 is not supported yet)
245
if vllm_config.model_config.runner_type == "pooling":
246
hf_config = vllm_config.model_config.hf_config
247
hf_config.is_causal = False
vllm/model_executor/models/interfaces.py
+7
-4
@@ -3,7 +3,7 @@
3
4
from collections.abc import Iterable, Mapping, MutableSequence
5
from typing import (TYPE_CHECKING, ClassVar, Literal, Optional, Protocol,
6
- Union, overload, runtime_checkable)
6
+ TypeVar, Union, overload, runtime_checkable)
7
8
import numpy as np
9
import torch
@@ -641,11 +641,14 @@ def supports_cross_encoding(
641
return is_pooling_model(model) and _supports_cross_encoding(model)
642
643
644
-def default_pooling_type(pooling_type: str) -> object:
644
+_T = TypeVar("_T", bound=type[torch.nn.Module])
645
+
646
+
647
+def default_pooling_type(pooling_type: str):
648
"""Set default_pooling_type decorator. """
649
647
- def func(model: object):
648
- model.default_pooling_type = pooling_type
650
+ def func(model: _T) -> _T:
651
+ model.default_pooling_type = pooling_type # type: ignore
652
return model
653
654
return func