A high-throughput and memory-efficient inference and serving engine for LLMs
Model/colbert autoweightsloader (#44999)
Signed-off-by: Furkan Fidan <dev@yufufi.com> Co-authored-by: wang.yuqi <yuqi.wang@daocloud.io>
Furkan F committed
Jun 10, 2026 at 09:58 UTC
7fdfa6441db1395434f0b95eb375969c119476c5
1 file changed
+27
-79
vllm/model_executor/models/colbert.py
+27
-79
@@ -25,6 +25,7 @@ from torch import nn
25
from vllm.config import PoolerConfig, VllmConfig
26
from vllm.model_executor.layers.pooler import Pooler
27
from vllm.model_executor.layers.pooler.tokwise import pooler_for_token_embed
28
+from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper
29
30
from .bert import BertEmbeddingModel, BertModel
31
from .interfaces import HasInnerState, IsHybrid, SupportsLateInteraction
@@ -217,38 +218,12 @@ class ColBERTModel(ColBERTMixin, BertEmbeddingModel):
218
return self._build_colbert_pooler(pooler_config)
219
220
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
220
- def _strip(name: str) -> str:
221
- for p in ("model.", "bert."):
222
- if name.startswith(p):
223
- name = name[len(p) :]
224
- return name
225
-
226
- weights_list = list(weights)
227
- model_side: list[tuple[str, torch.Tensor]] = []
228
- colbert_side: list[tuple[str, torch.Tensor]] = []
229
-
230
- for name, weight in weights_list:
231
- stripped = _strip(name)
232
- # Handle different checkpoint naming conventions
233
- if stripped in ("linear.weight", "colbert_linear.weight"):
234
- colbert_side.append(("colbert_linear.weight", weight))
235
- elif stripped.startswith("linear.") or stripped.startswith(
236
- "colbert_linear."
237
- ):
238
- new_name = stripped.replace("linear.", "colbert_linear.")
239
- colbert_side.append((new_name, weight))
240
- else:
241
- model_side.append((stripped, weight))
242
-
243
- loaded: set[str] = set()
244
- loaded_model = self.model.load_weights(model_side)
245
- loaded.update({"model." + n for n in loaded_model})
246
-
247
- if colbert_side:
248
- _, colbert_loaded = self._load_colbert_weights(colbert_side)
249
- loaded.update(colbert_loaded)
250
-
251
- return loaded
221
+ other_weights, colbert_loaded = self._load_colbert_weights(weights)
222
+ # Force "bert." to become "model."
223
+ mapper = WeightsMapper(orig_to_new_prefix={"bert.": "model."})
224
+ loader = AutoWeightsLoader(self)
225
+ loaded = loader.load_weights(other_weights, mapper=mapper)
226
+ return loaded | colbert_loaded
227
228
229
# -----------------------------------------------------------------------
@@ -309,18 +284,14 @@ class ColBERTModernBertModel(ColBERTMixin, nn.Module):
284
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
285
other_weights, colbert_loaded = self._load_colbert_weights(weights)
286
312
- # Strip "model." prefix added by the embedding adapter
313
- model_weights = [
314
- (n[len("model.") :] if n.startswith("model.") else n, w)
315
- for n, w in other_weights
316
- ]
317
-
318
- loaded_model = self.model.load_weights(model_weights)
319
- loaded = {"model." + n for n in loaded_model} | colbert_loaded
287
+ loaded_model = self.model.load_weights(other_weights)
288
+ loaded = {f"model.{name}" for name in loaded_model} | colbert_loaded
289
321
- # When the ST projector was auto-loaded during init
322
- # (not from the main checkpoint), mark its params as loaded
323
- # so the weight validator doesn't complain.
290
+ # When the ST projector is loaded via `_build_colbert_pooler`, the weights
291
+ # might come from `colbert_loaded` or the pooler automatically falls back to
292
+ # load from `/1_Dense` etc.
293
+ # We need to mark its params as loaded so the weight validator doesn't complain
294
+ # when they are loaded via fallback.
295
if hasattr(self.pooler, "head"):
296
head = self.pooler.head
297
projector = getattr(head, "projector", None)
@@ -385,36 +356,15 @@ class ColBERTJinaRobertaModel(ColBERTMixin, nn.Module):
356
)
357
358
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
388
- weights_list = list(weights)
389
- model_side: list[tuple[str, torch.Tensor]] = []
390
- colbert_side: list[tuple[str, torch.Tensor]] = []
391
-
392
- for name, weight in weights_list:
393
- stripped = name
394
- # Strip "model." prefix added by the embedding adapter
395
- if stripped.startswith("model."):
396
- stripped = stripped[len("model.") :]
397
- # Strip "roberta." prefix from checkpoint
398
- if stripped.startswith("roberta."):
399
- stripped = stripped[len("roberta.") :]
400
-
401
- if stripped in ("linear.weight", "colbert_linear.weight"):
402
- colbert_side.append(("colbert_linear.weight", weight))
403
- elif stripped.startswith("pooler."):
404
- # Skip HF pooler weights (not used in ColBERT)
405
- continue
406
- else:
407
- model_side.append((stripped, weight))
359
+ other_weights, colbert_loaded = self._load_colbert_weights(weights)
360
409
- loaded: set[str] = set()
410
- loaded_model = self.model.load_weights(model_side)
411
- loaded.update({"model." + n for n in loaded_model})
361
+ mapper = WeightsMapper(orig_to_new_prefix={"roberta.": "model."})
362
413
- if colbert_side:
414
- _, colbert_loaded = self._load_colbert_weights(colbert_side)
415
- loaded.update(colbert_loaded)
363
+ # Skip HF pooler weights (model.pooler.*) as they not used in ColBERT
364
+ loader = AutoWeightsLoader(self, skip_prefixes=["model.pooler."])
365
417
- return loaded
366
+ loaded = loader.load_weights(other_weights, mapper=mapper)
367
+ return loaded | colbert_loaded
368
369
370
# -----------------------------------------------------------------------
@@ -491,17 +441,15 @@ class ColBERTLfm2Model(ColBERTMixin, nn.Module, HasInnerState, IsHybrid):
441
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
442
other_weights, colbert_loaded = self._load_colbert_weights(weights)
443
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)
444
+ loaded_model = self.model.load_weights(other_weights)
445
+
446
loaded = {f"model.{name}" for name in loaded_model} | colbert_loaded
447
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.
448
+ # When the ST projector is loaded via `_build_colbert_pooler`, the weights
449
+ # might come from `colbert_loaded` or the pooler automatically falls back to
450
+ # load from `/1_Dense` etc.
451
+ # We need to mark its params as loaded so the weight validator doesn't complain
452
+ # when they are loaded via fallback.
453
if hasattr(self.pooler, "head"):
454
head = self.pooler.head
455
projector = getattr(head, "projector", None)