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
[New Model]: google/embeddinggemma-300m (#24318)
Signed-off-by: wang.yuqi <noooop@126.com>
wang.yuqi committed
Sep 6, 2025 at 13:58 UTC
6d6c6b05d37f12abe18cffdd9ade9f1fab864749
9 files changed
+73
-29
docs/models/supported_models.md
+1
@@ -440,6 +440,7 @@ These models primarily support the [`LLM.embed`](./pooling_models.md#llmembed) A
440
|--------------|--------|-------------------|----------------------|---------------------------|---------------------|
441
| `BertModel`<sup>C</sup> | BERT-based | `BAAI/bge-base-en-v1.5`, `Snowflake/snowflake-arctic-embed-xs`, etc. | | | ✅︎ |
442
| `Gemma2Model`<sup>C</sup> | Gemma 2-based | `BAAI/bge-multilingual-gemma2`, etc. | ✅︎ | ✅︎ | ✅︎ |
443
+| `Gemma3TextModel`<sup>C</sup> | Gemma 3-based | `google/embeddinggemma-300m`, etc. | ✅︎ | ✅︎ | ✅︎ |
444
| `GritLM` | GritLM | `parasail-ai/GritLM-7B-vllm`. | ✅︎ | ✅︎ | ✅︎ |
445
| `GteModel`<sup>C</sup> | Arctic-Embed-2.0-M | `Snowflake/snowflake-arctic-embed-m-v2.0`. | | | ✅︎ |
446
| `GteNewModel`<sup>C</sup> | mGTE-TRM (see note) | `Alibaba-NLP/gte-multilingual-base`, etc. | | | ✅︎ |
tests/models/language/pooling/mteb_utils.py
+16
-2
@@ -10,7 +10,8 @@ import numpy as np
10
import pytest
11
import requests
12
13
-from tests.models.utils import EmbedModelInfo, RerankModelInfo
13
+from tests.models.utils import (EmbedModelInfo, RerankModelInfo,
14
+ check_embeddings_close)
15
16
# Most embedding models on the STS12 task (See #17175):
17
# - Model implementation and minor changes in tensor dtype
@@ -163,12 +164,14 @@ def mteb_test_embed_models(hf_runner,
164
model_info: EmbedModelInfo,
165
vllm_extra_kwargs=None,
166
hf_model_callback=None,
166
- atol=MTEB_RERANK_TOL):
167
+ atol=MTEB_EMBED_TOL):
168
if not model_info.enable_test:
169
# A model family has many models with the same architecture,
170
# and we don't need to test each one.
171
pytest.skip("Skipping test.")
172
173
+ example_prompts = ["The chef prepared a delicious meal."]
174
+
175
vllm_extra_kwargs = vllm_extra_kwargs or {}
176
vllm_extra_kwargs["dtype"] = model_info.dtype
177
@@ -191,6 +194,7 @@ def mteb_test_embed_models(hf_runner,
194
vllm_main_score = run_mteb_embed_task(VllmMtebEncoder(vllm_model),
195
MTEB_EMBED_TASKS)
196
vllm_dtype = vllm_model.llm.llm_engine.model_config.dtype
197
+ vllm_outputs = vllm_model.embed(example_prompts)
198
199
if model_info.mteb_score is None:
200
with hf_runner(model_info.name,
@@ -202,6 +206,16 @@ def mteb_test_embed_models(hf_runner,
206
207
st_main_score = run_mteb_embed_task(hf_model, MTEB_EMBED_TASKS)
208
st_dtype = next(hf_model.model.parameters()).dtype
209
+
210
+ # Test embed_dims and whether to use normalize
211
+ hf_outputs = hf_model.encode(example_prompts)
212
+ check_embeddings_close(
213
+ embeddings_0_lst=hf_outputs,
214
+ embeddings_1_lst=vllm_outputs,
215
+ name_0="hf",
216
+ name_1="vllm",
217
+ tol=1e-2,
218
+ )
219
else:
220
st_main_score = model_info.mteb_score
221
st_dtype = "Constant"
tests/models/language/pooling/test_st_projector.py
+6
-1
@@ -2,7 +2,8 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import pytest
4
5
-from ...utils import CLSPoolingEmbedModelInfo, EmbedModelInfo
5
+from ...utils import (CLSPoolingEmbedModelInfo, EmbedModelInfo,
6
+ LASTPoolingEmbedModelInfo)
7
from .mteb_utils import mteb_test_embed_models
8
9
# ST models with projector (Dense) layers
@@ -13,6 +14,10 @@ ST_PROJECTOR_MODELS = [
14
mteb_score=0.688611955,
15
enable_test=True,
16
),
17
+ LASTPoolingEmbedModelInfo("google/embeddinggemma-300m",
18
+ architecture="Gemma3TextModel",
19
+ mteb_score=0.7473819294684156,
20
+ enable_test=True)
21
]
22
23
tests/models/registry.py
+1
@@ -352,6 +352,7 @@ _EMBEDDING_EXAMPLE_MODELS = {
352
# [Text-only]
353
"BertModel": _HfExamplesInfo("BAAI/bge-base-en-v1.5"),
354
"Gemma2Model": _HfExamplesInfo("BAAI/bge-multilingual-gemma2"), # noqa: E501
355
+ "Gemma3TextModel": _HfExamplesInfo("google/embeddinggemma-300m"),
356
"GritLM": _HfExamplesInfo("parasail-ai/GritLM-7B-vllm"),
357
"GteModel": _HfExamplesInfo("Snowflake/snowflake-arctic-embed-m-v2.0",
358
trust_remote_code=True),
vllm/config/__init__.py
+2
@@ -2750,6 +2750,8 @@ _STR_DTYPE_TO_TORCH_DTYPE = {
2750
_FLOAT16_NOT_SUPPORTED_MODELS = {
2751
"gemma2": "Numerical instability. Please use bfloat16 or float32 instead.",
2752
"gemma3": "Numerical instability. Please use bfloat16 or float32 instead.",
2753
+ "gemma3_text":
2754
+ "Numerical instability. Please use bfloat16 or float32 instead.",
2755
"plamo2": "Numerical instability. Please use bfloat16 or float32 instead.",
2756
"glm4": "Numerical instability. Please use bfloat16 or float32 instead.",
2757
}
vllm/model_executor/models/adapters.py
+17
-15
@@ -49,26 +49,28 @@ def _load_st_projector(model_config: "ModelConfig") -> Optional[nn.Module]:
49
if not dense_modules:
50
return None
51
52
- module = dense_modules[0]
53
- folder = module.get("path", "")
52
+ layers = []
53
+ for module in dense_modules:
54
+ folder = module.get("path", "")
55
+
56
+ config_path = f"{folder}/config.json" if folder else "config.json"
57
+ layer_config = get_hf_file_to_dict(config_path, model_config.model,
58
+ model_config.revision)
59
+ if not layer_config:
60
+ continue
61
55
- config_path = f"{folder}/config.json" if folder else "config.json"
56
- layer_config = get_hf_file_to_dict(config_path, model_config.model,
57
- model_config.revision)
58
- if not layer_config:
59
- return None
62
+ linear = nn.Linear(layer_config.get("in_features", 768),
63
+ layer_config.get("out_features", 768),
64
+ bias=layer_config.get("bias", True),
65
+ dtype=torch.float32)
66
61
- linear = nn.Linear(layer_config.get("in_features", 768),
62
- layer_config.get("out_features", 768),
63
- bias=layer_config.get("bias", True),
64
- dtype=torch.float32)
67
+ if not _load_dense_weights(linear, folder, model_config):
68
+ continue
69
66
- if _load_dense_weights(linear, folder, model_config):
67
- layers = [linear]
70
+ layers.append(linear)
71
if act_name := layer_config.get("activation_function"):
72
layers.append(get_act_fn(act_name))
70
- return nn.Sequential(*layers).to(dtype=torch.float32)
71
-
73
+ return nn.Sequential(*layers).to(dtype=torch.float32)
74
except Exception:
75
logger.exception("ST projector loading failed")
76
vllm/model_executor/models/config.py
+9
@@ -24,6 +24,14 @@ class VerifyAndUpdateConfig:
24
raise NotImplementedError
25
26
27
+class Gemma3TextModelConfig:
28
+
29
+ @staticmethod
30
+ def verify_and_update_config(vllm_config: "VllmConfig") -> None:
31
+ hf_config = vllm_config.model_config.hf_config
32
+ hf_config.is_causal = not hf_config.use_bidirectional_attention
33
+
34
+
35
class GteNewModelConfig(VerifyAndUpdateConfig):
36
37
@staticmethod
@@ -409,6 +417,7 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
417
"GteModel": SnowflakeGteNewModelConfig,
418
"GteNewModel": GteNewModelConfig,
419
"GteNewForSequenceClassification": GteNewModelConfig,
420
+ "Gemma3TextModel": Gemma3TextModelConfig,
421
"NomicBertModel": NomicBertModelConfig,
422
"Qwen2ForProcessRewardModel": Qwen2ForProcessRewardModelConfig,
423
"Qwen2ForRewardModel": Qwen2ForRewardModelConfig,
vllm/model_executor/models/gemma3.py
+20
-11
@@ -24,7 +24,7 @@ import torch.nn.functional as F
24
from torch import nn
25
from transformers import Gemma3TextConfig
26
27
-from vllm.attention import Attention
27
+from vllm.attention import Attention, AttentionType
28
from vllm.compilation.decorators import support_torch_compile
29
from vllm.config import CacheConfig, VllmConfig
30
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
@@ -44,6 +44,7 @@ from vllm.model_executor.model_loader.weight_utils import (
44
from vllm.model_executor.sampling_metadata import SamplingMetadata
45
from vllm.sequence import IntermediateTensors
46
47
+from ...attention.layers.encoder_only_attention import EncoderOnlyAttention
48
from .interfaces import SupportsLoRA, SupportsPP
49
from .utils import (AutoWeightsLoader, extract_layer_index,
50
is_pp_missing_parameter,
@@ -169,16 +170,24 @@ class Gemma3Attention(nn.Module):
170
rope_scaling=self.rope_scaling,
171
)
172
172
- # Initialize the attention.
173
- self.attn = Attention(self.num_heads,
174
- self.head_dim,
175
- self.scaling,
176
- num_kv_heads=self.num_kv_heads,
177
- cache_config=cache_config,
178
- quant_config=quant_config,
179
- logits_soft_cap=attn_logits_soft_cap,
180
- per_layer_sliding_window=sliding_window,
181
- prefix=f"{prefix}.attn")
173
+ if getattr(config, "is_causal", True):
174
+ attn_type = AttentionType.DECODER
175
+ else:
176
+ attn_type = AttentionType.ENCODER_ONLY
177
+
178
+ attn_cls = (EncoderOnlyAttention
179
+ if attn_type == AttentionType.ENCODER_ONLY else Attention)
180
+
181
+ self.attn = attn_cls(self.num_heads,
182
+ self.head_dim,
183
+ self.scaling,
184
+ num_kv_heads=self.num_kv_heads,
185
+ cache_config=cache_config,
186
+ quant_config=quant_config,
187
+ attn_type=attn_type,
188
+ logits_soft_cap=attn_logits_soft_cap,
189
+ per_layer_sliding_window=sliding_window,
190
+ prefix=f"{prefix}.attn")
191
192
def forward(
193
self,
vllm/model_executor/models/registry.py
+1
@@ -155,6 +155,7 @@ _EMBEDDING_MODELS = {
155
"BertModel": ("bert", "BertEmbeddingModel"),
156
"DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"),
157
"Gemma2Model": ("gemma2", "Gemma2ForCausalLM"),
158
+ "Gemma3TextModel": ("gemma3", "Gemma3Model"),
159
"GlmForCausalLM": ("glm", "GlmForCausalLM"),
160
"GPT2ForSequenceClassification": ("gpt2", "GPT2ForSequenceClassification"),
161
"GritLM": ("gritlm", "GritLM"),