@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 071d863e20
Fork 0

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 8, 07:30 UTC

Extend ColBERT support to non-standard BERT backbones (#34170)

Signed-off-by: Ilya Boytsov <ilya.boytsov@aleph-alpha.com>

Ilya Boytsov committed Feb 13, 2026 at 10:53 UTC 071d863e208b40fa1bb986ad230e322b2bbbbcf5
9 files changed +793 -309
docs/models/pooling_models.md
+18 -7
@@ -311,20 +311,31 @@ An OpenAI client example can be found here: [examples/pooling/embed/openai_embed
311
312 [ColBERT](https://arxiv.org/abs/2004.12832) (Contextualized Late Interaction over BERT) is a retrieval model that uses per-token embeddings and MaxSim scoring for document ranking. Unlike single-vector embedding models, ColBERT retains token-level representations and computes relevance scores through late interaction, providing better accuracy while being more efficient than cross-encoders.
313
314 -vLLM supports ColBERT models for reranking tasks, automatically applying MaxSim scoring for query-document relevance:
314 +vLLM supports ColBERT models with multiple encoder backbones:
315 +
316 +| Architecture | Backbone | Example HF Models |
317 +|---|---|---|
318 +| `HF_ColBERT` | BERT | `answerdotai/answerai-colbert-small-v1`, `colbert-ir/colbertv2.0` |
319 +| `ColBERTModernBertModel` | ModernBERT | `lightonai/GTE-ModernColBERT-v1` |
320 +| `ColBERTJinaRobertaModel` | Jina XLM-RoBERTa | `jinaai/jina-colbert-v2` |
321 +
322 +**BERT-based ColBERT** models work out of the box:
323
324 ```shell
325 vllm serve answerdotai/answerai-colbert-small-v1
326 ```
327
320 -Currently supports ColBERT models with standard BERT encoders (e.g., `answerdotai/answerai-colbert-small-v1`, `colbert-ir/colbertv2.0`).
321 -
322 -ColBERT models with modified encoder architectures are not yet supported, including BERT variants with rotary embeddings (e.g., `jinaai/jina-colbert-v2`) or other custom encoders (e.g., `LiquidAI/LFM2-ColBERT-350M`).
323 -
324 -If your standard BERT ColBERT model's config doesn't specify the architecture as `HF_ColBERT`, override it with:
328 +For **non-BERT backbones**, use `--hf-overrides` to set the correct architecture:
329
330 ```shell
327 -vllm serve your-colbert-model --hf-overrides '{"architectures": ["HF_ColBERT"]}'
331 +# ModernBERT backbone
332 +vllm serve lightonai/GTE-ModernColBERT-v1 \
333 + --hf-overrides '{"architectures": ["ColBERTModernBertModel"]}'
334 +
335 +# Jina XLM-RoBERTa backbone
336 +vllm serve jinaai/jina-colbert-v2 \
337 + --hf-overrides '{"architectures": ["ColBERTJinaRobertaModel"]}' \
338 + --trust-remote-code
339 ```
340
341 Then you can use the rerank endpoint:
examples/pooling/score/colbert_rerank_online.py
+64 -29
@@ -1,15 +1,27 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """
4 -Example of using ColBERT late interaction model for reranking.
4 +Example of using ColBERT late interaction models for reranking and scoring.
5
6 ColBERT (Contextualized Late Interaction over BERT) uses per-token embeddings
7 and MaxSim scoring for document reranking, providing better accuracy than
8 single-vector models while being more efficient than cross-encoders.
9
10 -Start the server with:
10 +vLLM supports ColBERT with multiple encoder backbones. Start the server
11 +with one of the following:
12 +
13 + # BERT backbone (works out of the box)
14 vllm serve answerdotai/answerai-colbert-small-v1
15
16 + # ModernBERT backbone
17 + vllm serve lightonai/GTE-ModernColBERT-v1 \
18 + --hf-overrides '{"architectures": ["ColBERTModernBertModel"]}'
19 +
20 + # Jina XLM-RoBERTa backbone
21 + vllm serve jinaai/jina-colbert-v2 \
22 + --hf-overrides '{"architectures": ["ColBERTJinaRobertaModel"]}' \
23 + --trust-remote-code
24 +
25 Then run this script:
26 python colbert_rerank_online.py
27 """
@@ -18,39 +30,62 @@ import json
30
31 import requests
32
21 -url = "http://127.0.0.1:8000/rerank"
33 +# Change this to match the model you started the server with
34 +MODEL = "answerdotai/answerai-colbert-small-v1"
35 +BASE_URL = "http://127.0.0.1:8000"
36
37 headers = {"accept": "application/json", "Content-Type": "application/json"}
38
25 -data = {
26 - "model": "answerdotai/answerai-colbert-small-v1",
27 - "query": "What is machine learning?",
28 - "documents": [
29 - "Machine learning is a subset of artificial intelligence.",
30 - "Python is a programming language.",
31 - "Deep learning uses neural networks for complex tasks.",
32 - "The weather today is sunny.",
33 - ],
34 -}
39 +documents = [
40 + "Machine learning is a subset of artificial intelligence.",
41 + "Python is a programming language.",
42 + "Deep learning uses neural networks for complex tasks.",
43 + "The weather today is sunny.",
44 +]
45 +
46 +
47 +def rerank_example():
48 + """Use the /rerank endpoint to rank documents by query relevance."""
49 + print("=== Rerank Example ===")
50 +
51 + data = {
52 + "model": MODEL,
53 + "query": "What is machine learning?",
54 + "documents": documents,
55 + }
56 +
57 + response = requests.post(f"{BASE_URL}/rerank", headers=headers, json=data)
58 + result = response.json()
59 + print(json.dumps(result, indent=2))
60 +
61 + print("\nRanked documents (most relevant first):")
62 + for item in result["results"]:
63 + doc_idx = item["index"]
64 + score = item["relevance_score"]
65 + print(f" Score {score:.4f}: {documents[doc_idx]}")
66 +
67 +
68 +def score_example():
69 + """Use the /score endpoint for pairwise query-document scoring."""
70 + print("\n=== Score Example ===")
71 +
72 + data = {
73 + "model": MODEL,
74 + "text_1": "What is machine learning?",
75 + "text_2": [
76 + "Machine learning is a subset of AI.",
77 + "The weather is sunny.",
78 + ],
79 + }
80 +
81 + response = requests.post(f"{BASE_URL}/score", headers=headers, json=data)
82 + result = response.json()
83 + print(json.dumps(result, indent=2))
84
85
86 def main():
38 - response = requests.post(url, headers=headers, json=data)
39 -
40 - if response.status_code == 200:
41 - print("ColBERT Rerank Request successful!")
42 - result = response.json()
43 - print(json.dumps(result, indent=2))
44 -
45 - # Show ranked results
46 - print("\nRanked documents (most relevant first):")
47 - for item in result["results"]:
48 - doc_idx = item["index"]
49 - score = item["relevance_score"]
50 - print(f" Score {score:.4f}: {data['documents'][doc_idx]}")
51 - else:
52 - print(f"Request failed with status code: {response.status_code}")
53 - print(response.text)
87 + rerank_example()
88 + score_example()
89
90
91 if __name__ == "__main__":
tests/entrypoints/pooling/score/test_online_colbert.py
+117 -129
@@ -8,10 +8,8 @@ import requests
8 from tests.utils import RemoteOpenAIServer
9 from vllm.entrypoints.pooling.score.protocol import RerankResponse, ScoreResponse
10
11 -# ColBERT model - using answerai-colbert-small-v1 as it's a smaller model
11 MODEL_NAME = "answerdotai/answerai-colbert-small-v1"
13 -COLBERT_DIM = 96 # This model uses 96-dimensional output
14 -DTYPE = "half"
12 +COLBERT_DIM = 96
13 MAX_MODEL_LEN = 512
14
15
@@ -26,129 +24,119 @@ def server():
24 yield remote_server
25
26
29 -@pytest.mark.parametrize("model_name", [MODEL_NAME])
30 -def test_colbert_rerank(server: RemoteOpenAIServer, model_name: str):
31 - """Test ColBERT rerank endpoint."""
32 - query = "What is the capital of France?"
33 - documents = [
34 - "The capital of Brazil is Brasilia.",
35 - "The capital of France is Paris.",
36 - ]
37 -
38 - rerank_response = requests.post(
39 - server.url_for("rerank"),
40 - json={
41 - "model": model_name,
42 - "query": query,
43 - "documents": documents,
44 - },
45 - )
46 - rerank_response.raise_for_status()
47 - rerank = RerankResponse.model_validate(rerank_response.json())
48 -
49 - assert rerank.id is not None
50 - assert rerank.results is not None
51 - assert len(rerank.results) == 2
52 -
53 - # The relevant document (Paris) should have higher score
54 - paris_result = next(r for r in rerank.results if r.index == 1)
55 - brazil_result = next(r for r in rerank.results if r.index == 0)
56 -
57 - assert paris_result.relevance_score > brazil_result.relevance_score
58 -
59 -
60 -@pytest.mark.parametrize("model_name", [MODEL_NAME])
61 -def test_colbert_rerank_top_n(server: RemoteOpenAIServer, model_name: str):
62 - """Test ColBERT rerank with top_n parameter."""
63 - query = "What is the capital of France?"
64 - documents = [
65 - "The capital of Brazil is Brasilia.",
66 - "The capital of France is Paris.",
67 - "Machine learning is a field of AI.",
68 - ]
69 -
70 - rerank_response = requests.post(
71 - server.url_for("rerank"),
72 - json={
73 - "model": model_name,
74 - "query": query,
75 - "documents": documents,
76 - "top_n": 2,
77 - },
78 - )
79 - rerank_response.raise_for_status()
80 - rerank = RerankResponse.model_validate(rerank_response.json())
81 -
82 - assert len(rerank.results) == 2
83 - # Top result should be about Paris (index 1)
84 - assert rerank.results[0].index == 1
85 -
86 -
87 -@pytest.mark.parametrize("model_name", [MODEL_NAME])
88 -def test_colbert_score(server: RemoteOpenAIServer, model_name: str):
89 - """Test ColBERT score endpoint."""
90 - text_1 = "What is the capital of France?"
91 - text_2 = ["The capital of France is Paris.", "Python is a language."]
92 -
93 - score_response = requests.post(
94 - server.url_for("score"),
95 - json={
96 - "model": model_name,
97 - "text_1": text_1,
98 - "text_2": text_2,
99 - },
100 - )
101 - score_response.raise_for_status()
102 - score = ScoreResponse.model_validate(score_response.json())
103 -
104 - assert score.id is not None
105 - assert score.data is not None
106 - assert len(score.data) == 2
107 -
108 - # The relevant document should have higher score
109 - assert score.data[0].score > score.data[1].score
110 -
111 -
112 -@pytest.mark.parametrize("model_name", [MODEL_NAME])
113 -def test_colbert_token_embed(server: RemoteOpenAIServer, model_name: str):
114 - """Test ColBERT token_embed task via pooling endpoint."""
115 - text = "What is the capital of France?"
116 -
117 - pooling_response = requests.post(
118 - server.url_for("pooling"),
119 - json={
120 - "model": model_name,
121 - "input": text,
122 - "task": "token_embed",
123 - },
124 - )
125 - pooling_response.raise_for_status()
126 - pooling = pooling_response.json()
127 -
128 - assert "data" in pooling
129 - assert len(pooling["data"]) == 1
130 -
131 - # Token embeddings should be 2D
132 - embeddings = pooling["data"][0]["data"]
133 - assert isinstance(embeddings, list)
134 - assert len(embeddings) > 0 # Should have tokens
135 - assert len(embeddings[0]) == COLBERT_DIM
136 -
137 -
138 -@pytest.mark.parametrize("model_name", [MODEL_NAME])
139 -def test_colbert_embed_not_supported(server: RemoteOpenAIServer, model_name: str):
140 - """Test that ColBERT model does not support 'embed' task."""
141 - task = "embed"
142 - text = "What is the capital of France?"
143 -
144 - response = requests.post(
145 - server.url_for("pooling"),
146 - json={
147 - "model": model_name,
148 - "input": text,
149 - "task": task,
150 - },
151 - )
152 -
153 - assert response.json()["error"]["type"] == "BadRequestError"
154 - assert response.json()["error"]["message"].startswith(f"Unsupported task: {task!r}")
27 +class TestColBERTOnline:
28 + def test_rerank(self, server: RemoteOpenAIServer):
29 + """Test ColBERT rerank endpoint."""
30 + query = "What is the capital of France?"
31 + documents = [
32 + "The capital of Brazil is Brasilia.",
33 + "The capital of France is Paris.",
34 + ]
35 +
36 + rerank_response = requests.post(
37 + server.url_for("rerank"),
38 + json={
39 + "model": MODEL_NAME,
40 + "query": query,
41 + "documents": documents,
42 + },
43 + )
44 + rerank_response.raise_for_status()
45 + rerank = RerankResponse.model_validate(rerank_response.json())
46 +
47 + assert rerank.id is not None
48 + assert rerank.results is not None
49 + assert len(rerank.results) == 2
50 +
51 + paris_result = next(r for r in rerank.results if r.index == 1)
52 + brazil_result = next(r for r in rerank.results if r.index == 0)
53 +
54 + assert paris_result.relevance_score > brazil_result.relevance_score
55 +
56 + def test_rerank_top_n(self, server: RemoteOpenAIServer):
57 + """Test ColBERT rerank with top_n parameter."""
58 + query = "What is the capital of France?"
59 + documents = [
60 + "The capital of Brazil is Brasilia.",
61 + "The capital of France is Paris.",
62 + "Machine learning is a field of AI.",
63 + ]
64 +
65 + rerank_response = requests.post(
66 + server.url_for("rerank"),
67 + json={
68 + "model": MODEL_NAME,
69 + "query": query,
70 + "documents": documents,
71 + "top_n": 2,
72 + },
73 + )
74 + rerank_response.raise_for_status()
75 + rerank = RerankResponse.model_validate(rerank_response.json())
76 +
77 + assert len(rerank.results) == 2
78 + assert rerank.results[0].index == 1
79 +
80 + def test_score(self, server: RemoteOpenAIServer):
81 + """Test ColBERT score endpoint."""
82 + text_1 = "What is the capital of France?"
83 + text_2 = ["The capital of France is Paris.", "Python is a language."]
84 +
85 + score_response = requests.post(
86 + server.url_for("score"),
87 + json={
88 + "model": MODEL_NAME,
89 + "text_1": text_1,
90 + "text_2": text_2,
91 + },
92 + )
93 + score_response.raise_for_status()
94 + score = ScoreResponse.model_validate(score_response.json())
95 +
96 + assert score.id is not None
97 + assert score.data is not None
98 + assert len(score.data) == 2
99 +
100 + assert score.data[0].score > score.data[1].score
101 +
102 + def test_token_embed(self, server: RemoteOpenAIServer):
103 + """Test ColBERT token_embed task via pooling endpoint."""
104 + text = "What is the capital of France?"
105 +
106 + pooling_response = requests.post(
107 + server.url_for("pooling"),
108 + json={
109 + "model": MODEL_NAME,
110 + "input": text,
111 + "task": "token_embed",
112 + },
113 + )
114 + pooling_response.raise_for_status()
115 + pooling = pooling_response.json()
116 +
117 + assert "data" in pooling
118 + assert len(pooling["data"]) == 1
119 +
120 + embeddings = pooling["data"][0]["data"]
121 + assert isinstance(embeddings, list)
122 + assert len(embeddings) > 0
123 + assert len(embeddings[0]) == COLBERT_DIM
124 +
125 + def test_embed_not_supported(self, server: RemoteOpenAIServer):
126 + """Test that ColBERT model does not support 'embed' task."""
127 + task = "embed"
128 + text = "What is the capital of France?"
129 +
130 + response = requests.post(
131 + server.url_for("pooling"),
132 + json={
133 + "model": MODEL_NAME,
134 + "input": text,
135 + "task": task,
136 + },
137 + )
138 +
139 + assert response.json()["error"]["type"] == "BadRequestError"
140 + assert response.json()["error"]["message"].startswith(
141 + f"Unsupported task: {task!r}"
142 + )
tests/models/language/pooling/test_colbert.py
+242 -73
@@ -1,16 +1,47 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 -"""Tests for ColBERT late interaction scoring."""
3 +"""Tests for ColBERT late interaction scoring.
4 +
5 +Tests are parametrized across multiple ColBERT backbones to ensure the
6 +generic ColBERT support works with different encoder architectures.
7 +"""
8
9 import pytest
10 import torch
11
12 from vllm.entrypoints.pooling.score.utils import compute_maxsim_score
13
10 -# ColBERT model - using answerai-colbert-small-v1 as it's a smaller model
11 -# suitable for testing (based on BERT-base)
12 -COLBERT_MODEL = "answerdotai/answerai-colbert-small-v1"
13 -COLBERT_DIM = 96 # This model uses 96-dimensional output
14 +# -----------------------------------------------------------------------
15 +# Model definitions: (model_name, colbert_dim, extra vllm_runner kwargs)
16 +# -----------------------------------------------------------------------
17 +COLBERT_MODELS = {
18 + "bert": {
19 + "model": "answerdotai/answerai-colbert-small-v1",
20 + "colbert_dim": 96,
21 + "max_model_len": 512,
22 + "extra_kwargs": {},
23 + },
24 + "modernbert": {
25 + "model": "lightonai/GTE-ModernColBERT-v1",
26 + "colbert_dim": 128,
27 + "max_model_len": 299,
28 + "extra_kwargs": {
29 + "hf_overrides": {
30 + "architectures": ["ColBERTModernBertModel"],
31 + },
32 + },
33 + },
34 + "jina": {
35 + "model": "jinaai/jina-colbert-v2",
36 + "colbert_dim": 128,
37 + "max_model_len": 8192,
38 + "extra_kwargs": {
39 + "hf_overrides": {
40 + "architectures": ["ColBERTJinaRobertaModel"],
41 + },
42 + },
43 + },
44 +}
45
46 TEXTS_1 = [
47 "What is the capital of France?",
@@ -25,80 +56,121 @@ TEXTS_2 = [
56 DTYPE = "half"
57
58
59 +# -----------------------------------------------------------------------
60 +# Fixtures
61 +# -----------------------------------------------------------------------
62 +
63 +
64 +@pytest.fixture(params=list(COLBERT_MODELS.keys()), scope="module")
65 +def colbert_spec(request):
66 + """Return the model spec dict for the current parametrization."""
67 + return COLBERT_MODELS[request.param]
68 +
69 +
70 @pytest.fixture(scope="module")
29 -def colbert_model_name():
30 - return COLBERT_MODEL
71 +def colbert_model_name(colbert_spec):
72 + return colbert_spec["model"]
73 +
74
75 +@pytest.fixture(scope="module")
76 +def colbert_dim(colbert_spec):
77 + return colbert_spec["colbert_dim"]
78 +
79 +
80 +@pytest.fixture(scope="module")
81 +def colbert_max_model_len(colbert_spec):
82 + return colbert_spec["max_model_len"]
83 +
84 +
85 +@pytest.fixture(scope="module")
86 +def colbert_extra_kwargs(colbert_spec):
87 + return colbert_spec["extra_kwargs"]
88
33 -def test_colbert_token_embed(vllm_runner, colbert_model_name):
89 +
90 +# -----------------------------------------------------------------------
91 +# Tests
92 +# -----------------------------------------------------------------------
93 +
94 +
95 +def test_colbert_token_embed(
96 + vllm_runner,
97 + colbert_model_name,
98 + colbert_dim,
99 + colbert_max_model_len,
100 + colbert_extra_kwargs,
101 +):
102 """Test that ColBERT model produces token embeddings."""
103 with vllm_runner(
104 colbert_model_name,
105 runner="pooling",
106 dtype=DTYPE,
39 - max_model_len=512,
107 + max_model_len=colbert_max_model_len,
108 enforce_eager=True,
109 + **colbert_extra_kwargs,
110 ) as vllm_model:
42 - # Get token embeddings for a single text
111 outputs = vllm_model.token_embed([TEXTS_1[0]])
112
113 assert len(outputs) == 1
46 - # Token embeddings should be 2D: [num_tokens, colbert_dim]
114 emb = torch.tensor(outputs[0])
115 assert emb.dim() == 2
49 - assert emb.shape[1] == COLBERT_DIM
50 - # Should have at least a few tokens
116 + assert emb.shape[1] == colbert_dim
117 assert emb.shape[0] > 1
118
119
54 -def test_colbert_late_interaction_1_to_1(vllm_runner, colbert_model_name):
120 +def test_colbert_late_interaction_1_to_1(
121 + vllm_runner,
122 + colbert_model_name,
123 + colbert_max_model_len,
124 + colbert_extra_kwargs,
125 +):
126 """Test ColBERT late interaction scoring with 1:1 query-document pair."""
127 with vllm_runner(
128 colbert_model_name,
129 runner="pooling",
130 dtype=DTYPE,
60 - max_model_len=512,
131 + max_model_len=colbert_max_model_len,
132 enforce_eager=True,
133 + **colbert_extra_kwargs,
134 ) as vllm_model:
63 - # Get token embeddings
135 q_outputs = vllm_model.token_embed([TEXTS_1[0]])
136 d_outputs = vllm_model.token_embed([TEXTS_2[0]])
137
138 q_emb = torch.tensor(q_outputs[0])
139 d_emb = torch.tensor(d_outputs[0])
140
70 - # Compute MaxSim manually
141 manual_score = compute_maxsim_score(q_emb, d_emb).item()
142
73 - # Use the score API (which should internally use _late_interaction_score)
143 vllm_scores = vllm_model.score(TEXTS_1[0], TEXTS_2[0])
144
145 assert len(vllm_scores) == 1
146 assert vllm_scores[0] == pytest.approx(manual_score, rel=0.01)
147
148
80 -def test_colbert_late_interaction_1_to_N(vllm_runner, colbert_model_name):
149 +def test_colbert_late_interaction_1_to_N(
150 + vllm_runner,
151 + colbert_model_name,
152 + colbert_max_model_len,
153 + colbert_extra_kwargs,
154 +):
155 """Test ColBERT late interaction scoring with 1:N query-documents."""
156 with vllm_runner(
157 colbert_model_name,
158 runner="pooling",
159 dtype=DTYPE,
86 - max_model_len=512,
160 + max_model_len=colbert_max_model_len,
161 enforce_eager=True,
162 + **colbert_extra_kwargs,
163 ) as vllm_model:
89 - # Get token embeddings
164 q_outputs = vllm_model.token_embed([TEXTS_1[0]])
165 d_outputs = vllm_model.token_embed(TEXTS_2)
166
167 q_emb = torch.tensor(q_outputs[0])
168
95 - # Compute MaxSim manually for each document
169 manual_scores = []
170 for d_out in d_outputs:
171 d_emb = torch.tensor(d_out)
172 manual_scores.append(compute_maxsim_score(q_emb, d_emb).item())
173
101 - # Use the score API
174 vllm_scores = vllm_model.score(TEXTS_1[0], TEXTS_2)
175
176 assert len(vllm_scores) == 2
@@ -106,27 +178,30 @@ def test_colbert_late_interaction_1_to_N(vllm_runner, colbert_model_name):
178 assert vllm_scores[i] == pytest.approx(manual_scores[i], rel=0.01)
179
180
109 -def test_colbert_late_interaction_N_to_N(vllm_runner, colbert_model_name):
181 +def test_colbert_late_interaction_N_to_N(
182 + vllm_runner,
183 + colbert_model_name,
184 + colbert_max_model_len,
185 + colbert_extra_kwargs,
186 +):
187 """Test ColBERT late interaction scoring with N:N query-documents."""
188 with vllm_runner(
189 colbert_model_name,
190 runner="pooling",
191 dtype=DTYPE,
115 - max_model_len=512,
192 + max_model_len=colbert_max_model_len,
193 enforce_eager=True,
194 + **colbert_extra_kwargs,
195 ) as vllm_model:
118 - # Get token embeddings
196 q_outputs = vllm_model.token_embed(TEXTS_1)
197 d_outputs = vllm_model.token_embed(TEXTS_2)
198
122 - # Compute MaxSim manually for each pair
199 manual_scores = []
200 for q_out, d_out in zip(q_outputs, d_outputs):
201 q_emb = torch.tensor(q_out)
202 d_emb = torch.tensor(d_out)
203 manual_scores.append(compute_maxsim_score(q_emb, d_emb).item())
204
129 - # Use the score API
205 vllm_scores = vllm_model.score(TEXTS_1, TEXTS_2)
206
207 assert len(vllm_scores) == 2
@@ -134,8 +209,13 @@ def test_colbert_late_interaction_N_to_N(vllm_runner, colbert_model_name):
209 assert vllm_scores[i] == pytest.approx(manual_scores[i], rel=0.01)
210
211
137 -def test_colbert_relevance_ordering(vllm_runner, colbert_model_name):
138 - """Test that ColBERT scores relevant documents higher than irrelevant ones."""
212 +def test_colbert_relevance_ordering(
213 + vllm_runner,
214 + colbert_model_name,
215 + colbert_max_model_len,
216 + colbert_extra_kwargs,
217 +):
218 + """Test that ColBERT scores relevant documents higher than irrelevant."""
219 query = "What is machine learning?"
220 documents = [
221 "Machine learning is a subset of artificial intelligence.",
@@ -147,48 +227,73 @@ def test_colbert_relevance_ordering(vllm_runner, colbert_model_name):
227 colbert_model_name,
228 runner="pooling",
229 dtype=DTYPE,
150 - max_model_len=512,
230 + max_model_len=colbert_max_model_len,
231 enforce_eager=True,
232 + **colbert_extra_kwargs,
233 ) as vllm_model:
234 scores = vllm_model.score(query, documents)
235
236 assert len(scores) == 3
156 - # ML-related documents should score higher than unrelated Python doc
157 - # Document 0 (ML definition) should be most relevant
158 - # Document 2 (Deep learning) should also be relevant
159 - # Document 1 (Python) should be least relevant
237 assert scores[0] > scores[1], "ML doc should score higher than Python doc"
238 assert scores[2] > scores[1], "DL doc should score higher than Python doc"
239
240
164 -def test_colbert_embed_not_supported(vllm_runner, colbert_model_name):
241 +def test_colbert_embed_not_supported(
242 + vllm_runner,
243 + colbert_model_name,
244 + colbert_max_model_len,
245 + colbert_extra_kwargs,
246 +):
247 """Test that ColBERT model does not support 'embed' task."""
248 with (
249 vllm_runner(
250 colbert_model_name,
251 runner="pooling",
252 dtype=DTYPE,
171 - max_model_len=512,
253 + max_model_len=colbert_max_model_len,
254 enforce_eager=True,
255 + **colbert_extra_kwargs,
256 ) as vllm_model,
257 pytest.raises(ValueError, match="Embedding API is not supported"),
258 ):
259 vllm_model.embed([TEXTS_1[0]])
260
261
179 -def test_colbert_hf_comparison(vllm_runner, colbert_model_name):
180 - """Test that vLLM ColBERT produces same embeddings as HuggingFace."""
262 +# -----------------------------------------------------------------------
263 +# Per-model HuggingFace comparison tests
264 +# -----------------------------------------------------------------------
265 +
266 +
267 +def _assert_embeddings_close(vllm_outputs, hf_embeddings):
268 + """Assert that vLLM and HuggingFace embeddings match."""
269 + for i, (hf_emb, vllm_out) in enumerate(zip(hf_embeddings, vllm_outputs)):
270 + vllm_emb = torch.tensor(vllm_out).float()
271 +
272 + assert hf_emb.shape == vllm_emb.shape, (
273 + f"Shape mismatch for text {i}: HF {hf_emb.shape} vs vLLM {vllm_emb.shape}"
274 + )
275 +
276 + torch.testing.assert_close(
277 + vllm_emb,
278 + hf_emb,
279 + rtol=1e-2,
280 + atol=1e-2,
281 + msg=f"Embedding mismatch for text {i}",
282 + )
283 +
284 +
285 +def test_colbert_hf_comparison_bert(vllm_runner):
286 + """Test that vLLM ColBERT produces same embeddings as HuggingFace (BERT)."""
287 import torch.nn.functional as F
288 from huggingface_hub import hf_hub_download
289 from safetensors.torch import load_file
290 from transformers import AutoTokenizer, BertModel
291
292 + model_name = COLBERT_MODELS["bert"]["model"]
293 test_texts = [TEXTS_1[0], TEXTS_2[0]]
294
188 - # Get vLLM embeddings first (to avoid GPU memory contention)
189 - # Use fp32 to match HuggingFace default precision for fair comparison
295 with vllm_runner(
191 - colbert_model_name,
296 + model_name,
297 runner="pooling",
298 dtype="float32",
299 max_model_len=512,
@@ -196,14 +301,11 @@ def test_colbert_hf_comparison(vllm_runner, colbert_model_name):
301 ) as vllm_model:
302 vllm_outputs = vllm_model.token_embed(test_texts)
303
199 - # Get HuggingFace reference embeddings on CPU
200 - # Load the base BERT model and manually apply the ColBERT linear projection
201 - hf_tokenizer = AutoTokenizer.from_pretrained(colbert_model_name)
202 - hf_bert = BertModel.from_pretrained(colbert_model_name)
304 + hf_tokenizer = AutoTokenizer.from_pretrained(model_name)
305 + hf_bert = BertModel.from_pretrained(model_name)
306 hf_bert.eval()
307
205 - # Load the ColBERT linear weights from safetensors
206 - weights_path = hf_hub_download(colbert_model_name, filename="model.safetensors")
308 + weights_path = hf_hub_download(model_name, filename="model.safetensors")
309 weights = load_file(weights_path)
310 linear_weight = weights["linear.weight"] # [96, 384]
311
@@ -212,36 +314,103 @@ def test_colbert_hf_comparison(vllm_runner, colbert_model_name):
314 inputs = hf_tokenizer(text, return_tensors="pt")
315 with torch.no_grad():
316 outputs = hf_bert(**inputs)
215 - # Get last hidden state: [1, seq_len, 384]
317 hidden_states = outputs.last_hidden_state
217 - # Apply ColBERT linear projection: [1, seq_len, 96]
318 token_emb = F.linear(hidden_states, linear_weight)
219 - # L2 normalize
319 token_emb = F.normalize(token_emb, p=2, dim=-1)
320 hf_embeddings.append(token_emb.squeeze(0).float())
321
223 - # Compare embeddings
224 - for i, (hf_emb, vllm_out) in enumerate(zip(hf_embeddings, vllm_outputs)):
225 - vllm_emb = torch.tensor(vllm_out).float()
322 + _assert_embeddings_close(vllm_outputs, hf_embeddings)
323
227 - # Print first few components for debugging
228 - print(f"\n=== Text {i}: '{test_texts[i][:30]}...' ===")
229 - print(f"HF shape: {hf_emb.shape}, vLLM shape: {vllm_emb.shape}")
230 - print(f"HF first token, first 10 dims: {hf_emb[0, :10].tolist()}")
231 - print(f"vLLM first token, first 10 dims: {vllm_emb[0, :10].tolist()}")
232 - print(f"HF last token, first 10 dims: {hf_emb[-1, :10].tolist()}")
233 - print(f"vLLM last token, first 10 dims: {vllm_emb[-1, :10].tolist()}")
324
235 - # Should have same shape
236 - assert hf_emb.shape == vllm_emb.shape, (
237 - f"Shape mismatch for text {i}: HF {hf_emb.shape} vs vLLM {vllm_emb.shape}"
238 - )
325 +def test_colbert_hf_comparison_modernbert(vllm_runner):
326 + """Test that vLLM ColBERT produces same embeddings as HuggingFace
327 + (ModernBERT)."""
328 + import torch.nn.functional as F
329 + from huggingface_hub import hf_hub_download
330 + from safetensors.torch import load_file
331 + from transformers import AutoModel, AutoTokenizer
332
240 - # Should have same values (with tolerance for fp16)
241 - torch.testing.assert_close(
242 - vllm_emb,
243 - hf_emb,
244 - rtol=1e-2,
245 - atol=1e-2,
246 - msg=f"Embedding mismatch for text {i}",
247 - )
333 + spec = COLBERT_MODELS["modernbert"]
334 + model_name = spec["model"]
335 + test_texts = [TEXTS_1[0], TEXTS_2[0]]
336 +
337 + with vllm_runner(
338 + model_name,
339 + runner="pooling",
340 + dtype="float32",
341 + max_model_len=spec["max_model_len"],
342 + enforce_eager=True,
343 + **spec["extra_kwargs"],
344 + ) as vllm_model:
345 + vllm_outputs = vllm_model.token_embed(test_texts)
346 +
347 + hf_tokenizer = AutoTokenizer.from_pretrained(model_name)
348 + hf_model = AutoModel.from_pretrained(model_name)
349 + hf_model.eval()
350 +
351 + # Load projection from sentence-transformers 1_Dense layer
352 + dense_path = hf_hub_download(model_name, filename="1_Dense/model.safetensors")
353 + dense_weights = load_file(dense_path)
354 + linear_weight = dense_weights["linear.weight"] # [128, 768]
355 +
356 + hf_embeddings = []
357 + for text in test_texts:
358 + inputs = hf_tokenizer(text, return_tensors="pt")
359 + with torch.no_grad():
360 + outputs = hf_model(**inputs)
361 + hidden_states = outputs.last_hidden_state
362 + token_emb = F.linear(hidden_states, linear_weight)
363 + token_emb = F.normalize(token_emb, p=2, dim=-1)
364 + hf_embeddings.append(token_emb.squeeze(0).float())
365 +
366 + _assert_embeddings_close(vllm_outputs, hf_embeddings)
367 +
368 +
369 +def test_colbert_hf_comparison_jina(vllm_runner):
370 + """Test that vLLM ColBERT produces same embeddings as HuggingFace
371 + (Jina XLM-RoBERTa)."""
372 + import torch.nn.functional as F
373 + from huggingface_hub import hf_hub_download
374 + from safetensors.torch import load_file
375 + from transformers import AutoModel, AutoTokenizer
376 +
377 + spec = COLBERT_MODELS["jina"]
378 + model_name = spec["model"]
379 + test_texts = [TEXTS_1[0], TEXTS_2[0]]
380 +
381 + with vllm_runner(
382 + model_name,
383 + runner="pooling",
384 + dtype="float32",
385 + max_model_len=spec["max_model_len"],
386 + enforce_eager=True,
387 + **spec["extra_kwargs"],
388 + ) as vllm_model:
389 + vllm_outputs = vllm_model.token_embed(test_texts)
390 +
391 + hf_tokenizer = AutoTokenizer.from_pretrained(
392 + model_name,
393 + trust_remote_code=True,
394 + )
395 + hf_model = AutoModel.from_pretrained(
396 + model_name,
397 + trust_remote_code=True,
398 + )
399 + hf_model.eval()
400 +
401 + # Load projection from main checkpoint
402 + weights_path = hf_hub_download(model_name, filename="model.safetensors")
403 + weights = load_file(weights_path)
404 + linear_weight = weights["linear.weight"] # [128, 1024]
405 +
406 + hf_embeddings = []
407 + for text in test_texts:
408 + inputs = hf_tokenizer(text, return_tensors="pt")
409 + with torch.no_grad():
410 + outputs = hf_model(**inputs)
411 + hidden_states = outputs.last_hidden_state
412 + token_emb = F.linear(hidden_states.float(), linear_weight.float())
413 + token_emb = F.normalize(token_emb, p=2, dim=-1)
414 + hf_embeddings.append(token_emb.squeeze(0).float())
415 +
416 + _assert_embeddings_close(vllm_outputs, hf_embeddings)
tests/models/registry.py
+9
@@ -529,6 +529,15 @@ _EMBEDDING_EXAMPLE_MODELS = {
529 # [Text-only]
530 "BertModel": _HfExamplesInfo("BAAI/bge-base-en-v1.5"),
531 "HF_ColBERT": _HfExamplesInfo("answerdotai/answerai-colbert-small-v1"),
532 + "ColBERTModernBertModel": _HfExamplesInfo(
533 + "lightonai/GTE-ModernColBERT-v1",
534 + hf_overrides={"architectures": ["ColBERTModernBertModel"]},
535 + ),
536 + "ColBERTJinaRobertaModel": _HfExamplesInfo(
537 + "jinaai/jina-colbert-v2",
538 + trust_remote_code=True,
539 + hf_overrides={"architectures": ["ColBERTJinaRobertaModel"]},
540 + ),
541 "BgeM3EmbeddingModel": _HfExamplesInfo("BAAI/bge-m3"),
542 "Gemma2Model": _HfExamplesInfo("BAAI/bge-multilingual-gemma2"),
543 "Gemma3TextModel": _HfExamplesInfo("google/embeddinggemma-300m"),
vllm/model_executor/models/colbert.py
+335 -68
@@ -6,6 +6,14 @@ ColBERT late interaction model for retrieval and reranking.
6 ColBERT uses per-token embeddings and late interaction (MaxSim) scoring
7 instead of single-vector representations or cross-encoder concatenation.
8
9 +This module provides:
10 +
11 +- :class:`ColBERTMixin` — mixin that adds ColBERT late-interaction support
12 + to any embedding model.
13 +- :class:`ColBERTModel` — ColBERT with BERT backbone (original architecture).
14 +- :class:`ColBERTModernBertModel` — ColBERT with ModernBERT backbone.
15 +- :class:`ColBERTJinaRobertaModel` — ColBERT with Jina XLM-RoBERTa backbone.
16 +
17 Reference: https://arxiv.org/abs/2004.12832
18 """
19
@@ -23,51 +31,60 @@ from .bert import BertEmbeddingModel, BertModel
31 from .interfaces_base import default_pooling_type
32
33
26 -@default_pooling_type(seq_pooling_type="CLS", tok_pooling_type="ALL")
27 -class ColBERTModel(BertEmbeddingModel):
28 - """ColBERT late interaction model for retrieval/reranking.
34 +class ColBERTMixin:
35 + """Mixin that adds ColBERT late interaction support to any embedding model.
36
30 - This model extends BertEmbeddingModel with a ColBERT-style linear
31 - projection layer for per-token embeddings. It supports only:
32 - - "token_embed" task: Per-token embeddings for late interaction
37 + ColBERT (Contextualized Late Interaction over BERT) uses per-token
38 + embeddings with a linear projection layer. This mixin provides:
39
34 - ColBERT is fundamentally a per-token embedding model - the linear
35 - projection is trained for per-token representations, not for CLS
36 - pooling. Use a dedicated dense embedding model if you need single-
37 - vector representations.
40 + - ``supports_late_interaction`` class-var
41 + - ColBERT linear projection initialisation / lazy creation
42 + - Weight loading helpers for the projection layer
43 + - A builder for the token-embedding pooler
44
39 - The ColBERT scoring (MaxSim) is computed externally, either client-side
40 - or via the late interaction scoring path in ServingScores.
45 + **Integration:**
46
42 - Attributes:
43 - colbert_linear: Linear projection from hidden_size to colbert_dim
44 - supports_late_interaction: Flag indicating this model uses late
45 - interaction scoring
47 + 1. Inherit from both ``ColBERTMixin`` and ``nn.Module``.
48 + 2. In ``__init__``: call ``super().__init__()``, then
49 + :meth:`_init_colbert_components`, then create ``self.model``
50 + (the backbone) and ``self.pooler`` via :meth:`_build_colbert_pooler`.
51 + 3. In ``load_weights``: use :meth:`_load_colbert_weights` to separate
52 + the ColBERT projection weight, then delegate the rest to the backbone.
53 """
54
48 - # Mark this model as supporting late interaction scoring
55 supports_late_interaction: ClassVar[Literal[True]] = True
56
51 - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
52 - # Get config before calling super().__init__
53 - config = vllm_config.model_config.hf_config
54 - self.hidden_size = config.hidden_size
55 - self.head_dtype = vllm_config.model_config.head_dtype
56 -
57 - # ColBERT dimension - check various config field names used by different
58 - # ColBERT implementations. If not found in config, will be inferred
59 - # from loaded weights in load_weights()
60 - self.colbert_dim: int | None = (
61 - getattr(config, "colbert_dim", None)
62 - or getattr(config, "dim", None)
63 - or getattr(config, "projection_dim", None)
64 - )
57 + # Set during _init_colbert_components
58 + colbert_dim: int | None
59 + colbert_linear: nn.Linear | None
60 + hidden_size: int
61 + head_dtype: torch.dtype
62
66 - # Initialize parent (this will call _build_pooler)
67 - super().__init__(vllm_config=vllm_config, prefix=prefix)
63 + # ------------------------------------------------------------------ init
64
69 - def _build_model(self, vllm_config: VllmConfig, prefix: str = "") -> BertModel:
70 - return BertModel(vllm_config=vllm_config, prefix=prefix)
65 + def _init_colbert_components(
66 + self,
67 + hidden_size: int,
68 + colbert_dim: int | None,
69 + head_dtype: torch.dtype,
70 + ) -> None:
71 + """Initialise ColBERT projection layer.
72 +
73 + Args:
74 + hidden_size: Hidden dimension of the encoder backbone.
75 + colbert_dim: Output dimension for ColBERT embeddings. If
76 + ``None``, will be inferred from weights during loading (or
77 + auto-loaded from sentence-transformers Dense layers).
78 + head_dtype: Data type for the projection layer.
79 + """
80 + self.hidden_size = hidden_size
81 + self.colbert_dim = colbert_dim
82 + self.head_dtype = head_dtype
83 +
84 + if colbert_dim is not None:
85 + self.colbert_linear = self._build_colbert_linear()
86 + else:
87 + self.colbert_linear = None
88
89 def _build_colbert_linear(self) -> nn.Linear:
90 """Build the ColBERT linear projection layer."""
@@ -80,24 +97,127 @@ class ColBERTModel(BertEmbeddingModel):
97 dtype=self.head_dtype,
98 )
99
83 - def _build_pooler(self, pooler_config: PoolerConfig) -> Pooler:
84 - # ColBERT linear projection: hidden_size -> colbert_dim
85 - # Original ColBERT uses bias=False
86 - # If colbert_dim is not set from config, it will be inferred during
87 - # load_weights and the linear layer will be created there
88 - if self.colbert_dim is not None:
89 - self.colbert_linear = self._build_colbert_linear()
90 - else:
91 - # Placeholder - will be created when weights are loaded
92 - self.colbert_linear = None
100 + # ---------------------------------------------------------------- pooler
101
94 - # ColBERT only supports token_embed - it's fundamentally a per-token
95 - # embedding model.
102 + def _build_colbert_pooler(self, pooler_config: PoolerConfig) -> Pooler:
103 + """Build pooler for ColBERT token embeddings.
104 +
105 + When ``colbert_linear`` is set, it is used as the projector.
106 + Otherwise ``pooler_for_token_embed`` falls back to auto-loading
107 + sentence-transformers Dense layers (``1_Dense/`` etc.).
108 + """
109 return pooler_for_token_embed(
110 pooler_config,
111 projector=self.colbert_linear,
112 )
113
114 + # --------------------------------------------------------- config helper
115 +
116 + @classmethod
117 + def get_colbert_dim_from_config(cls, hf_config) -> int | None:
118 + """Extract ColBERT dimension from a HuggingFace config.
119 +
120 + Checks ``colbert_dim``, ``dim`` and ``projection_dim`` in that order.
121 + """
122 + return (
123 + getattr(hf_config, "colbert_dim", None)
124 + or getattr(hf_config, "dim", None)
125 + or getattr(hf_config, "projection_dim", None)
126 + )
127 +
128 + # -------------------------------------------------------- weight loading
129 +
130 + def _load_colbert_weights(
131 + self,
132 + weights: Iterable[tuple[str, torch.Tensor]],
133 + colbert_weight_names: tuple[str, ...] = (
134 + "linear.weight",
135 + "colbert_linear.weight",
136 + ),
137 + ) -> tuple[list[tuple[str, torch.Tensor]], set[str]]:
138 + """Separate and load ColBERT projection weights.
139 +
140 + Scans *weights* for entries whose name ends with one of
141 + *colbert_weight_names*. The matching weight is loaded into
142 + ``self.colbert_linear`` (creating it first if ``colbert_dim`` was
143 + not known at init time).
144 +
145 + Args:
146 + weights: Iterable of ``(name, tensor)`` weight pairs.
147 + colbert_weight_names: Suffixes that identify the ColBERT linear
148 + weight.
149 +
150 + Returns:
151 + ``(remaining_weights, loaded_names)`` — the weights that were
152 + **not** consumed and the set of names that were loaded.
153 + """
154 + weights_list = list(weights)
155 + other_weights: list[tuple[str, torch.Tensor]] = []
156 + colbert_weight: tuple[str, torch.Tensor] | None = None
157 +
158 + for name, weight in weights_list:
159 + if any(name.endswith(cw) for cw in colbert_weight_names):
160 + colbert_weight = (name, weight)
161 + else:
162 + other_weights.append((name, weight))
163 +
164 + loaded: set[str] = set()
165 + if colbert_weight is not None:
166 + _name, weight = colbert_weight
167 + if weight.dim() == 2:
168 + # Infer colbert_dim from weight shape if not set
169 + if self.colbert_dim is None:
170 + self.colbert_dim = weight.shape[0]
171 + self.colbert_linear = self._build_colbert_linear()
172 + # Update the pooler's projector
173 + if hasattr(self, "pooler") and hasattr(self.pooler, "head"):
174 + self.pooler.head.projector = self.colbert_linear
175 +
176 + assert self.colbert_linear is not None
177 + # Move to same device as model
178 + if hasattr(self, "model"):
179 + device = next(self.model.parameters()).device
180 + self.colbert_linear.to(device)
181 +
182 + weight = weight.to(self.colbert_linear.weight.device)
183 + self.colbert_linear.weight.data.copy_(weight)
184 + loaded.add("pooler.head.projector.weight")
185 +
186 + return other_weights, loaded
187 +
188 +
189 +# -----------------------------------------------------------------------
190 +# Concrete model: ColBERT + BERT backbone (original architecture)
191 +# -----------------------------------------------------------------------
192 +
193 +
194 +@default_pooling_type(seq_pooling_type="CLS", tok_pooling_type="ALL")
195 +class ColBERTModel(ColBERTMixin, BertEmbeddingModel):
196 + """ColBERT late interaction model with BERT backbone.
197 +
198 + Supports the ``token_embed`` task (per-token embeddings for late
199 + interaction). MaxSim scoring is computed externally.
200 + """
201 +
202 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
203 + config = vllm_config.model_config.hf_config
204 +
205 + # Must run before super().__init__ because _build_pooler reads these.
206 + colbert_dim = self.get_colbert_dim_from_config(config)
207 + self._init_colbert_components(
208 + hidden_size=config.hidden_size,
209 + colbert_dim=colbert_dim,
210 + head_dtype=vllm_config.model_config.head_dtype,
211 + )
212 +
213 + super().__init__(vllm_config=vllm_config, prefix=prefix)
214 +
215 + def _build_model(self, vllm_config: VllmConfig, prefix: str = "") -> BertModel:
216 + return BertModel(vllm_config=vllm_config, prefix=prefix)
217 +
218 + def _build_pooler(self, pooler_config: PoolerConfig) -> Pooler:
219 + return self._build_colbert_pooler(pooler_config)
220 +
221 def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
222 def _strip(name: str) -> str:
223 for p in ("model.", "bert."):
@@ -111,7 +231,7 @@ class ColBERTModel(BertEmbeddingModel):
231
232 for name, weight in weights_list:
233 stripped = _strip(name)
114 - # Handle different checkpoint naming conventions for ColBERT linear
234 + # Handle different checkpoint naming conventions
235 if stripped in ("linear.weight", "colbert_linear.weight"):
236 colbert_side.append(("colbert_linear.weight", weight))
237 elif stripped.startswith("linear.") or stripped.startswith(
@@ -122,31 +242,178 @@ class ColBERTModel(BertEmbeddingModel):
242 else:
243 model_side.append((stripped, weight))
244
125 - # Load base BERT weights using BertModel.load_weights which handles QKV fusion
245 loaded: set[str] = set()
246 loaded_model = self.model.load_weights(model_side)
247 loaded.update({"model." + n for n in loaded_model})
248
130 - # Load ColBERT linear weights
249 if colbert_side:
132 - for name, weight in colbert_side:
133 - if name == "colbert_linear.weight":
134 - # Infer colbert_dim from weights if not set in config
135 - if self.colbert_dim is None:
136 - # Weight shape is [colbert_dim, hidden_size]
137 - self.colbert_dim = weight.shape[0]
138 - # Create the linear layer now that we know the dimension
139 - self.colbert_linear = self._build_colbert_linear()
140 - # Move to the same device as the model's existing parameters
141 - device = next(self.model.parameters()).device
142 - self.colbert_linear.to(device)
143 - # Update the pooler's projector to use the new linear layer
144 - self.pooler.head.projector = self.colbert_linear
250 + _, colbert_loaded = self._load_colbert_weights(colbert_side)
251 + loaded.update(colbert_loaded)
252 +
253 + return loaded
254 +
255 +
256 +# -----------------------------------------------------------------------
257 +# Concrete model: ColBERT + ModernBERT backbone
258 +# -----------------------------------------------------------------------
259
146 - # Load weights directly into the pooler's projector
147 - weight = weight.to(self.pooler.head.projector.weight.device)
148 - self.pooler.head.projector.weight.data.copy_(weight)
149 - loaded.add("pooler.head.projector.weight")
150 - break
260 +from .modernbert import ModernBertModel # noqa: E402
261 +
262 +
263 +@default_pooling_type(seq_pooling_type="CLS", tok_pooling_type="ALL")
264 +class ColBERTModernBertModel(ColBERTMixin, nn.Module):
265 + """ColBERT late interaction model with ModernBERT backbone.
266 +
267 + For ``lightonai/GTE-ModernColBERT-v1`` and similar models.
268 + The projection is auto-loaded from sentence-transformers ``1_Dense/``
269 + when not present in the main checkpoint.
270 + """
271 +
272 + is_pooling_model = True
273 +
274 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
275 + super().__init__()
276 + config = vllm_config.model_config.hf_config
277 +
278 + colbert_dim = self.get_colbert_dim_from_config(config)
279 + self._init_colbert_components(
280 + hidden_size=config.hidden_size,
281 + colbert_dim=colbert_dim,
282 + head_dtype=vllm_config.model_config.head_dtype,
283 + )
284 +
285 + self.model = ModernBertModel(
286 + vllm_config=vllm_config,
287 + prefix=prefix,
288 + )
289 +
290 + pooler_config = vllm_config.model_config.pooler_config
291 + assert pooler_config is not None
292 + self.pooler = self._build_colbert_pooler(pooler_config)
293 +
294 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
295 + return self.model.embed_input_ids(input_ids)
296 +
297 + def forward(
298 + self,
299 + input_ids: torch.Tensor,
300 + positions: torch.Tensor,
301 + intermediate_tensors=None,
302 + inputs_embeds: torch.Tensor | None = None,
303 + ) -> torch.Tensor:
304 + return self.model(
305 + input_ids=input_ids,
306 + positions=positions,
307 + inputs_embeds=inputs_embeds,
308 + intermediate_tensors=intermediate_tensors,
309 + )
310 +
311 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
312 + other_weights, colbert_loaded = self._load_colbert_weights(weights)
313 +
314 + # Strip "model." prefix added by the embedding adapter
315 + model_weights = [
316 + (n[len("model.") :] if n.startswith("model.") else n, w)
317 + for n, w in other_weights
318 + ]
319 +
320 + loaded_model = self.model.load_weights(model_weights)
321 + loaded = {"model." + n for n in loaded_model} | colbert_loaded
322 +
323 + # When the ST projector was auto-loaded during init
324 + # (not from the main checkpoint), mark its params as loaded
325 + # so the weight validator doesn't complain.
326 + if hasattr(self.pooler, "head"):
327 + head = self.pooler.head
328 + projector = getattr(head, "projector", None)
329 + if projector is not None and isinstance(projector, nn.Module):
330 + for name, _ in projector.named_parameters():
331 + loaded.add(f"pooler.head.projector.{name}")
332 +
333 + return loaded
334 +
335 +
336 +# -----------------------------------------------------------------------
337 +# Concrete model: ColBERT + Jina XLM-RoBERTa backbone
338 +# -----------------------------------------------------------------------
339 +
340 +from .bert_with_rope import JinaRobertaModel # noqa: E402
341 +
342 +
343 +@default_pooling_type(seq_pooling_type="CLS", tok_pooling_type="ALL")
344 +class ColBERTJinaRobertaModel(ColBERTMixin, nn.Module):
345 + """ColBERT late interaction model with Jina XLM-RoBERTa backbone.
346 +
347 + For ``jinaai/jina-colbert-v2`` and similar models.
348 + """
349 +
350 + is_pooling_model = True
351 +
352 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
353 + super().__init__()
354 + config = vllm_config.model_config.hf_config
355 +
356 + colbert_dim = self.get_colbert_dim_from_config(config)
357 + self._init_colbert_components(
358 + hidden_size=config.hidden_size,
359 + colbert_dim=colbert_dim,
360 + head_dtype=vllm_config.model_config.head_dtype,
361 + )
362 +
363 + self.model = JinaRobertaModel(
364 + vllm_config=vllm_config,
365 + prefix=prefix,
366 + )
367 +
368 + pooler_config = vllm_config.model_config.pooler_config
369 + assert pooler_config is not None
370 + self.pooler = self._build_colbert_pooler(pooler_config)
371 +
372 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
373 + return self.model.embed_input_ids(input_ids)
374 +
375 + def forward(
376 + self,
377 + input_ids: torch.Tensor,
378 + positions: torch.Tensor,
379 + intermediate_tensors=None,
380 + inputs_embeds: torch.Tensor | None = None,
381 + ) -> torch.Tensor:
382 + return self.model(
383 + input_ids=input_ids,
384 + positions=positions,
385 + inputs_embeds=inputs_embeds,
386 + intermediate_tensors=intermediate_tensors,
387 + )
388 +
389 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
390 + weights_list = list(weights)
391 + model_side: list[tuple[str, torch.Tensor]] = []
392 + colbert_side: list[tuple[str, torch.Tensor]] = []
393 +
394 + for name, weight in weights_list:
395 + stripped = name
396 + # Strip "model." prefix added by the embedding adapter
397 + if stripped.startswith("model."):
398 + stripped = stripped[len("model.") :]
399 + # Strip "roberta." prefix from checkpoint
400 + if stripped.startswith("roberta."):
401 + stripped = stripped[len("roberta.") :]
402 +
403 + if stripped in ("linear.weight", "colbert_linear.weight"):
404 + colbert_side.append(("colbert_linear.weight", weight))
405 + elif stripped.startswith("pooler."):
406 + # Skip HF pooler weights (not used in ColBERT)
407 + continue
408 + else:
409 + model_side.append((stripped, weight))
410 +
411 + loaded: set[str] = set()
412 + loaded_model = self.model.load_weights(model_side)
413 + loaded.update({"model." + n for n in loaded_model})
414 +
415 + if colbert_side:
416 + _, colbert_loaded = self._load_colbert_weights(colbert_side)
417 + loaded.update(colbert_loaded)
418
419 return loaded
vllm/model_executor/models/config.py
+1
@@ -629,6 +629,7 @@ MODELS_CONFIG_MAP: dict[str, type[VerifyAndUpdateConfig]] = {
629 "Qwen3ForSequenceClassification": Qwen3ForSequenceClassificationConfig,
630 "Qwen3VLForSequenceClassification": Qwen3VLForSequenceClassificationConfig,
631 "XLMRobertaModel": JinaRobertaModelConfig,
632 + "ColBERTJinaRobertaModel": JinaRobertaModelConfig,
633 "JinaVLForRanking": JinaVLForSequenceClassificationConfig,
634 "JambaForSequenceClassification": JambaForSequenceClassificationConfig,
635 "GptOssForCausalLM": GptOssForCausalLMConfig,
vllm/model_executor/models/registry.py
+2
@@ -208,6 +208,8 @@ _EMBEDDING_MODELS = {
208 "BertModel": ("bert", "BertEmbeddingModel"),
209 "BertSpladeSparseEmbeddingModel": ("bert", "BertSpladeSparseEmbeddingModel"),
210 "HF_ColBERT": ("colbert", "ColBERTModel"),
211 + "ColBERTModernBertModel": ("colbert", "ColBERTModernBertModel"),
212 + "ColBERTJinaRobertaModel": ("colbert", "ColBERTJinaRobertaModel"),
213 "DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"),
214 "Gemma2Model": ("gemma2", "Gemma2ForCausalLM"),
215 "Gemma3TextModel": ("gemma3", "Gemma3Model"),
vllm/transformers_utils/config.py
+5 -3
@@ -1068,9 +1068,11 @@ def try_get_dense_modules(
1068 if isinstance(modules, dict):
1069 modules = modules.get("modules", [])
1070
1071 - dense_modules = [
1072 - m for m in modules if m.get("type") == "sentence_transformers.models.Dense"
1073 - ]
1071 + _DENSE_MODULE_TYPES = {
1072 + "sentence_transformers.models.Dense",
1073 + "pylate.models.Dense.Dense",
1074 + }
1075 + dense_modules = [m for m in modules if m.get("type") in _DENSE_MODULE_TYPES]
1076 if not dense_modules:
1077 return None
1078