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 15, 02:15 UTC
[Model] Add AutoWeightsLoader support for BERT, RoBERTa (#20534)
Signed-off-by: Jennifer He <islandhe@gmail.com> Signed-off-by: <islandhe@gmail.com> Signed-off-by: Jen H <islandhe@gmail.com>
Jennifer He committed
Jul 15, 2025 at 01:34 UTC
85bd6599e4574f405c6b7944380c16358d5949bc
2 files changed
+59
-100
vllm/model_executor/models/bert.py
+37
-48
@@ -22,12 +22,11 @@ from vllm.model_executor.layers.pooler import (ClassifierPooler, Pooler,
22
from vllm.model_executor.layers.quantization import QuantizationConfig
23
from vllm.model_executor.layers.vocab_parallel_embedding import (
24
VocabParallelEmbedding)
25
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
25
from vllm.model_executor.pooling_metadata import PoolingMetadata
26
from vllm.sequence import IntermediateTensors, PoolerOutput
27
28
from .interfaces import SupportsCrossEncoding, SupportsQuant, SupportsV0Only
30
-from .utils import WeightsMapper, maybe_prefix
29
+from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
30
31
32
class BertEmbedding(nn.Module):
@@ -44,9 +43,11 @@ class BertEmbedding(nn.Module):
43
config.type_vocab_size, config.hidden_size)
44
self.LayerNorm = nn.LayerNorm(config.hidden_size,
45
eps=config.layer_norm_eps)
47
- self.position_ids = nn.Parameter(
48
- torch.empty((1, config.max_position_embeddings)), )
46
47
+ self.register_buffer(
48
+ "position_ids",
49
+ torch.arange(config.max_position_embeddings).unsqueeze(0),
50
+ )
51
self.position_embedding_type = config.position_embedding_type
52
if self.position_embedding_type != "absolute":
53
raise ValueError("Only 'absolute' position_embedding_type" +
@@ -358,45 +359,45 @@ class BertModel(nn.Module, SupportsQuant):
359
("qkv_proj", "value", "v"),
360
]
361
362
+ loaded_stacked_params = []
363
+ other_weights = []
364
params_dict = dict(self.named_parameters())
362
- loaded_params: set[str] = set()
365
for name, loaded_weight in weights:
364
- if self.pooler is None and "pooler" in name:
365
- continue
366
for (param_name, weight_name, shard_id) in stacked_params_mapping:
367
if weight_name not in name:
368
continue
369
+
370
name = name.replace(weight_name, param_name)
370
- # Skip loading extra bias for GPTQ models.
371
- if name.endswith(".bias") and name not in params_dict:
371
+ if name not in params_dict:
372
continue
373
param = params_dict[name]
374
weight_loader = param.weight_loader
375
weight_loader(param, loaded_weight, shard_id)
376
+ loaded_stacked_params.append(name)
377
break
378
else:
378
- # Skip loading extra bias for GPTQ models.
379
- if name.endswith(".bias") and name not in params_dict:
380
- continue
381
- param = params_dict[name]
382
- weight_loader = getattr(param, "weight_loader",
383
- default_weight_loader)
384
- weight_loader(param, loaded_weight)
385
- loaded_params.add(name)
379
+ if name in params_dict:
380
+ other_weights.append((name, loaded_weight))
381
+
382
+ loader = AutoWeightsLoader(
383
+ self,
384
+ skip_prefixes=(["pooler."] if self.pooler is None else []),
385
+ )
386
+ loaded_params = loader.load_weights(other_weights)
387
+ loaded_params.update(loaded_stacked_params)
388
return loaded_params
389
390
391
class BertEmbeddingModel(nn.Module, SupportsV0Only, SupportsQuant):
392
"""A model that uses Bert to provide embedding functionalities.
393
392
- This class encapsulates the BertModel and provides an interface for
393
- embedding operations and customized pooling functions.
394
+ This class encapsulates the BertModel and provides an interface for
395
+ embedding operations and customized pooling functions.
396
395
- Attributes:
396
- model: An instance of BertModel used for forward operations.
397
- _pooler: An instance of Pooler used for pooling operations.
398
- """
399
- hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"model.": ""})
397
+ Attributes:
398
+ model: An instance of BertModel used for forward operations.
399
+ _pooler: An instance of Pooler used for pooling operations.
400
+ """
401
402
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
403
super().__init__()
@@ -425,10 +426,15 @@ class BertEmbeddingModel(nn.Module, SupportsV0Only, SupportsQuant):
426
return self._pooler(hidden_states, pooling_metadata)
427
428
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
428
- weights = self.hf_to_vllm_mapper.apply(weights)
429
- weights = ((name, data) for name, data in weights
430
- if not name.startswith("lm_head."))
431
- self.model.load_weights(weights)
429
+ weights_list = list(weights)
430
+
431
+ has_model_prefix = any(
432
+ name.startswith("model.") for name, _ in weights_list)
433
+ if not has_model_prefix:
434
+ mapper = WeightsMapper(orig_to_new_prefix={"": "model."})
435
+
436
+ loader = AutoWeightsLoader(self, skip_prefixes=["lm_head."])
437
+ return loader.load_weights(weights_list, mapper=mapper)
438
439
def _build_model(self,
440
vllm_config: VllmConfig,
@@ -470,26 +476,9 @@ class BertForSequenceClassification(nn.Module, SupportsV0Only,
476
self.classifier, self.bert.pooler)
477
478
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
473
-
474
- self_weights = []
475
-
476
- def weight_filter():
477
- for name, weight in weights:
478
- if name.startswith("bert."):
479
- yield (name[len("bert."):], weight)
480
- else:
481
- self_weights.append((name, weight))
482
-
483
- self.bert.load_weights(weight_filter())
484
-
485
- params_dict = dict(self.named_parameters())
486
-
487
- for name, loaded_weight in self_weights:
488
- if name.startswith("classifier"):
489
- param = params_dict[name]
490
- weight_loader = getattr(param, "weight_loader",
491
- default_weight_loader)
492
- weight_loader(param, loaded_weight)
479
+ loader = AutoWeightsLoader(self)
480
+ loaded_params = loader.load_weights(weights)
481
+ return loaded_params
482
483
def pooler(
484
self,
vllm/model_executor/models/roberta.py
+22
-52
@@ -1,7 +1,6 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
-import itertools
4
from collections.abc import Iterable
5
from typing import Optional, Union
6
@@ -13,9 +12,9 @@ from vllm.config import VllmConfig
12
from vllm.model_executor.layers.pooler import ClassifierPooler
13
from vllm.model_executor.layers.vocab_parallel_embedding import (
14
VocabParallelEmbedding)
16
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
15
from vllm.model_executor.models.bert import BertEmbeddingModel, BertModel
18
-from vllm.model_executor.models.utils import WeightsMapper, maybe_prefix
16
+from vllm.model_executor.models.utils import (AutoWeightsLoader, WeightsMapper,
17
+ maybe_prefix)
18
from vllm.model_executor.pooling_metadata import PoolingMetadata
19
from vllm.sequence import IntermediateTensors, PoolerOutput
20
@@ -39,8 +38,10 @@ class RobertaEmbedding(nn.Module):
38
config.hidden_size)
39
self.LayerNorm = nn.LayerNorm(config.hidden_size,
40
eps=config.layer_norm_eps)
42
- self.position_ids = nn.Parameter(
43
- torch.empty((1, config.max_position_embeddings)), )
41
+ self.register_buffer(
42
+ "position_ids",
43
+ torch.arange(config.max_position_embeddings).unsqueeze(0),
44
+ )
45
46
self.position_embedding_type = config.position_embedding_type
47
if self.position_embedding_type != "absolute":
@@ -136,16 +137,20 @@ class RobertaEmbeddingModel(BertEmbeddingModel):
137
embedding_class=RobertaEmbedding)
138
139
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
139
- weights = self.hf_to_vllm_mapper.apply(weights)
140
- # Separate weights in "roberta"-prefixed and all else (not in memory).
141
- # For use with models like FacebookAI/roberta-base.
142
- bert_weights, task_weights = roberta_task_weights_filter(weights)
143
- loaded = self.model.load_weights(bert_weights)
144
- if not len(loaded):
145
- # Fix for models like `sentence-transformers/stsb-roberta-base-v2`
146
- # which use the same architecture, but have no "roberta" prefix.
147
- loaded = self.model.load_weights(task_weights)
148
- assert len(loaded), "Unable to load RobertaEmbeddingModel"
140
+ weights_list = list(weights)
141
+ has_roberta_prefix = any(
142
+ name.startswith("roberta.") for name, _ in weights_list)
143
+ if has_roberta_prefix:
144
+ # For models with the `roberta.` prefix e.g.
145
+ # `FacebookAI/roberta-base`
146
+ mapper = WeightsMapper(orig_to_new_prefix={"roberta.": "model."})
147
+ else:
148
+ # For models without the `roberta.` prefix e.g.
149
+ # `sentence-transformers/stsb-roberta-base-v2`
150
+ mapper = WeightsMapper(orig_to_new_prefix={"": "model."})
151
+
152
+ loader = AutoWeightsLoader(self, skip_prefixes=["lm_head."])
153
+ return loader.load_weights(weights_list, mapper=mapper)
154
155
156
class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding,
@@ -187,19 +192,8 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding,
192
self.classifier)
193
194
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
190
- bert_weights, task_weights = roberta_task_weights_filter(weights)
191
- bert_weights = self.jina_to_vllm_mapper.apply(bert_weights)
192
-
193
- self.roberta.load_weights(bert_weights)
194
-
195
- params_dict = dict(self.named_parameters())
196
-
197
- for name, loaded_weight in task_weights:
198
- if name.startswith("classifier"):
199
- param = params_dict[name]
200
- weight_loader = getattr(param, "weight_loader",
201
- default_weight_loader)
202
- weight_loader(param, loaded_weight)
195
+ loader = AutoWeightsLoader(self)
196
+ return loader.load_weights(weights, mapper=self.jina_to_vllm_mapper)
197
198
def pooler(
199
self,
@@ -245,27 +239,3 @@ def create_position_ids_from_input_ids(input_ids,
239
past_key_values_length) * mask
240
241
return incremental_indices.long() + padding_idx
248
-
249
-
250
-def roberta_task_weights_filter(
251
- all_weights: Iterable[tuple[str, torch.Tensor]]
252
-) -> tuple[Iterable[tuple[str, torch.Tensor]], Iterable[tuple[str,
253
- torch.Tensor]]]:
254
- """
255
- Separate task-specific weights that are applied on top
256
- of the encoder-decoder bert base.
257
- To do so, return two generators over the original iterator.
258
- Also, remove the "roberta." prefix to make it loadable
259
- from vanilla BertModel.
260
- """
261
- # Copy of a lazy iterator without in-memory overhead so both
262
- # iterators can be iterated upon independently.
263
- all_weights1, all_weights2 = itertools.tee(all_weights)
264
-
265
- def encoder_decoder_weights():
266
- for name, weight in all_weights1:
267
- if name.startswith("roberta."):
268
- yield (name[len("roberta."):], weight)
269
-
270
- return encoder_decoder_weights(), ((n, w) for n, w in all_weights2
271
- if not n.startswith("roberta."))