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 5, 09:45 UTC
[Model] Update ColModernVBERT to support latest HF checkpoint (#39307)
Signed-off-by: Ilya Boytsov <ilyaboytsov1805@gmail.com>
Ilya Boytsov committed
Apr 9, 2026 at 04:48 UTC
d37b378762e44c964ad6993c6e9612102388db4c
5 files changed
+29
-88
tests/models/multimodal/pooling/test_colmodernvbert.py
-8
@@ -15,10 +15,6 @@ from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score
15
MODEL_NAME = "ModernVBERT/colmodernvbert-merged"
16
COLBERT_DIM = 128
17
DTYPE = "half"
18
-# Fixme:
19
-# Update colmodernvbert code to support the latest HF version
20
-# and remove revision set.
21
-REVISION = "4a0a9f3ac7a7992fec410bfa8e3d080ac9a5bcee"
18
19
20
# -----------------------------------------------------------------------
@@ -30,7 +26,6 @@ def test_colmodernvbert_text_token_embed(vllm_runner):
26
"""Text query produces per-token embeddings with shape (seq_len, 128)."""
27
with vllm_runner(
28
MODEL_NAME,
33
- revision=REVISION,
29
runner="pooling",
30
dtype=DTYPE,
31
enforce_eager=True,
@@ -54,7 +49,6 @@ def test_colmodernvbert_text_relevance_ordering(vllm_runner):
49
50
with vllm_runner(
51
MODEL_NAME,
57
- revision=REVISION,
52
runner="pooling",
53
dtype=DTYPE,
54
enforce_eager=True,
@@ -72,7 +66,6 @@ def test_colmodernvbert_text_late_interaction(vllm_runner):
66
67
with vllm_runner(
68
MODEL_NAME,
75
- revision=REVISION,
69
runner="pooling",
70
dtype=DTYPE,
71
enforce_eager=True,
@@ -99,7 +92,6 @@ def test_colmodernvbert_image_token_embed(vllm_runner, image_assets):
92
"""Image input produces per-token embeddings including vision tokens."""
93
with vllm_runner(
94
MODEL_NAME,
102
- revision=REVISION,
95
runner="pooling",
96
dtype=DTYPE,
97
enforce_eager=True,
tests/models/registry.py
-1
@@ -648,7 +648,6 @@ _LATE_INTERACTION_EXAMPLE_MODELS = {
648
# [Multimodal]
649
"ColModernVBertForRetrieval": _HfExamplesInfo(
650
"ModernVBERT/colmodernvbert-merged",
651
- revision="4a0a9f3ac7a7992fec410bfa8e3d080ac9a5bcee",
651
),
652
"ColPaliForRetrieval": _HfExamplesInfo("vidore/colpali-v1.3-hf"),
653
"ColQwen3": _HfExamplesInfo(
vllm/model_executor/models/colmodernvbert.py
+3
-51
@@ -18,7 +18,6 @@ from vllm.config import VllmConfig
18
from vllm.config.multimodal import BaseDummyOptions
19
from vllm.inputs import MultiModalDataDict
20
from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_embed
21
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
21
from vllm.multimodal import MULTIMODAL_REGISTRY
22
from vllm.multimodal.inputs import (
23
MultiModalFieldConfig,
@@ -358,70 +357,23 @@ class ColModernVBertForRetrieval(
357
"model.text_model.layers.": "text_layers.",
358
"model.text_model.embeddings.": "text_embeddings.",
359
"model.text_model.final_norm.": "text_final_norm.",
361
- "model.connector.modality_projection.": "connector.",
360
+ "model.connector.modality_projection.": "connector.proj.",
361
"model.custom_text_proj.": "custom_text_proj.",
363
- "model.vision_model.": "vision_model.vision_model.",
362
+ "model.vision_model.vision_model.": "vision_model.vision_model.",
363
"model.": "",
364
},
365
)
366
368
- # Checkpoint names for DecoupledEmbedding parts
369
- _BASE_EMB = "model.text_model.embeddings.tok_embeddings.weight"
370
- _EXTRA_EMB = (
371
- "model.text_model.embeddings.tok_embeddings.additional_embedding.weight"
372
- )
373
-
367
def load_weights(
368
self,
369
weights: Iterable[tuple[str, torch.Tensor]],
370
) -> set[str]:
378
- # DecoupledEmbedding requires concatenating base + additional
379
- # embedding tensors before loading, so we extract them first.
380
- base_embedding_weight: torch.Tensor | None = None
381
- additional_embedding_weight: torch.Tensor | None = None
382
- remaining: list[tuple[str, torch.Tensor]] = []
383
-
384
- for name, tensor in weights:
385
- if name == self._BASE_EMB:
386
- base_embedding_weight = tensor
387
- elif name == self._EXTRA_EMB:
388
- additional_embedding_weight = tensor
389
- else:
390
- remaining.append((name, tensor))
391
-
392
- # Load all non-embedding weights via AutoWeightsLoader
371
loader = AutoWeightsLoader(self)
372
loaded_params = loader.load_weights(
395
- remaining,
373
+ weights,
374
mapper=self.hf_to_vllm_mapper,
375
)
376
399
- # Concatenate and load DecoupledEmbedding weights
400
- if base_embedding_weight is not None:
401
- combined = base_embedding_weight
402
- if additional_embedding_weight is not None:
403
- combined = torch.cat(
404
- [base_embedding_weight, additional_embedding_weight],
405
- dim=0,
406
- )
407
- param_name = "text_embeddings.tok_embeddings.weight"
408
- params_dict = dict(self.named_parameters())
409
- if param_name in params_dict:
410
- param = params_dict[param_name]
411
- weight_loader = getattr(
412
- param,
413
- "weight_loader",
414
- default_weight_loader,
415
- )
416
- weight_loader(param, combined)
417
- loaded_params.add(param_name)
418
- elif additional_embedding_weight is not None:
419
- raise ValueError(
420
- "Found 'text_model.embeddings.tok_embeddings"
421
- ".additional_embedding.weight' but not "
422
- "'text_model.embeddings.tok_embeddings.weight'"
423
- )
424
-
377
# The pooler wraps ``custom_text_proj`` as its head projector.
378
# Mark those params as loaded under the pooler path too.
379
if hasattr(self, "pooler") and hasattr(self.pooler, "head"):
vllm/transformers_utils/config.py
+1
-1
@@ -82,7 +82,7 @@ _CONFIG_REGISTRY: dict[str, type[PretrainedConfig]] = LazyConfigDict(
82
bagel="BagelConfig",
83
umm="CheersConfig",
84
chatglm="ChatGLMConfig",
85
- colmodernvbert="ColModernVBertConfig",
85
+ modernvbert="ColModernVBertConfig",
86
colpali="ColPaliConfig",
87
colqwen3="ColQwen3Config",
88
ops_colqwen3="OpsColQwen3Config",
vllm/transformers_utils/configs/colmodernvbert.py
+25
-27
@@ -17,43 +17,41 @@ class ColModernVBertConfig(PretrainedConfig):
17
def __init__(
18
self,
19
embedding_dim: int = 128,
20
- vlm_config: dict | None = None,
20
+ image_token_id: int = 50407,
21
+ pixel_shuffle_factor: int = 4,
22
+ text_config: dict | None = None,
23
+ vision_config: dict | None = None,
24
**kwargs,
25
):
26
self.embedding_dim = embedding_dim
27
+ self.image_token_id = image_token_id
28
+ self.pixel_shuffle_factor = pixel_shuffle_factor
29
25
- if vlm_config is None:
26
- vlm_config = {}
30
+ text_config = text_config or {}
31
+ self.hidden_size = text_config.get("hidden_size", 768)
32
28
- # Top-level VLM fields
29
- self.image_token_id = vlm_config.get("image_token_id", 50407)
30
- self.pixel_shuffle_factor = vlm_config.get("pixel_shuffle_factor", 4)
31
- self.hidden_size = vlm_config.get("hidden_size", 768)
32
- additional_vocab_size = vlm_config.get("additional_vocab_size", 40)
33
-
34
- # Text config (ModernBERT)
35
- text_cfg = vlm_config.get("text_config", {})
36
- base_vocab = text_cfg.get("vocab_size", 50368)
33
self.text_config = ModernBertConfig(
38
- vocab_size=base_vocab + additional_vocab_size,
39
- hidden_size=text_cfg.get("hidden_size", 768),
40
- intermediate_size=text_cfg.get("intermediate_size", 1152),
41
- num_hidden_layers=text_cfg.get("num_hidden_layers", 22),
42
- num_attention_heads=text_cfg.get("num_attention_heads", 12),
43
- mlp_bias=text_cfg.get("mlp_bias", False),
44
- max_position_embeddings=vlm_config.get("max_position_embeddings", 8192),
34
+ vocab_size=text_config.get("vocab_size", 50408),
35
+ hidden_size=text_config.get("hidden_size", 768),
36
+ intermediate_size=text_config.get("intermediate_size", 1152),
37
+ num_hidden_layers=text_config.get("num_hidden_layers", 22),
38
+ num_attention_heads=text_config.get("num_attention_heads", 12),
39
+ mlp_bias=text_config.get("mlp_bias", False),
40
+ max_position_embeddings=text_config.get("max_position_embeddings", 8192),
41
)
42
47
- # Vision config (SigLIP)
48
- vis_cfg = vlm_config.get("vision_config", {})
43
+ vision_config = vision_config or {}
44
self.vision_config = SiglipVisionConfig(
50
- hidden_size=vis_cfg.get("embed_dim", 768),
51
- image_size=vis_cfg.get("image_size", 512),
52
- patch_size=vis_cfg.get("patch_size", 16),
53
- num_hidden_layers=vis_cfg.get("num_hidden_layers", 12),
54
- intermediate_size=vis_cfg.get("intermediate_size", 3072),
55
- num_attention_heads=vis_cfg.get("num_attention_heads", 12),
45
+ hidden_size=vision_config.get("hidden_size", 768),
46
+ image_size=vision_config.get("image_size", 512),
47
+ patch_size=vision_config.get("patch_size", 16),
48
+ num_hidden_layers=vision_config.get("num_hidden_layers", 12),
49
+ intermediate_size=vision_config.get("intermediate_size", 3072),
50
+ num_attention_heads=vision_config.get("num_attention_heads", 12),
51
)
52
+
53
+ # Ensure architectures is set so vLLM routes to our model class
54
+ kwargs.setdefault("architectures", ["ColModernVBertForRetrieval"])
55
super().__init__(**kwargs)
56
57
@property