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
[Bugfix] Fix Dense module loading for sentence-transformers embedding models (simplified V2) (#23408)
Signed-off-by: FFFfff1FFFfff <yifanli0919@gmail.com>
LIYIFAN_liyifan committed
Aug 24, 2025 at 22:39 UTC
c9abb104893b16958d8c73f927ea37e65985f4ca
5 files changed
+175
-2
tests/models/language/pooling/test_st_projector.py
new
+22
@@ -0,0 +1,22 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import pytest
4
+
5
+from ...utils import CLSPoolingEmbedModelInfo, EmbedModelInfo
6
+from .mteb_utils import mteb_test_embed_models
7
+
8
+# ST models with projector (Dense) layers
9
+ST_PROJECTOR_MODELS = [
10
+ CLSPoolingEmbedModelInfo(
11
+ "TencentBAC/Conan-embedding-v1",
12
+ architecture="BertModel",
13
+ enable_test=True,
14
+ ),
15
+]
16
+
17
+
18
+@pytest.mark.parametrize("model_info", ST_PROJECTOR_MODELS)
19
+def test_embed_models_mteb(hf_runner, vllm_runner,
20
+ model_info: EmbedModelInfo) -> None:
21
+
22
+ mteb_test_embed_models(hf_runner, vllm_runner, model_info)
vllm/model_executor/layers/activation.py
+11
@@ -422,12 +422,23 @@ _ACTIVATION_REGISTRY = LazyDict({
422
lambda: nn.SiLU(),
423
"quick_gelu":
424
lambda: QuickGELU(),
425
+ "tanh":
426
+ lambda: nn.Tanh(),
427
+ "sigmoid":
428
+ lambda: nn.Sigmoid(),
429
})
430
431
432
def get_act_fn(act_fn_name: str) -> nn.Module:
433
"""Get an activation function by name."""
434
act_fn_name = act_fn_name.lower()
435
+
436
+ if act_fn_name.startswith("torch.nn.modules."):
437
+ activation_name = act_fn_name.split(".")[-1]
438
+ if activation_name == "identity":
439
+ return nn.Identity()
440
+ act_fn_name = activation_name
441
+
442
if act_fn_name not in _ACTIVATION_REGISTRY:
443
raise ValueError(
444
f"Activation function {act_fn_name!r} is not supported.")
vllm/model_executor/layers/pooler.py
+23
-1
@@ -5,7 +5,7 @@ from collections.abc import Mapping, Set
5
from dataclasses import dataclass
6
from enum import IntEnum
7
from itertools import groupby
8
-from typing import Callable, Optional, TypeVar, Union
8
+from typing import Callable, Optional, TypeVar, Union, cast
9
10
import torch
11
import torch.nn as nn
@@ -435,9 +435,31 @@ class EmbeddingPoolerHead(PoolerHead):
435
def __init__(self) -> None:
436
super().__init__(activation=PoolerNormalize())
437
438
+ # Load ST projector if available
439
+ from vllm.config import get_current_vllm_config
440
+ from vllm.model_executor.models.adapters import _load_st_projector
441
+
442
+ vllm_config = get_current_vllm_config()
443
+ self.projector = _load_st_projector(
444
+ vllm_config.model_config) if vllm_config else None
445
+
446
def forward(self, pooled_data: Union[list[torch.Tensor], torch.Tensor],
447
pooling_metadata: PoolingMetadata):
448
449
+ # Apply ST projector
450
+ if self.projector is not None:
451
+ projector = cast(nn.Module, self.projector)
452
+
453
+ def _proj(x: torch.Tensor) -> torch.Tensor:
454
+ orig_dtype = x.dtype
455
+ y = projector(x.to(torch.float32))
456
+ return y.to(orig_dtype)
457
+
458
+ if isinstance(pooled_data, torch.Tensor):
459
+ pooled_data = _proj(pooled_data)
460
+ else:
461
+ pooled_data = [_proj(t) for t in pooled_data]
462
+
463
pooling_params = get_pooling_params(pooling_metadata)
464
465
if isinstance(pooled_data, list):
vllm/model_executor/models/adapters.py
+97
-1
@@ -7,15 +7,21 @@ from typing import TYPE_CHECKING, Any, Optional, TypeVar, cast
7
import torch
8
import torch.nn as nn
9
10
+from vllm.logger import init_logger
11
+from vllm.model_executor.layers.activation import get_act_fn
12
from vllm.model_executor.models.config import VerifyAndUpdateConfig
13
+from vllm.transformers_utils.config import (get_hf_file_bytes,
14
+ get_hf_file_to_dict)
15
16
from .interfaces_base import VllmModelForPooling, is_pooling_model
17
18
if TYPE_CHECKING:
15
- from vllm.config import VllmConfig
19
+ from vllm.config import ModelConfig, VllmConfig
20
21
_T = TypeVar("_T", bound=type[nn.Module])
22
23
+logger = init_logger(__name__)
24
+
25
_GENERATE_SUFFIXES = [
26
"ForCausalLM",
27
"ForConditionalGeneration",
@@ -24,6 +30,96 @@ _GENERATE_SUFFIXES = [
30
]
31
32
33
+def _load_st_projector(model_config: "ModelConfig") -> Optional[nn.Module]:
34
+ """Load Sentence-Transformers Dense projection layers."""
35
+
36
+ try:
37
+ modules = get_hf_file_to_dict("modules.json", model_config.model,
38
+ model_config.revision)
39
+ if not modules:
40
+ return None
41
+
42
+ if isinstance(modules, dict):
43
+ modules = modules.get("modules", [])
44
+
45
+ dense_modules = [
46
+ m for m in modules
47
+ if m.get("type") == "sentence_transformers.models.Dense"
48
+ ]
49
+ if not dense_modules:
50
+ return None
51
+
52
+ module = dense_modules[0]
53
+ folder = module.get("path", "")
54
+
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
60
+
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)
65
+
66
+ if _load_dense_weights(linear, folder, model_config):
67
+ layers = [linear]
68
+ if act_name := layer_config.get("activation_function"):
69
+ layers.append(get_act_fn(act_name))
70
+ return nn.Sequential(*layers).to(dtype=torch.float32)
71
+
72
+ except Exception:
73
+ logger.exception("ST projector loading failed")
74
+
75
+ return None
76
+
77
+
78
+def _load_dense_weights(linear: nn.Linear, folder: str,
79
+ model_config: "ModelConfig") -> bool:
80
+ """Load weights using vLLM's weight_loader pattern."""
81
+ from vllm.model_executor.model_loader.weight_utils import (
82
+ default_weight_loader)
83
+
84
+ for filename in ["model.safetensors", "pytorch_model.bin"]:
85
+ file_path = f"{folder}/{filename}" if folder else filename
86
+
87
+ try:
88
+ file_bytes = get_hf_file_bytes(file_path, model_config.model,
89
+ model_config.revision)
90
+ if not file_bytes:
91
+ continue
92
+
93
+ if filename.endswith(".safetensors"):
94
+ from safetensors.torch import load as load_safetensors
95
+ state_dict = load_safetensors(file_bytes)
96
+ else:
97
+ import io
98
+ state_dict = torch.load(io.BytesIO(file_bytes),
99
+ map_location="cpu",
100
+ weights_only=True)
101
+
102
+ for weight_key in ["weight", "linear.weight", "dense.weight"]:
103
+ if weight_key in state_dict:
104
+ weight_loader = getattr(linear.weight, "weight_loader",
105
+ default_weight_loader)
106
+ weight_loader(linear.weight,
107
+ state_dict[weight_key].to(torch.float32))
108
+
109
+ bias_key = weight_key.replace("weight", "bias")
110
+ if linear.bias is not None and bias_key in state_dict:
111
+ bias_loader = getattr(linear.bias, "weight_loader",
112
+ default_weight_loader)
113
+ bias_loader(linear.bias,
114
+ state_dict[bias_key].to(torch.float32))
115
+ return True
116
+ except Exception:
117
+ logger.exception("Failed to load %s", filename)
118
+ continue
119
+
120
+ return False
121
+
122
+
123
def _get_pooling_model_name(orig_model_name: str, pooling_suffix: str) -> str:
124
model_name = orig_model_name
125
vllm/transformers_utils/config.py
+22
@@ -927,3 +927,25 @@ def get_model_path(model: Union[str, Path], revision: Optional[str] = None):
927
928
from huggingface_hub import snapshot_download
929
return snapshot_download(repo_id=model, **common_kwargs)
930
+
931
+
932
+def get_hf_file_bytes(file_name: str,
933
+ model: Union[str, Path],
934
+ revision: Optional[str] = 'main') -> Optional[bytes]:
935
+ """Get file contents from HuggingFace repository as bytes."""
936
+ file_path = try_get_local_file(model=model,
937
+ file_name=file_name,
938
+ revision=revision)
939
+
940
+ if file_path is None:
941
+ hf_hub_file = hf_hub_download(model,
942
+ file_name,
943
+ revision=revision,
944
+ token=_get_hf_token())
945
+ file_path = Path(hf_hub_file)
946
+
947
+ if file_path is not None and file_path.is_file():
948
+ with open(file_path, 'rb') as file:
949
+ return file.read()
950
+
951
+ return None