[Bugfix] Fix Dense module loading for sentence-transformers embedding models v2

Signed-off-by: FFFfff1FFFfff <[email protected]>

Author: FFFfff1FFFfff

requirements/test.txt

@@ -968,7 +968,6 @@ setuptools==77.0.3
     # via
     #   lightning-utilities
     #   pytablewriter
-    #   torch
     #   triton
 shapely==2.1.1
     # via

vllm/model_executor/layers/pooler.py

@@ -72,6 +72,7 @@ def for_encode(
         pooler_config: PoolerConfig,
         *,
         default_pooling_type: PoolingType = PoolingType.ALL,
+        projector: Optional[nn.Module] = None,
     ):
         resolved_config = ResolvedPoolingConfig.from_config_with_defaults(
             task="encode",
@@ -82,21 +83,22 @@ def for_encode(
         if resolved_config.pooling_type == PoolingType.STEP:
             return StepPooler()
 
-        return SimplePooler.from_config(resolved_config)
+        return SimplePooler.from_config(resolved_config, projector=projector)
 
     @staticmethod
     def for_embed(
         pooler_config: PoolerConfig,
         *,
         default_pooling_type: PoolingType = PoolingType.LAST,
+        projector: Optional[nn.Module] = None,
     ):
         resolved_config = ResolvedPoolingConfig.from_config_with_defaults(
             task="embed",
             pooler_config=pooler_config,
             pooling_type=default_pooling_type,
         )
 
-        return SimplePooler.from_config(resolved_config)
+        return SimplePooler.from_config(resolved_config, projector=projector)
 
     @staticmethod
     def for_classify(
@@ -470,12 +472,32 @@ def forward(self, pooled_data: Union[list[torch.Tensor], torch.Tensor],
 
 class EmbeddingPoolerHead(PoolerHead):
 
-    def __init__(self) -> None:
+    def __init__(self, projector: Optional[nn.Module] = None) -> None:
         super().__init__(activation=PoolerNormalize())
+        self.projector = projector
+        self._device_set = False
 
     def forward(self, pooled_data: Union[list[torch.Tensor], torch.Tensor],
                 pooling_metadata: PoolingMetadata):
 
+        if self.projector is not None:
+            ref = pooled_data[0] if isinstance(pooled_data,
+                                               list) else pooled_data
+
+            if not self._device_set:
+                self.projector.to(device=ref.device, dtype=torch.float32)
+                self._device_set = True
+
+            def _proj(x: torch.Tensor) -> torch.Tensor:
+                y = self.projector(x.to(torch.float32))
+                return y.to(x.dtype)
+
+            if isinstance(pooled_data, torch.Tensor):
+                pooled_data = _proj(pooled_data)
+            else:
+                pooled_data = [_proj(t) for t in pooled_data]
+        # else: keep as is
+
         pooling_params = get_pooling_params(pooling_metadata)
 
         # for matryoshka representation
@@ -546,10 +568,11 @@ class SimplePooler(Pooler):
     def from_config(
         cls,
         pooler_config: ResolvedPoolingConfig,
+        projector: Optional[nn.Module] = None,
     ) -> "SimplePooler":
         pooling = PoolingMethod.from_pooling_type(pooler_config.pooling_type)
         if pooler_config.task == "embed":
-            head = EmbeddingPoolerHead()
+            head = EmbeddingPoolerHead(projector=projector)
         elif pooler_config.task == "encode":
             head = RewardPoolerHead()
         else:

vllm/model_executor/models/adapters.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import io
 from collections.abc import Iterable
 from typing import TYPE_CHECKING, Any, Optional, TypeVar, cast
 
@@ -23,6 +24,25 @@
     "LMHeadModel",
 ]
 
+# Note: projector uses standard nn.Linear to avoid sharding edge-cases
+
+
+def st_activation(name: Optional[str]) -> nn.Module:
+    m = (name or "").lower()
+    if m in ("gelu", "gelu_new"):
+        return nn.GELU()
+    if m == "relu":
+        return nn.ReLU()
+    if m == "tanh":
+        return nn.Tanh()
+    if m == "sigmoid":
+        return nn.Sigmoid()
+    if m == "swish":
+        return nn.SiLU()
+    if m == "identity":
+        return nn.Identity()
+    return nn.Identity()
+
 
 def _get_pooling_model_name(orig_model_name: str, pooling_suffix: str) -> str:
     model_name = orig_model_name
@@ -99,38 +119,137 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
     return ModelForPooling  # type: ignore
 
 
-def as_embedding_model(cls: _T) -> _T:
+def _load_st_projector(vllm_config: "VllmConfig") -> Optional[nn.Module]:
+    """Load Sentence Transformers projector from modules.json
+    and Dense folders.
     """
-    Subclass an existing vLLM model to support embeddings.
+    from vllm.transformers_utils.config import (file_or_path_exists,
+                                                get_hf_file_bytes,
+                                                get_hf_file_to_dict)
+
+    model_path = vllm_config.model_config.model
+    revision = vllm_config.model_config.revision
+
+    # Check if modules.json exists and contains Dense modules
+    if not file_or_path_exists(model_path, "modules.json", revision):
+        return None
+
+    modules = get_hf_file_to_dict("modules.json", model_path, revision)
+    if not isinstance(modules, list):
+        return None
+
+    dense_entries = [
+        m for m in modules
+        if m.get("type") == "sentence_transformers.models.Dense"
+    ]
+    if not dense_entries:
+        return None
+
+    # Get dtype and quant config
+    raw_dtype = getattr(vllm_config.model_config, "dtype", None)
+    if isinstance(raw_dtype, str):
+        desired_dtype = getattr(torch, raw_dtype, torch.float32)
+    elif isinstance(raw_dtype, torch.dtype):
+        desired_dtype = raw_dtype
+    else:
+        desired_dtype = torch.float32
+
+    def _load_config_json(path: str) -> Optional[dict]:
+        """Load config.json from a Dense folder."""
+        try:
+            return get_hf_file_to_dict(path, model_path, revision)
+        except Exception:
+            return None
+
+    def _load_dense_weights(
+        folder: str,
+    ) -> tuple[Optional[torch.Tensor], Optional[torch.Tensor]]:
+        """Load weights from a Dense folder, trying safetensors first."""
+        # Try safetensors first
+        try:
+            b = get_hf_file_bytes(f"{folder}/model.safetensors", model_path,
+                                  revision)
+            if b is not None:
+                from safetensors.torch import load as st_load
+                sd = st_load(io.BytesIO(b))
+                w = (sd.get("linear.weight") or sd.get("dense.weight")
+                     or sd.get("weight"))
+                bias = (sd.get("linear.bias") or sd.get("dense.bias")
+                        or sd.get("bias"))
+                return w, bias
+        except Exception:
+            pass
+
+        # Fallback to pytorch_model.bin
+        try:
+            b = get_hf_file_bytes(f"{folder}/pytorch_model.bin", model_path,
+                                  revision)
+            if b is not None:
+                sd = torch.load(io.BytesIO(b), map_location="cpu")
+                w = (sd.get("linear.weight") or sd.get("dense.weight")
+                     or sd.get("weight"))
+                bias = (sd.get("linear.bias") or sd.get("dense.bias")
+                        or sd.get("bias"))
+                return w, bias
+        except Exception:
+            pass
+        return None, None
+
+    # Build projector layers
+    layers: list[nn.Module] = []
+    for i, entry in enumerate(dense_entries):
+        folder = entry.get("path")
+        if not folder:
+            continue
+
+        cfg = _load_config_json(f"{folder}/config.json")
+        if not cfg:
+            continue
+
+        in_features = cfg.get("in_features")
+        out_features = cfg.get("out_features")
+        if in_features is None or out_features is None:
+            continue
+
+        use_bias = cfg.get("bias", True)
+        activation = st_activation(cfg.get("activation_function"))
+
+        # Create a simple nn.Linear for projector to avoid sharding edge-cases
+        linear = nn.Linear(in_features, out_features, bias=use_bias)
+        linear = linear.to(dtype=desired_dtype)
+
+        # Load weights
+        weight, bias = _load_dense_weights(folder)
+        if weight is not None:
+            with torch.no_grad():
+                # weight is expected in [out_features, in_features]
+                linear.weight.copy_(weight.to(dtype=linear.weight.dtype))
+                if use_bias and bias is not None and linear.bias is not None:
+                    linear.bias.copy_(bias.to(dtype=linear.bias.dtype))
+
+        layers.append(linear)
+        layers.append(activation)
+
+    if not layers:
+        return None
+    return nn.Sequential(*layers)
 
-    By default, the embeddings of the whole prompt are extracted from the
-    normalized hidden state corresponding to the last token.
 
-    Note:
-        We assume that no extra layers are added to the original model;
-        please implement your own model if this is not the case.
-    """
-    # Avoid modifying existing embedding models
-    if is_pooling_model(cls):
-        return cls
-
-    # Lazy import
-    from vllm.model_executor.layers.pooler import DispatchPooler, Pooler
+def as_embedding_model(cls: _T) -> _T:
+    """Convert a model class to support embedding tasks."""
 
     class ModelForEmbedding(_create_pooling_model_cls(cls)):
 
         def _init_pooler(self, vllm_config: "VllmConfig", prefix: str = ""):
-            pooler_config = vllm_config.model_config.pooler_config
-            assert pooler_config is not None
+            from vllm.model_executor.layers.pooler import Pooler
 
-            self.pooler = DispatchPooler(
-                {
-                    "encode": Pooler.for_encode(pooler_config),
-                    "embed": Pooler.for_embed(pooler_config),
-                }, )
+            # Load ST projector if available
+            projector = _load_st_projector(vllm_config)
 
-    ModelForEmbedding.__name__ = \
-        _get_pooling_model_name(cls.__name__, "ForEmbedding")
+            # Use existing pooler_config instead of creating new one
+            pooler_config = vllm_config.model_config.pooler_config
+            assert pooler_config is not None
+            self.pooler = Pooler.for_embed(pooler_config, projector=projector)
 
     return ModelForEmbedding  # type: ignore
 
@@ -229,7 +348,6 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
                 # ForSequenceClassification model.
                 return seq_cls_model_loader(self, weights)
 
-
     ModelForSequenceClassification.__name__ = \
         _get_pooling_model_name(cls.__name__, "ForSequenceClassification")
 

vllm/model_executor/models/bert.py

@@ -24,6 +24,7 @@
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     VocabParallelEmbedding)
+from vllm.model_executor.models.adapters import _load_st_projector
 from vllm.model_executor.pooling_metadata import PoolingMetadata
 from vllm.sequence import IntermediateTensors
 from vllm.tasks import PoolingTask
@@ -457,6 +458,9 @@ class BertEmbeddingModel(nn.Module, SupportsQuant):
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
 
+        # Save vllm_config for projector loading
+        self.vllm_config = vllm_config
+
         pooler_config = vllm_config.model_config.pooler_config
         assert pooler_config is not None
 
@@ -497,13 +501,19 @@ def _build_model(self,
                          embedding_class=BertEmbedding)
 
     def _build_pooler(self, pooler_config: PoolerConfig) -> Pooler:
+        # Load projector using the stored vllm_config
+        projector = None
+        if hasattr(self, 'vllm_config'):
+            projector = _load_st_projector(self.vllm_config)
+
         return DispatchPooler({
             "encode":
             Pooler.for_encode(pooler_config),
             "embed":
             Pooler.for_embed(
                 pooler_config,
                 default_pooling_type=PoolingType.CLS,
+                projector=projector,
             ),
         })