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, 00:00 UTC
[CI] Fix ColBERT HF comparison tests on AMD CI + refactor (#34567)
Signed-off-by: Andreas Karatzas <akaratza@amd.com>
Andreas Karatzas committed
Feb 20, 2026 at 22:12 UTC
89358f0d35e7923cf1554d4d652094c4ad2e80de
1 file changed
+107
-149
tests/models/language/pooling/test_colbert.py
+107
-149
@@ -20,6 +20,12 @@ COLBERT_MODELS = {
20
"colbert_dim": 96,
21
"max_model_len": 512,
22
"extra_kwargs": {},
23
+ "hf_comparison": {
24
+ "weights_file": "model.safetensors",
25
+ "weights_key": "linear.weight",
26
+ "trust_remote_code": False,
27
+ "model_cls": "BertModel",
28
+ },
29
},
30
"modernbert": {
31
"model": "lightonai/GTE-ModernColBERT-v1",
@@ -30,6 +36,12 @@ COLBERT_MODELS = {
36
"architectures": ["ColBERTModernBertModel"],
37
},
38
},
39
+ "hf_comparison": {
40
+ "weights_file": "1_Dense/model.safetensors",
41
+ "weights_key": "linear.weight",
42
+ "trust_remote_code": False,
43
+ "model_cls": "AutoModel",
44
+ },
45
},
46
"jina": {
47
"model": "jinaai/jina-colbert-v2",
@@ -40,9 +52,16 @@ COLBERT_MODELS = {
52
"architectures": ["ColBERTJinaRobertaModel"],
53
},
54
},
55
+ "hf_comparison": {
56
+ "weights_file": "model.safetensors",
57
+ "weights_key": "linear.weight",
58
+ "trust_remote_code": True,
59
+ "model_cls": "AutoModel",
60
+ },
61
},
62
}
63
64
+
65
TEXTS_1 = [
66
"What is the capital of France?",
67
"What is the capital of Germany?",
@@ -56,9 +75,68 @@ TEXTS_2 = [
75
DTYPE = "half"
76
77
59
-# -----------------------------------------------------------------------
60
-# Fixtures
61
-# -----------------------------------------------------------------------
78
+def _load_hf_model(model_name: str, hf_spec: dict, device: torch.device):
79
+ """Load HF model on the given device with a compatible attention impl."""
80
+ from transformers import AutoModel, BertModel
81
+
82
+ cls = BertModel if hf_spec["model_cls"] == "BertModel" else AutoModel
83
+ trust = hf_spec.get("trust_remote_code", False)
84
+
85
+ # Flash / Triton kernels require GPU tensors; fall back to eager on CPU.
86
+ extra = {}
87
+ if device.type == "cpu":
88
+ extra["attn_implementation"] = "eager"
89
+
90
+ model = cls.from_pretrained(
91
+ model_name,
92
+ trust_remote_code=trust,
93
+ **extra,
94
+ ).to(device)
95
+ model.eval()
96
+ return model
97
+
98
+
99
+def _load_projection_weight(model_name: str, hf_spec: dict, device: torch.device):
100
+ """Download and return the ColBERT linear projection weight."""
101
+ from huggingface_hub import hf_hub_download
102
+ from safetensors.torch import load_file
103
+
104
+ path = hf_hub_download(model_name, filename=hf_spec["weights_file"])
105
+ weights = load_file(path)
106
+ return weights[hf_spec["weights_key"]].to(device)
107
+
108
+
109
+def _compute_hf_colbert_embeddings(model, tokenizer, linear_weight, texts, device):
110
+ """Run HF model + projection and return L2-normalised token embeddings."""
111
+ import torch.nn.functional as F
112
+
113
+ embeddings = []
114
+ for text in texts:
115
+ inputs = tokenizer(text, return_tensors="pt").to(device)
116
+ with torch.no_grad():
117
+ hidden = model(**inputs).last_hidden_state.float()
118
+ projected = F.linear(hidden, linear_weight.float())
119
+ normalised = F.normalize(projected, p=2, dim=-1)
120
+ embeddings.append(normalised.squeeze(0).cpu())
121
+ return embeddings
122
+
123
+
124
+def _assert_embeddings_close(vllm_outputs, hf_embeddings):
125
+ """Assert that vLLM and HuggingFace embeddings match."""
126
+ for i, (hf_emb, vllm_out) in enumerate(zip(hf_embeddings, vllm_outputs)):
127
+ vllm_emb = torch.as_tensor(vllm_out).float()
128
+
129
+ assert hf_emb.shape == vllm_emb.shape, (
130
+ f"Shape mismatch for text {i}: HF {hf_emb.shape} vs vLLM {vllm_emb.shape}"
131
+ )
132
+
133
+ torch.testing.assert_close(
134
+ vllm_emb,
135
+ hf_emb,
136
+ rtol=1e-2,
137
+ atol=1e-2,
138
+ msg=f"Embedding mismatch for text {i}",
139
+ )
140
141
142
@pytest.fixture(params=list(COLBERT_MODELS.keys()), scope="module")
@@ -87,11 +165,6 @@ def colbert_extra_kwargs(colbert_spec):
165
return colbert_spec["extra_kwargs"]
166
167
90
-# -----------------------------------------------------------------------
91
-# Tests
92
-# -----------------------------------------------------------------------
93
-
94
-
168
def test_colbert_token_embed(
169
vllm_runner,
170
colbert_model_name,
@@ -111,7 +184,7 @@ def test_colbert_token_embed(
184
outputs = vllm_model.token_embed([TEXTS_1[0]])
185
186
assert len(outputs) == 1
114
- emb = torch.tensor(outputs[0])
187
+ emb = torch.as_tensor(outputs[0])
188
assert emb.dim() == 2
189
assert emb.shape[1] == colbert_dim
190
assert emb.shape[0] > 1
@@ -135,8 +208,8 @@ def test_colbert_late_interaction_1_to_1(
208
q_outputs = vllm_model.token_embed([TEXTS_1[0]])
209
d_outputs = vllm_model.token_embed([TEXTS_2[0]])
210
138
- q_emb = torch.tensor(q_outputs[0])
139
- d_emb = torch.tensor(d_outputs[0])
211
+ q_emb = torch.as_tensor(q_outputs[0])
212
+ d_emb = torch.as_tensor(d_outputs[0])
213
214
manual_score = compute_maxsim_score(q_emb, d_emb).item()
215
@@ -164,11 +237,11 @@ def test_colbert_late_interaction_1_to_N(
237
q_outputs = vllm_model.token_embed([TEXTS_1[0]])
238
d_outputs = vllm_model.token_embed(TEXTS_2)
239
167
- q_emb = torch.tensor(q_outputs[0])
240
+ q_emb = torch.as_tensor(q_outputs[0])
241
242
manual_scores = []
243
for d_out in d_outputs:
171
- d_emb = torch.tensor(d_out)
244
+ d_emb = torch.as_tensor(d_out)
245
manual_scores.append(compute_maxsim_score(q_emb, d_emb).item())
246
247
vllm_scores = vllm_model.score(TEXTS_1[0], TEXTS_2)
@@ -198,8 +271,8 @@ def test_colbert_late_interaction_N_to_N(
271
272
manual_scores = []
273
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)
274
+ q_emb = torch.as_tensor(q_out)
275
+ d_emb = torch.as_tensor(d_out)
276
manual_scores.append(compute_maxsim_score(q_emb, d_emb).item())
277
278
vllm_scores = vllm_model.score(TEXTS_1, TEXTS_2)
@@ -259,79 +332,16 @@ def test_colbert_embed_not_supported(
332
vllm_model.embed([TEXTS_1[0]])
333
334
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
- )
335
+@pytest.mark.parametrize("backend", list(COLBERT_MODELS.keys()))
336
+def test_colbert_hf_comparison(vllm_runner, backend):
337
+ """Test that vLLM ColBERT embeddings match HuggingFace for each backend."""
338
+ from transformers import AutoTokenizer
339
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
-
295
- with vllm_runner(
296
- model_name,
297
- runner="pooling",
298
- dtype="float32",
299
- max_model_len=512,
300
- enforce_eager=True,
301
- ) as vllm_model:
302
- vllm_outputs = vllm_model.token_embed(test_texts)
303
-
304
- hf_tokenizer = AutoTokenizer.from_pretrained(model_name)
305
- hf_bert = BertModel.from_pretrained(model_name)
306
- hf_bert.eval()
307
-
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
-
312
- hf_embeddings = []
313
- for text in test_texts:
314
- inputs = hf_tokenizer(text, return_tensors="pt")
315
- with torch.no_grad():
316
- outputs = hf_bert(**inputs)
317
- hidden_states = outputs.last_hidden_state
318
- token_emb = F.linear(hidden_states, linear_weight)
319
- token_emb = F.normalize(token_emb, p=2, dim=-1)
320
- hf_embeddings.append(token_emb.squeeze(0).float())
321
-
322
- _assert_embeddings_close(vllm_outputs, hf_embeddings)
323
-
324
-
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
-
333
- spec = COLBERT_MODELS["modernbert"]
340
+ spec = COLBERT_MODELS[backend]
341
+ hf_spec = spec["hf_comparison"]
342
model_name = spec["model"]
343
+ assert isinstance(model_name, str)
344
+ assert isinstance(hf_spec, dict)
345
test_texts = [TEXTS_1[0], TEXTS_2[0]]
346
347
with vllm_runner(
@@ -344,73 +354,21 @@ def test_colbert_hf_comparison_modernbert(vllm_runner):
354
) as vllm_model:
355
vllm_outputs = vllm_model.token_embed(test_texts)
356
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)
357
+ device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
358
359
hf_tokenizer = AutoTokenizer.from_pretrained(
360
model_name,
393
- trust_remote_code=True,
361
+ trust_remote_code=hf_spec.get("trust_remote_code", False),
362
)
395
- hf_model = AutoModel.from_pretrained(
396
- model_name,
397
- trust_remote_code=True,
363
+ hf_model = _load_hf_model(model_name, hf_spec, device)
364
+ linear_weight = _load_projection_weight(model_name, hf_spec, device)
365
+
366
+ hf_embeddings = _compute_hf_colbert_embeddings(
367
+ hf_model,
368
+ hf_tokenizer,
369
+ linear_weight,
370
+ test_texts,
371
+ device,
372
)
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())
373
374
_assert_embeddings_close(vllm_outputs, hf_embeddings)