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 6, 08:15 UTC
[Model] Add LFM2-ColBERT-350M support (#37528)
Signed-off-by: Ilya Boytsov <ilyaboytsov1805@gmail.com>
Ilya Boytsov committed
Mar 20, 2026 at 15:57 UTC
8b6c6b950579778e222b24f501cd81a0c1d719d8
6 files changed
+125
-1
docs/models/pooling_models/specific_models.md
+5
@@ -11,6 +11,7 @@ vLLM supports ColBERT models with multiple encoder backbones:
11
| `HF_ColBERT` | BERT | `answerdotai/answerai-colbert-small-v1`, `colbert-ir/colbertv2.0` |
12
| `ColBERTModernBertModel` | ModernBERT | `lightonai/GTE-ModernColBERT-v1` |
13
| `ColBERTJinaRobertaModel` | Jina XLM-RoBERTa | `jinaai/jina-colbert-v2` |
14
+| `ColBERTLfm2Model` | LFM2 | `LiquidAI/LFM2-ColBERT-350M` |
15
16
**BERT-based ColBERT** models work out of the box:
17
@@ -29,6 +30,10 @@ vllm serve lightonai/GTE-ModernColBERT-v1 \
30
vllm serve jinaai/jina-colbert-v2 \
31
--hf-overrides '{"architectures": ["ColBERTJinaRobertaModel"]}' \
32
--trust-remote-code
33
+
34
+# LFM2 backbone
35
+vllm serve LiquidAI/LFM2-ColBERT-350M \
36
+ --hf-overrides '{"architectures": ["ColBERTLfm2Model"]}'
37
```
38
39
Then you can use the rerank API:
docs/models/pooling_models/token_embed.md
+1
@@ -39,6 +39,7 @@ Models of any architecture can be converted into embedding models using `--conve
39
40
| Architecture | Models | Example HF Models | [LoRA](../../features/lora.md) | [PP](../../serving/parallelism_scaling.md) |
41
| ------------ | ------ | ----------------- | -------------------- | ------------------------- |
42
+| `ColBERTLfm2Model` | LFM2 | `LiquidAI/LFM2-ColBERT-350M` | | |
43
| `ColBERTModernBertModel` | ModernBERT | `lightonai/GTE-ModernColBERT-v1` | | |
44
| `ColBERTJinaRobertaModel` | Jina XLM-RoBERTa | `jinaai/jina-colbert-v2` | | |
45
| `HF_ColBERT` | BERT | `answerdotai/answerai-colbert-small-v1`, `colbert-ir/colbertv2.0` | | |
tests/models/language/pooling/test_colbert.py
+16
@@ -59,6 +59,22 @@ COLBERT_MODELS = {
59
"model_cls": "AutoModel",
60
},
61
},
62
+ "lfm2": {
63
+ "model": "LiquidAI/LFM2-ColBERT-350M",
64
+ "colbert_dim": 128,
65
+ "max_model_len": 511,
66
+ "extra_kwargs": {
67
+ "hf_overrides": {
68
+ "architectures": ["ColBERTLfm2Model"],
69
+ },
70
+ },
71
+ "hf_comparison": {
72
+ "weights_file": "1_Dense/model.safetensors",
73
+ "weights_key": "linear.weight",
74
+ "trust_remote_code": False,
75
+ "model_cls": "AutoModel",
76
+ },
77
+ },
78
}
79
80
tests/models/registry.py
+5
@@ -628,6 +628,11 @@ _LATE_INTERACTION_EXAMPLE_MODELS = {
628
trust_remote_code=True,
629
hf_overrides={"architectures": ["ColBERTJinaRobertaModel"]},
630
),
631
+ "ColBERTLfm2Model": _HfExamplesInfo(
632
+ "LiquidAI/LFM2-ColBERT-350M",
633
+ trust_remote_code=True,
634
+ hf_overrides={"architectures": ["ColBERTLfm2Model"]},
635
+ ),
636
# [Multimodal]
637
"ColModernVBertForRetrieval": _HfExamplesInfo(
638
"ModernVBERT/colmodernvbert-merged",
vllm/model_executor/models/colbert.py
+97
-1
@@ -27,8 +27,9 @@ from vllm.model_executor.layers.pooler import Pooler
27
from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_embed
28
29
from .bert import BertEmbeddingModel, BertModel
30
-from .interfaces import SupportsLateInteraction
30
+from .interfaces import HasInnerState, IsHybrid, SupportsLateInteraction
31
from .interfaces_base import default_pooling_type
32
+from .lfm2 import Lfm2ForCausalLM, Lfm2Model
33
34
35
class ColBERTMixin(nn.Module, SupportsLateInteraction):
@@ -414,3 +415,98 @@ class ColBERTJinaRobertaModel(ColBERTMixin, nn.Module):
415
loaded.update(colbert_loaded)
416
417
return loaded
418
+
419
+
420
+# -----------------------------------------------------------------------
421
+# Concrete model: ColBERT + LFM2 backbone
422
+# -----------------------------------------------------------------------
423
+
424
+
425
+@default_pooling_type(seq_pooling_type="CLS", tok_pooling_type="ALL")
426
+class ColBERTLfm2Model(ColBERTMixin, nn.Module, HasInnerState, IsHybrid):
427
+ """ColBERT late interaction model with LFM2 backbone.
428
+
429
+ For ``LiquidAI/LFM2-ColBERT-350M`` and similar models.
430
+
431
+ The projection is auto-loaded from sentence-transformers ``1_Dense/``
432
+ when not present in the main checkpoint.
433
+ """
434
+
435
+ is_pooling_model = True
436
+ # LFM2 is a hybrid model (attention + SSM layers); these flags ensure
437
+ # HybridAttentionMambaModelConfig.verify_and_update_config runs so that
438
+ # mamba_block_size and related cache settings are correctly initialised.
439
+ is_hybrid = True
440
+ has_inner_state = True
441
+
442
+ @classmethod
443
+ def get_mamba_state_shape_from_config(cls, vllm_config: VllmConfig):
444
+ return Lfm2ForCausalLM.get_mamba_state_shape_from_config(vllm_config)
445
+
446
+ @classmethod
447
+ def get_mamba_state_dtype_from_config(cls, vllm_config: VllmConfig):
448
+ return Lfm2ForCausalLM.get_mamba_state_dtype_from_config(vllm_config)
449
+
450
+ @classmethod
451
+ def get_mamba_state_copy_func(cls):
452
+ return Lfm2ForCausalLM.get_mamba_state_copy_func()
453
+
454
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
455
+ super().__init__()
456
+ config = vllm_config.model_config.hf_config
457
+
458
+ colbert_dim = self.get_colbert_dim_from_config(config)
459
+ self._init_colbert_components(
460
+ hidden_size=config.hidden_size,
461
+ colbert_dim=colbert_dim,
462
+ head_dtype=vllm_config.model_config.head_dtype,
463
+ )
464
+
465
+ self.model = Lfm2Model(
466
+ vllm_config=vllm_config,
467
+ prefix=prefix,
468
+ )
469
+
470
+ pooler_config = vllm_config.model_config.pooler_config
471
+ assert pooler_config is not None
472
+ self.pooler = self._build_colbert_pooler(pooler_config)
473
+
474
+ def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
475
+ return self.model.embed_input_ids(input_ids)
476
+
477
+ def forward(
478
+ self,
479
+ input_ids: torch.Tensor,
480
+ positions: torch.Tensor,
481
+ intermediate_tensors=None,
482
+ inputs_embeds: torch.Tensor | None = None,
483
+ ) -> torch.Tensor:
484
+ return self.model(
485
+ input_ids=input_ids,
486
+ positions=positions,
487
+ inputs_embeds=inputs_embeds,
488
+ intermediate_tensors=intermediate_tensors,
489
+ )
490
+
491
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
492
+ other_weights, colbert_loaded = self._load_colbert_weights(weights)
493
+
494
+ # Strip "model." prefix added by the embedding adapter
495
+ model_weights = [
496
+ (n[len("model.") :] if n.startswith("model.") else n, w)
497
+ for n, w in other_weights
498
+ ]
499
+ loaded_model = self.model.load_weights(model_weights)
500
+ loaded = {f"model.{name}" for name in loaded_model} | colbert_loaded
501
+
502
+ # When the ST projector was auto-loaded during init
503
+ # (not from the main checkpoint), mark its params as loaded
504
+ # so the weight validator doesn't complain.
505
+ if hasattr(self.pooler, "head"):
506
+ head = self.pooler.head
507
+ projector = getattr(head, "projector", None)
508
+ if projector is not None and isinstance(projector, nn.Module):
509
+ for name, _ in projector.named_parameters():
510
+ loaded.add(f"pooler.head.projector.{name}")
511
+
512
+ return loaded
vllm/model_executor/models/registry.py
+1
@@ -269,6 +269,7 @@ _LATE_INTERACTION_MODELS = {
269
"HF_ColBERT": ("colbert", "ColBERTModel"),
270
"ColBERTModernBertModel": ("colbert", "ColBERTModernBertModel"),
271
"ColBERTJinaRobertaModel": ("colbert", "ColBERTJinaRobertaModel"),
272
+ "ColBERTLfm2Model": ("colbert", "ColBERTLfm2Model"),
273
# [Multimodal]
274
"ColModernVBertForRetrieval": ("colmodernvbert", "ColModernVBertForRetrieval"),
275
"ColPaliForRetrieval": ("colpali", "ColPaliModel"),