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 9, 19:30 UTC
Support bge-m3 sparse embeddings and colbert embeddings (#14526)
Signed-off-by: Max de Bayser <mbayser@br.ibm.com> Signed-off-by: Max de Bayser <maxdebayser@gmail.com>
Maximilien de Bayser committed
Jan 22, 2026 at 12:52 UTC
ff365eea94add487cf91316087bcfa8f2b39c2b8
9 files changed
+393
-19
docs/models/pooling_models.md
+38
@@ -305,6 +305,44 @@ Expected output:
305
306
An OpenAI client example can be found here: [examples/pooling/embed/openai_embedding_matryoshka_fy_client.py](../../examples/pooling/embed/openai_embedding_matryoshka_fy_client.py)
307
308
+## Specific models
309
+
310
+### BAAI/bge-m3
311
+
312
+The `BAAI/bge-m3` model comes with extra weights for sparse and colbert embeddings but unfortunately in its `config.json`
313
+the architecture is declared as `XLMRobertaModel`, which makes `vLLM` load it as a vanilla ROBERTA model without the
314
+extra weights. To load the full model weights, override its architecture like this:
315
+
316
+```shell
317
+vllm serve BAAI/bge-m3 --hf-overrides '{"architectures": ["BgeM3EmbeddingModel"]}'
318
+```
319
+
320
+Then you obtain the sparse embeddings like this:
321
+
322
+```shell
323
+curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{
324
+ "model": "BAAI/bge-m3",
325
+ "task": "token_classify",
326
+ "input": ["What is BGE M3?", "Defination of BM25"]
327
+}'
328
+```
329
+
330
+Due to limitations in the the output schema, the output consists of a list of
331
+token scores for each token for each input. This means that you'll have to call
332
+`/tokenize` as well to be able to pair tokens with scores.
333
+Refer to the tests in `tests/models/language/pooling/test_bge_m3.py` to see how
334
+to do that.
335
+
336
+You can obtain the colbert embeddings like this:
337
+
338
+```shell
339
+curl -s http://localhost:8000/pooling -H "Content-Type: application/json" -d '{
340
+ "model": "BAAI/bge-m3",
341
+ "task": "token_embed",
342
+ "input": ["What is BGE M3?", "Defination of BM25"]
343
+}'
344
+```
345
+
346
## Deprecated Features
347
348
### Encode task
tests/models/language/pooling/embed_utils.py
+14
@@ -2,6 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
from collections.abc import Sequence
4
5
+import openai
6
import pytest
7
8
from tests.conftest import HfRunner
@@ -65,3 +66,16 @@ def correctness_test_embed_models(
66
hf_model_callback(hf_model)
67
68
run_embedding_correctness_test(hf_model, example_prompts, vllm_outputs)
69
+
70
+
71
+async def run_client_embeddings(
72
+ client: openai.AsyncOpenAI,
73
+ model_name: str,
74
+ queries: list[str],
75
+ instruction: str = "",
76
+) -> list[list[float]]:
77
+ outputs = await client.embeddings.create(
78
+ model=model_name,
79
+ input=[instruction + q for q in queries],
80
+ )
81
+ return [data.embedding for data in outputs.data]
tests/models/language/pooling/test_bge_m3.py
new
+170
@@ -0,0 +1,170 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import httpx
4
+import openai
5
+import pytest
6
+import pytest_asyncio
7
+import torch
8
+
9
+from ....utils import RemoteOpenAIServer
10
+from .embed_utils import run_client_embeddings
11
+
12
+MODEL_NAME = "BAAI/bge-m3"
13
+MAX_MODEL_LEN = 512
14
+
15
+
16
+# Example from https://huggingface.co/BAAI/bge-m3
17
+sentences_1 = ["What is BGE M3?", "Defination of BM25"]
18
+sentences_2 = [
19
+ "BGE M3 is an embedding model supporting dense retrieval, "
20
+ "lexical matching and multi-vector interaction.",
21
+ "BM25 is a bag-of-words retrieval function that ranks a set "
22
+ "of documents based on the query terms appearing in each document",
23
+]
24
+
25
+similarity_reference = [[0.6265, 0.3477], [0.3499, 0.678]]
26
+lexical_score_reference = [0.19554901123046875, 0.0]
27
+colbert_score_reference = [0.7797, 0.4620]
28
+
29
+
30
+@pytest.fixture(scope="module")
31
+def server():
32
+ args = [
33
+ "--max-model-len",
34
+ str(MAX_MODEL_LEN),
35
+ "--hf-overrides",
36
+ '{"architectures": ["BgeM3EmbeddingModel"]}',
37
+ ]
38
+
39
+ with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
40
+ yield remote_server
41
+
42
+
43
+@pytest_asyncio.fixture
44
+async def client(server):
45
+ async with server.get_async_client() as async_client:
46
+ yield async_client
47
+
48
+
49
+@pytest.mark.asyncio
50
+async def test_bge_m3_api_server_embedding(client: openai.AsyncOpenAI):
51
+ embeddings_list_1 = await run_client_embeddings(
52
+ client,
53
+ MODEL_NAME,
54
+ sentences_1,
55
+ )
56
+ embeddings_list_2 = await run_client_embeddings(
57
+ client,
58
+ MODEL_NAME,
59
+ sentences_2,
60
+ )
61
+
62
+ embeddings_1 = torch.tensor(embeddings_list_1)
63
+ embeddings_2 = torch.tensor(embeddings_list_2)
64
+ similarity = embeddings_1 @ embeddings_2.T
65
+
66
+ # reference values from BAAI/bge-m3 documentation
67
+ reference = torch.tensor(similarity_reference)
68
+
69
+ assert torch.allclose(similarity, reference, rtol=0.01)
70
+
71
+
72
+async def tokenize(client: openai.AsyncOpenAI, sentences: list[str]) -> list[list[int]]:
73
+ futures = []
74
+ for sentence in sentences:
75
+ futures.append(
76
+ client.post(
77
+ "../tokenize",
78
+ body={"model": MODEL_NAME, "prompt": sentence},
79
+ cast_to=httpx.Response,
80
+ )
81
+ )
82
+ return [(await future).json()["tokens"] for future in futures]
83
+
84
+
85
+async def sparse_embeddings(
86
+ client: openai.AsyncOpenAI, sentences: list[str]
87
+) -> list[dict[int, float]]:
88
+ all_tokens = await tokenize(client, sentences)
89
+ result = await client.post(
90
+ "../pooling",
91
+ body={"model": MODEL_NAME, "input": sentences, "task": "token_classify"},
92
+ cast_to=httpx.Response,
93
+ )
94
+ all_embeddings = [data["data"] for data in result.json()["data"]]
95
+
96
+ ret = []
97
+
98
+ for sent_tokens, sent_emb in zip(all_tokens, all_embeddings):
99
+ token_embs = dict[int, float]()
100
+ if sent_tokens[0] == 0:
101
+ sent_tokens = sent_tokens[1:]
102
+ for token, val in zip(sent_tokens, sent_emb):
103
+ token_embs[token] = max(val, token_embs.get(token, 0.0))
104
+ ret.append(token_embs)
105
+ return ret
106
+
107
+
108
+# Based on https://github.com/FlagOpen/FlagEmbedding/blob/6fd176266f2382878bcc69cd656cff425d52f49b/FlagEmbedding/inference/embedder/encoder_only/m3.py#L129
109
+def compute_lexical_matching_score(
110
+ lw1: dict[int, float], lw2: dict[int, float]
111
+) -> float:
112
+ scores = 0.0
113
+ for token, weight in lw1.items():
114
+ if token in lw2:
115
+ scores += weight * lw2[token]
116
+ return scores
117
+
118
+
119
+@pytest.mark.asyncio
120
+async def test_bge_m3_api_server_sparse_embedding(client: openai.AsyncOpenAI):
121
+ embeddings_1 = await sparse_embeddings(client, sentences_1)
122
+ embeddings_2 = await sparse_embeddings(client, sentences_2)
123
+
124
+ lexical_scores_1_0_x_2_0 = compute_lexical_matching_score(
125
+ embeddings_1[0], embeddings_2[0]
126
+ )
127
+ assert lexical_scores_1_0_x_2_0 == pytest.approx(
128
+ lexical_score_reference[0], rel=0.01
129
+ )
130
+
131
+ lexical_scores_1_0_x_1_1 = compute_lexical_matching_score(
132
+ embeddings_1[0], embeddings_1[1]
133
+ )
134
+ assert lexical_scores_1_0_x_1_1 == pytest.approx(
135
+ lexical_score_reference[1], rel=0.01
136
+ )
137
+
138
+
139
+# https://github.com/FlagOpen/FlagEmbedding/blob/6fd176266f2382878bcc69cd656cff425d52f49b/FlagEmbedding/inference/embedder/encoder_only/m3.py#L163
140
+def colbert_score(q_reps: torch.Tensor, p_reps: torch.Tensor) -> torch.Tensor:
141
+ token_scores = torch.einsum("in,jn->ij", q_reps, p_reps)
142
+ scores, _ = token_scores.max(-1)
143
+ scores = torch.sum(scores) / q_reps.size(0)
144
+ return scores
145
+
146
+
147
+@pytest.mark.asyncio
148
+async def test_bge_m3_api_server_multi_vector(client: openai.AsyncOpenAI):
149
+ result_1 = await client.post(
150
+ "../pooling",
151
+ body={"model": MODEL_NAME, "input": sentences_1, "task": "token_embed"},
152
+ cast_to=httpx.Response,
153
+ )
154
+ embeddings_1 = [torch.tensor(data["data"]) for data in result_1.json()["data"]]
155
+
156
+ result_2 = await client.post(
157
+ "../pooling",
158
+ body={"model": MODEL_NAME, "input": sentences_2, "task": "token_embed"},
159
+ cast_to=httpx.Response,
160
+ )
161
+ embeddings_2 = [torch.tensor(data["data"]) for data in result_2.json()["data"]]
162
+
163
+ colbert_score_1_0_x_2_0 = colbert_score(embeddings_1[0], embeddings_2[0])
164
+ assert colbert_score_1_0_x_2_0 == pytest.approx(
165
+ colbert_score_reference[0], rel=0.01
166
+ )
167
+ colbert_score_1_0_x_2_1 = colbert_score(embeddings_1[0], embeddings_2[1])
168
+ assert colbert_score_1_0_x_2_1 == pytest.approx(
169
+ colbert_score_reference[1], rel=0.01
170
+ )
tests/models/language/pooling/test_gritlm.py
+3
-13
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import numpy as np
4
-import openai
4
import pytest
5
from scipy.spatial.distance import cosine
6
@@ -9,6 +8,7 @@ from vllm import LLM, SamplingParams
8
from vllm.config import ModelConfig
9
10
from ....utils import RemoteOpenAIServer
11
+from .embed_utils import run_client_embeddings
12
13
MODEL_NAME = "parasail-ai/GritLM-7B-vllm"
14
MAX_MODEL_LEN = 4000
@@ -55,18 +55,6 @@ def run_llm_encode(
55
return [output.outputs.embedding for output in outputs]
56
57
58
-async def run_client_embeddings(
59
- client: openai.AsyncOpenAI,
60
- queries: list[str],
61
- instruction: str,
62
-) -> list[list[float]]:
63
- outputs = await client.embeddings.create(
64
- model=MODEL_NAME,
65
- input=[instruction + q for q in queries],
66
- )
67
- return [data.embedding for data in outputs.data]
68
-
69
-
58
def gritlm_instruction(instruction):
59
return (
60
"<|user|>\n" + instruction + "\n<|embed|>\n" if instruction else "<|embed|>\n"
@@ -145,11 +133,13 @@ async def test_gritlm_api_server_embedding():
133
134
d_rep = await run_client_embeddings(
135
client_embedding,
136
+ MODEL_NAME,
137
documents,
138
d_instruction,
139
)
140
q_rep = await run_client_embeddings(
141
client_embedding,
142
+ MODEL_NAME,
143
queries,
144
q_instruction,
145
)
tests/models/registry.py
+1
@@ -513,6 +513,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
513
_EMBEDDING_EXAMPLE_MODELS = {
514
# [Text-only]
515
"BertModel": _HfExamplesInfo("BAAI/bge-base-en-v1.5"),
516
+ "BgeM3EmbeddingModel": _HfExamplesInfo("BAAI/bge-m3"),
517
"Gemma2Model": _HfExamplesInfo("BAAI/bge-multilingual-gemma2"),
518
"Gemma3TextModel": _HfExamplesInfo("google/embeddinggemma-300m"),
519
"GritLM": _HfExamplesInfo("parasail-ai/GritLM-7B-vllm"),
vllm/model_executor/layers/pooler/special.py
+46
-1
@@ -125,4 +125,49 @@ class IdentityPooler(Pooler):
125
return hidden_states
126
127
128
-__all__ = ["DispatchPooler", "IdentityPooler"]
128
+class BOSEOSFilter(Pooler):
129
+ """Filters the BOS and EOS token results from outputs."""
130
+
131
+ def __init__(
132
+ self,
133
+ pooler: Pooler,
134
+ bos_token_id: int = -1, # -1 disables the filtering
135
+ eos_token_id: int = -1,
136
+ ) -> None:
137
+ super().__init__()
138
+
139
+ self.pooler = pooler
140
+ self.bos_token_id = bos_token_id
141
+ self.eos_token_id = eos_token_id
142
+
143
+ def get_supported_tasks(self) -> Set[PoolingTask]:
144
+ return self.pooler.get_supported_tasks()
145
+
146
+ def get_pooling_updates(self, task: PoolingTask) -> PoolingParamsUpdate:
147
+ return PoolingParamsUpdate(requires_token_ids=True)
148
+
149
+ def forward(
150
+ self,
151
+ hidden_states: torch.Tensor | list[torch.Tensor],
152
+ pooling_metadata: PoolingMetadata,
153
+ ) -> PoolerOutput:
154
+ pooled_outputs = self.pooler(hidden_states, pooling_metadata)
155
+ assert isinstance(pooled_outputs, list)
156
+
157
+ for i, prompt_len in enumerate(pooling_metadata.prompt_lens):
158
+ pooled_data = pooled_outputs[i]
159
+ assert (
160
+ isinstance(pooled_data, torch.Tensor)
161
+ and pooled_data.shape[0] == prompt_len
162
+ )
163
+ token_ids = pooling_metadata.prompt_token_ids[i, :prompt_len]
164
+ if token_ids[0] == self.bos_token_id:
165
+ pooled_data = pooled_data[1:]
166
+ if token_ids[-1] == self.eos_token_id:
167
+ pooled_data = pooled_data[:-1]
168
+ pooled_outputs[i] = pooled_data.squeeze()
169
+
170
+ return pooled_outputs
171
+
172
+
173
+__all__ = ["BOSEOSFilter", "DispatchPooler", "IdentityPooler"]
vllm/model_executor/layers/pooler/tokwise/poolers.py
+11
-3
@@ -6,7 +6,11 @@ from typing import TypeAlias
6
import torch
7
8
from vllm.config import PoolerConfig, get_current_vllm_config
9
-from vllm.model_executor.layers.pooler import ClassifierFn, PoolingParamsUpdate
9
+from vllm.model_executor.layers.pooler import (
10
+ ClassifierFn,
11
+ PoolingParamsUpdate,
12
+ ProjectorFn,
13
+)
14
from vllm.model_executor.layers.pooler.abstract import Pooler
15
from vllm.model_executor.layers.pooler.activations import (
16
PoolerActivation,
@@ -89,14 +93,18 @@ class TokenPooler(Pooler):
93
return pooled_data
94
95
92
-def pooler_for_token_embed(pooler_config: PoolerConfig):
96
+def pooler_for_token_embed(
97
+ pooler_config: PoolerConfig, projector: ProjectorFn | None = None
98
+) -> TokenPooler:
99
pooling = get_tok_pooling_method(pooler_config.get_tok_pooling_type())
100
101
vllm_config = get_current_vllm_config()
102
model_config = vllm_config.model_config
103
head = TokenEmbeddingPoolerHead(
104
head_dtype=model_config.head_dtype,
99
- projector=_load_st_projector(model_config),
105
+ projector=projector
106
+ if projector is not None
107
+ else _load_st_projector(model_config),
108
activation=PoolerNormalize(),
109
)
110
vllm/model_executor/models/registry.py
+1
@@ -234,6 +234,7 @@ _EMBEDDING_MODELS = {
234
"TeleChatForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
235
"TeleChat2ForCausalLM": ("telechat2", "TeleChat2ForCausalLM"),
236
"XLMRobertaModel": ("roberta", "RobertaEmbeddingModel"),
237
+ "BgeM3EmbeddingModel": ("roberta", "BgeM3EmbeddingModel"),
238
# [Multimodal]
239
"CLIPModel": ("clip", "CLIPEmbeddingModel"),
240
"LlavaNextForConditionalGeneration": (
vllm/model_executor/models/roberta.py
+109
-2
@@ -1,15 +1,30 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import itertools
5
from collections.abc import Iterable
6
7
import torch
8
from torch import nn
9
from transformers import RobertaConfig
10
10
-from vllm.config import ModelConfig, VllmConfig
11
-from vllm.model_executor.layers.pooler import DispatchPooler
11
+from vllm.config import ModelConfig, PoolerConfig, VllmConfig
12
+from vllm.model_executor.layers.pooler import (
13
+ BOSEOSFilter,
14
+ DispatchPooler,
15
+ Pooler,
16
+)
17
+from vllm.model_executor.layers.pooler.seqwise import (
18
+ pooler_for_embed,
19
+)
20
+from vllm.model_executor.layers.pooler.tokwise import (
21
+ AllPool,
22
+ pooler_for_token_classify,
23
+ pooler_for_token_embed,
24
+)
25
from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
26
+from vllm.model_executor.model_loader.default_loader import DefaultModelLoader
27
+from vllm.model_executor.model_loader.weight_utils import default_weight_loader
28
from vllm.model_executor.models.bert import (
29
TOKEN_TYPE_SHIFT,
30
BertEmbeddingModel,
@@ -149,6 +164,98 @@ class RobertaEmbeddingModel(BertEmbeddingModel):
164
return loader.load_weights(weights_list, mapper=mapper)
165
166
167
+def filter_secondary_weights(
168
+ all_weights: Iterable[tuple[str, torch.Tensor]],
169
+ secondary_weights: list[str],
170
+) -> tuple[Iterable[tuple[str, torch.Tensor]], Iterable[tuple[str, torch.Tensor]]]:
171
+ all_weights1, all_weights2 = itertools.tee(all_weights)
172
+
173
+ def filtered(n):
174
+ return any(n.startswith(f) for f in secondary_weights)
175
+
176
+ return ((n, w) for n, w in all_weights1 if filtered(n)), (
177
+ (n, w) for n, w in all_weights2 if not filtered(n)
178
+ )
179
+
180
+
181
+class BgeM3EmbeddingModel(RobertaEmbeddingModel):
182
+ """A model that extends RobertaEmbeddingModel with sparse embeddings.
183
+
184
+ This class supports loading an additional sparse_linear.pt file
185
+ to create sparse embeddings as described in https://arxiv.org/abs/2402.03216
186
+ """
187
+
188
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
189
+ self.hidden_size = vllm_config.model_config.hf_config.hidden_size
190
+
191
+ model_config = vllm_config.model_config
192
+ self.head_dtype = model_config.head_dtype
193
+ self.bos_token_id = model_config.hf_config.bos_token_id
194
+ self.eos_token_id = model_config.hf_config.eos_token_id
195
+
196
+ super().__init__(vllm_config=vllm_config, prefix=prefix)
197
+ self.secondary_weight_prefixes = ["sparse_linear.", "colbert_linear."]
198
+ self.secondary_weight_files = [
199
+ prefix + "pt" for prefix in self.secondary_weight_prefixes
200
+ ]
201
+
202
+ self.secondary_weights = [
203
+ DefaultModelLoader.Source(
204
+ model_or_path=vllm_config.model_config.model,
205
+ revision=None,
206
+ prefix=prefix,
207
+ allow_patterns_overrides=[filename],
208
+ )
209
+ for filename, prefix in zip(
210
+ self.secondary_weight_files, self.secondary_weight_prefixes
211
+ )
212
+ ]
213
+
214
+ def _build_pooler(self, pooler_config: PoolerConfig) -> Pooler:
215
+ self.sparse_linear = nn.Linear(self.hidden_size, 1, dtype=self.head_dtype)
216
+ self.colbert_linear = nn.Linear(
217
+ self.hidden_size, self.hidden_size, dtype=self.head_dtype
218
+ )
219
+
220
+ return DispatchPooler(
221
+ {
222
+ "embed": pooler_for_embed(pooler_config),
223
+ "token_embed": BOSEOSFilter(
224
+ pooler_for_token_embed(pooler_config, self.colbert_linear),
225
+ self.bos_token_id,
226
+ # for some reason m3 only filters the bos for colbert vectors
227
+ ),
228
+ "token_classify": BOSEOSFilter(
229
+ pooler_for_token_classify(
230
+ pooler_config,
231
+ pooling=AllPool(),
232
+ classifier=self.sparse_linear,
233
+ act_fn=torch.relu,
234
+ ),
235
+ self.bos_token_id,
236
+ self.eos_token_id,
237
+ ),
238
+ }
239
+ )
240
+
241
+ def load_weights(self, all_weights: Iterable[tuple[str, torch.Tensor]]):
242
+ secondary, weights = filter_secondary_weights(
243
+ all_weights, self.secondary_weight_prefixes
244
+ )
245
+
246
+ super().load_weights(weights)
247
+
248
+ params_dict = dict(self.named_parameters())
249
+
250
+ for name, loaded_weight in secondary:
251
+ if any(
252
+ name.startswith(prefix) for prefix in self.secondary_weight_prefixes
253
+ ):
254
+ param = params_dict[name]
255
+ weight_loader = getattr(param, "weight_loader", default_weight_loader)
256
+ weight_loader(param, loaded_weight)
257
+
258
+
259
@default_pooling_type(seq_pooling_type="CLS")
260
class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding):
261
"""A model that uses Roberta to provide embedding functionalities.