[Bugfix] Support Sentence Transformers 5.4+ serialized configs (#54533)
Signed-off-by: maireneu <maireneu@gmail.com> Co-authored-by: Codex <codex@openai.com>
Eli Nabro Yeongmin Oh committed
Sep 2, 2026 at 00:53 UTC
cdefd9d4997f00da72dc6245cc60678b50761b7e
3 files changed
+150
-21
docs/models/pooling_models/README.md
+3
-3
@@ -292,12 +292,12 @@ classification activation is applied.
292
293
| Field | Source precedence | How to override |
294
| ----- | ----------------- | --------------- |
295
-| Pooling method (`pooling_type`) | `--pooler-config` > boolean `pooling_mode_*` fields in the Pooling module referenced by Sentence Transformers `modules.json` > architecture default (`LAST` for sequence pooling and `ALL` for token pooling unless the architecture overrides it) | Set `{"pooling_type": "CLS"}`, or set `seq_pooling_type` / `tok_pooling_type` explicitly. |
295
+| Pooling method (`pooling_type`) | `--pooler-config` > compact `pooling_mode` or legacy boolean `pooling_mode_*` fields in the Pooling module referenced by Sentence Transformers `modules.json` > architecture default (`LAST` for sequence pooling and `ALL` for token pooling unless the architecture overrides it) | Set `{"pooling_type": "CLS"}`, or set `seq_pooling_type` / `tok_pooling_type` explicitly. |
296
| Embedding normalization (`use_activation`) | `--pooler-config` > Sentence Transformers modules (`true` when a Normalize module is present, otherwise `false`) > pooling-task default (`true`) when no Sentence Transformers Pooling module is found | Set `{"use_activation": false}` to return unnormalized embeddings. |
297
| Classification activation function | Hugging Face `problem_type` > Sentence Transformers activation metadata > sigmoid or softmax selected from the label count | The function cannot be selected through `--pooler-config`; set `{"use_activation": false}` to return logits instead. |
298
299
-Sentence Transformers configurations using the newer compact `pooling_mode`
300
-string are not currently parsed; see [issue #45995](https://github.com/vllm-project/vllm/issues/45995).
299
+Both compact `pooling_mode` values written by Sentence Transformers 5.4+ and
300
+legacy boolean `pooling_mode_*` fields are supported.
301
302
For converted models and predefined models using the standard DispatchPooler
303
adapters, `embed` and `token_embed` construct an L2-normalization head, while
tests/test_config.py
+108
@@ -1,9 +1,11 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import json
5
import logging
6
import os
7
from dataclasses import MISSING, Field, asdict, dataclass, field
8
+from pathlib import Path
9
from types import SimpleNamespace
10
from typing import cast
11
from unittest.mock import patch
@@ -39,11 +41,20 @@ from vllm.config.speculative import _validate_qwen3_omni_dspark
41
from vllm.config.utils import get_field
42
from vllm.config.vllm import OPTIMIZATION_LEVEL_TO_CONFIG, OptimizationLevel
43
from vllm.platforms import current_platform
44
+from vllm.transformers_utils.config import (
45
+ get_pooling_config,
46
+ try_get_dense_modules,
47
+)
48
from vllm.v1.attention.backend import AttentionCGSupport
49
50
DEVICE_TYPE = current_platform.device_type
51
52
53
+def _write_json(path: Path, value: object) -> None:
54
+ path.parent.mkdir(parents=True, exist_ok=True)
55
+ path.write_text(json.dumps(value), encoding="utf-8")
56
+
57
+
58
def test_kda_recoverssm_derivation_is_revalidated():
59
config = SimpleNamespace(
60
cache_config=SimpleNamespace(
@@ -1031,6 +1042,103 @@ def test_get_pooling_config():
1042
assert model_config.pooler_config.tok_pooling_type == "ALL"
1043
1044
1045
+@pytest.mark.parametrize(
1046
+ ("pooling_module_type", "normalize_module_type", "pooling_config", "expected"),
1047
+ [
1048
+ (
1049
+ "sentence_transformers.models.Pooling",
1050
+ "sentence_transformers.models.Normalize",
1051
+ {"pooling_mode_mean_tokens": True},
1052
+ {"use_activation": True, "seq_pooling_type": "MEAN"},
1053
+ ),
1054
+ (
1055
+ "sentence_transformers.sentence_transformer.modules.pooling.Pooling",
1056
+ "sentence_transformers.sentence_transformer.modules.normalize.Normalize",
1057
+ {"pooling_mode": "lasttoken"},
1058
+ {"use_activation": True, "seq_pooling_type": "LAST"},
1059
+ ),
1060
+ (
1061
+ "sentence_transformers.sentence_transformer.modules.pooling.Pooling",
1062
+ "sentence_transformers.base.modules.normalize.Normalize",
1063
+ {"pooling_mode": "mean"},
1064
+ {"use_activation": True, "seq_pooling_type": "MEAN"},
1065
+ ),
1066
+ ],
1067
+)
1068
+def test_get_pooling_config_supports_sentence_transformers_schemas(
1069
+ tmp_path,
1070
+ caplog,
1071
+ pooling_module_type,
1072
+ normalize_module_type,
1073
+ pooling_config,
1074
+ expected,
1075
+):
1076
+ modules = [
1077
+ {"idx": 1, "name": "1", "path": "1_Pooling", "type": pooling_module_type},
1078
+ {
1079
+ "idx": 2,
1080
+ "name": "2",
1081
+ "path": "2_Normalize",
1082
+ "type": normalize_module_type,
1083
+ },
1084
+ ]
1085
+ _write_json(tmp_path / "modules.json", modules)
1086
+ _write_json(tmp_path / "1_Pooling" / "config.json", pooling_config)
1087
+
1088
+ with caplog.at_level(logging.WARNING):
1089
+ config = get_pooling_config(str(tmp_path))
1090
+
1091
+ assert config == expected
1092
+ assert "Unable to determine Sentence Transformers pooling type" not in caplog.text
1093
+
1094
+
1095
+def test_get_pooling_config_warns_when_pooling_mode_is_unknown(tmp_path, caplog):
1096
+ modules = [
1097
+ {
1098
+ "idx": 1,
1099
+ "name": "1",
1100
+ "path": "1_Pooling",
1101
+ "type": "sentence_transformers.models.Pooling",
1102
+ }
1103
+ ]
1104
+ _write_json(tmp_path / "modules.json", modules)
1105
+ _write_json(
1106
+ tmp_path / "1_Pooling" / "config.json",
1107
+ {"pooling_mode": "unsupported"},
1108
+ )
1109
+
1110
+ with caplog.at_level(logging.WARNING):
1111
+ config = get_pooling_config(str(tmp_path))
1112
+
1113
+ assert config == {"use_activation": False}
1114
+ assert "Unable to determine Sentence Transformers pooling type" in caplog.text
1115
+
1116
+
1117
+@pytest.mark.parametrize(
1118
+ "dense_module_type",
1119
+ [
1120
+ "sentence_transformers.models.Dense",
1121
+ "sentence_transformers.base.modules.dense.Dense",
1122
+ ],
1123
+)
1124
+def test_get_dense_modules_supports_sentence_transformers_schemas(
1125
+ tmp_path, dense_module_type
1126
+):
1127
+ modules = [{"idx": 2, "name": "2", "path": "2_Dense", "type": dense_module_type}]
1128
+ dense_config = {
1129
+ "in_features": 768,
1130
+ "out_features": 3072,
1131
+ "bias": False,
1132
+ "activation_function": "torch.nn.modules.linear.Identity",
1133
+ }
1134
+ _write_json(tmp_path / "modules.json", modules)
1135
+ _write_json(tmp_path / "2_Dense" / "config.json", dense_config)
1136
+
1137
+ assert try_get_dense_modules(str(tmp_path)) == [
1138
+ {**dense_config, "folder": "2_Dense"}
1139
+ ]
1140
+
1141
+
1142
@pytest.mark.skipif(
1143
current_platform.is_rocm(), reason="Xformers backend is not supported on ROCm."
1144
)
vllm/transformers_utils/config.py
+39
-18
@@ -52,6 +52,21 @@ MISTRAL_CONFIG_NAME = "params.json"
52
53
logger = init_logger(__name__)
54
55
+_ST_POOLING_MODULE_TYPES = {
56
+ "sentence_transformers.models.Pooling",
57
+ "sentence_transformers.sentence_transformer.modules.pooling.Pooling",
58
+}
59
+_ST_NORMALIZE_MODULE_TYPES = {
60
+ "sentence_transformers.models.Normalize",
61
+ "sentence_transformers.base.modules.normalize.Normalize",
62
+ "sentence_transformers.sentence_transformer.modules.normalize.Normalize",
63
+}
64
+_DENSE_MODULE_TYPES = {
65
+ "sentence_transformers.models.Dense",
66
+ "sentence_transformers.base.modules.dense.Dense",
67
+ "pylate.models.Dense.Dense",
68
+}
69
+
70
if Version(version("transformers")) < Version("5.0.0"):
71
raise ImportError(
72
"Support for Transformers v4 is deprecated and was removed in vLLM v0.24.0. "
@@ -852,11 +867,7 @@ def get_pooling_config(
867
logger.info("Found sentence-transformers modules configuration.")
868
869
pooling = next(
855
- (
856
- item
857
- for item in modules_dict
858
- if item["type"] == "sentence_transformers.models.Pooling"
859
- ),
870
+ (item for item in modules_dict if item["type"] in _ST_POOLING_MODULE_TYPES),
871
None,
872
)
873
normalize = bool(
@@ -864,7 +875,7 @@ def get_pooling_config(
875
(
876
item
877
for item in modules_dict
867
- if item["type"] == "sentence_transformers.models.Normalize"
878
+ if item["type"] in _ST_NORMALIZE_MODULE_TYPES
879
),
880
False,
881
)
@@ -880,14 +891,28 @@ def get_pooling_config(
891
892
config: dict[str, Any] = {"use_activation": normalize}
893
for key, val in pooling_dict.items():
883
- if val is True:
884
- pooling_type = parse_pooling_type(key)
885
- if pooling_type in SEQ_POOLING_TYPES:
886
- config["seq_pooling_type"] = pooling_type
887
- elif pooling_type in TOK_POOLING_TYPES:
888
- config["tok_pooling_type"] = pooling_type
889
- else:
890
- logger.debug("Skipping unrelated field: %r=%r", key, val)
894
+ if key == "pooling_mode" and isinstance(val, str):
895
+ pooling_name = val
896
+ elif val is True:
897
+ pooling_name = key
898
+ else:
899
+ continue
900
+
901
+ pooling_type = parse_pooling_type(pooling_name)
902
+ if pooling_type in SEQ_POOLING_TYPES:
903
+ config["seq_pooling_type"] = pooling_type
904
+ elif pooling_type in TOK_POOLING_TYPES:
905
+ config["tok_pooling_type"] = pooling_type
906
+ else:
907
+ logger.debug("Skipping unrelated field: %r=%r", key, val)
908
+
909
+ if not {"seq_pooling_type", "tok_pooling_type"} & config.keys():
910
+ logger.warning(
911
+ "Unable to determine Sentence Transformers pooling type from %s; "
912
+ "unless configured explicitly, vLLM will fall back to the model "
913
+ "architecture default.",
914
+ pooling_file_name,
915
+ )
916
917
return config
918
@@ -1151,10 +1176,6 @@ def try_get_dense_modules(
1176
if isinstance(modules, dict):
1177
modules = modules.get("modules", [])
1178
1154
- _DENSE_MODULE_TYPES = {
1155
- "sentence_transformers.models.Dense",
1156
- "pylate.models.Dense.Dense",
1157
- }
1179
dense_modules = [m for m in modules if m.get("type") in _DENSE_MODULE_TYPES]
1180
if not dense_modules:
1181
return None