Refactor base BertModel and Roberta*Model to use AutoWeightLoader.

RobertaForSequenceClassification:
- python3 -m vllm.entrypoints.cli.main serve cardiffnlp/twitter-roberta-base-sentiment-latest --served-model-name roberta-sentiment --trust-remote-code
- python3 -m vllm.entrypoints.cli.main serve jinaai/jina-embeddings-v3 --served-model-name jina-v3 --trust-remote-code

RobertaEmbeddingMode:
- python3 -m vllm.entrypoints.cli.main serve FacebookAI/roberta-base --served-model-name roberta-base --trust-remote-code
- python3 -m vllm.entrypoints.cli.main serve sentence-transformers/stsb-roberta-base-v2 --served-model-name stsb-roberta --trust-remote-code

BertEmbeddingModel:
- python3 -m vllm.entrypoints.cli.main serve sentence-transformers/all-MiniLM-L6-v2 --served-model-name bert-embeddings --trust-remote-code

Signed-off-by:  <islandhe@gmail.com>
Signed-off-by: Jen H <islandhe@gmail.com>

Author: jennifurhe

vllm/model_executor/models/bert.py

@@ -22,7 +22,6 @@
 from vllm.model_executor.layers.quantization import QuantizationConfig
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     VocabParallelEmbedding)
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.pooling_metadata import PoolingMetadata
 from vllm.sequence import IntermediateTensors, PoolerOutput
 
@@ -44,9 +43,13 @@ def __init__(self, config: BertConfig):
             config.type_vocab_size, config.hidden_size)
         self.LayerNorm = nn.LayerNorm(config.hidden_size,
                                       eps=config.layer_norm_eps)
-        self.register_buffer(
-            "position_ids",
-            torch.arange(config.max_position_embeddings).expand((1, -1)))
+
+        # Use nn.Parameter with requires_grad=False to maintain compatibility
+        # with existing HF checkpoints while ensuring position_ids are
+        # non-trainable.
+        self.position_ids = nn.Parameter(torch.empty(
+            (1, config.max_position_embeddings)),
+                                         requires_grad=False)
 
         self.position_embedding_type = config.position_embedding_type
         if self.position_embedding_type != "absolute":
@@ -359,45 +362,44 @@ def load_weights(self, weights: Iterable[tuple[str,
             ("qkv_proj", "value", "v"),
         ]
 
+        loaded_stacked_params = []
+        other_weights = []
         params_dict = dict(self.named_parameters())
-        loaded_params: set[str] = set()
         for name, loaded_weight in weights:
-            if self.pooler is None and "pooler" in name:
-                continue
             for (param_name, weight_name, shard_id) in stacked_params_mapping:
                 if weight_name not in name:
                     continue
+
                 name = name.replace(weight_name, param_name)
-                # Skip loading extra bias for GPTQ models.
-                if name.endswith(".bias") and name not in params_dict:
+                if name not in params_dict:
                     continue
                 param = params_dict[name]
                 weight_loader = param.weight_loader
                 weight_loader(param, loaded_weight, shard_id)
+                loaded_stacked_params.append(name)
                 break
             else:
-                # Skip loading extra bias for GPTQ models.
-                if name.endswith(".bias") and name not in params_dict:
-                    continue
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader",
-                                        default_weight_loader)
-                weight_loader(param, loaded_weight)
-            loaded_params.add(name)
+                other_weights.append((name, loaded_weight))
+
+        loader = AutoWeightsLoader(
+            self,
+            skip_prefixes=(["pooler."] if self.pooler is None else []),
+        )
+        loaded_params = loader.load_weights(other_weights)
+        loaded_params.update(loaded_stacked_params)
         return loaded_params
 
 
 class BertEmbeddingModel(nn.Module, SupportsV0Only, SupportsQuant):
     """A model that uses Bert to provide embedding functionalities.
 
-   This class encapsulates the BertModel and provides an interface for
-   embedding operations and customized pooling functions.
+    This class encapsulates the BertModel and provides an interface for
+    embedding operations and customized pooling functions.
 
-   Attributes:
-       model: An instance of BertModel used for forward operations.
-       _pooler: An instance of Pooler used for pooling operations.
-   """
-    hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"model.": ""})
+    Attributes:
+        model: An instance of BertModel used for forward operations.
+        _pooler: An instance of Pooler used for pooling operations.
+    """
 
     def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         super().__init__()
@@ -426,10 +428,15 @@ def pooler(
         return self._pooler(hidden_states, pooling_metadata)
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
-        weights = self.hf_to_vllm_mapper.apply(weights)
-        weights = ((name, data) for name, data in weights
-                   if not name.startswith("lm_head."))
-        self.model.load_weights(weights)
+        weights_list = list(weights)
+
+        has_model_prefix = any(
+            name.startswith("model.") for name, _ in weights_list)
+        if not has_model_prefix:
+            mapper = WeightsMapper(orig_to_new_prefix={"": "model."})
+
+        loader = AutoWeightsLoader(self, skip_prefixes=["lm_head."])
+        return loader.load_weights(weights_list, mapper=mapper)
 
     def _build_model(self,
                      vllm_config: VllmConfig,
@@ -471,27 +478,8 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
                                         self.classifier, self.bert.pooler)
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
-        bert_weights = []
-        classifier_weights = []
-
-        for name, weight in weights:
-            if name.startswith("bert."):
-                bert_weights.append((name, weight))
-            else:
-                classifier_weights.append((name, weight))
-
         loader = AutoWeightsLoader(self)
-        loaded_params = loader.load_weights(bert_weights)
-
-        params_dict = dict(self.named_parameters())
-        for name, loaded_weight in classifier_weights:
-            if name in params_dict:
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader",
-                                        default_weight_loader)
-                weight_loader(param, loaded_weight)
-                loaded_params.add(name)
-
+        loaded_params = loader.load_weights(weights)
         return loaded_params
 
     def pooler(

vllm/model_executor/models/roberta.py

@@ -13,9 +13,9 @@
 from vllm.model_executor.layers.pooler import ClassifierPooler
 from vllm.model_executor.layers.vocab_parallel_embedding import (
     VocabParallelEmbedding)
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
 from vllm.model_executor.models.bert import BertEmbeddingModel, BertModel
-from vllm.model_executor.models.utils import WeightsMapper, maybe_prefix
+from vllm.model_executor.models.utils import (AutoWeightsLoader, WeightsMapper,
+                                              maybe_prefix)
 from vllm.model_executor.pooling_metadata import PoolingMetadata
 from vllm.sequence import IntermediateTensors, PoolerOutput
 
@@ -136,16 +136,20 @@ def _build_model(self,
                              embedding_class=RobertaEmbedding)
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
-        weights = self.hf_to_vllm_mapper.apply(weights)
-        # Separate weights in "roberta"-prefixed and all else (not in memory).
-        # For use with models like FacebookAI/roberta-base.
-        bert_weights, task_weights = roberta_task_weights_filter(weights)
-        loaded = self.model.load_weights(bert_weights)
-        if not len(loaded):
-            # Fix for models like `sentence-transformers/stsb-roberta-base-v2`
-            # which use the same architecture, but have no "roberta" prefix.
-            loaded = self.model.load_weights(task_weights)
-        assert len(loaded), "Unable to load RobertaEmbeddingModel"
+        weights_list = list(weights)
+        has_roberta_prefix = any(
+            name.startswith("roberta.") for name, _ in weights_list)
+        if has_roberta_prefix:
+            # For models with the `roberta.` prefix e.g.
+            # `FacebookAI/roberta-base`
+            mapper = WeightsMapper(orig_to_new_prefix={"roberta.": "model."})
+        else:
+            # For models without the `roberta.` prefix e.g.
+            # `sentence-transformers/stsb-roberta-base-v2`
+            mapper = WeightsMapper(orig_to_new_prefix={"": "model."})
+
+        loader = AutoWeightsLoader(self, skip_prefixes=["lm_head."])
+        return loader.load_weights(weights_list, mapper=mapper)
 
 
 class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding,
@@ -187,19 +191,8 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
                                         self.classifier)
 
     def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
-        bert_weights, task_weights = roberta_task_weights_filter(weights)
-        bert_weights = self.jina_to_vllm_mapper.apply(bert_weights)
-
-        self.roberta.load_weights(bert_weights)
-
-        params_dict = dict(self.named_parameters())
-
-        for name, loaded_weight in task_weights:
-            if name.startswith("classifier"):
-                param = params_dict[name]
-                weight_loader = getattr(param, "weight_loader",
-                                        default_weight_loader)
-                weight_loader(param, loaded_weight)
+        loader = AutoWeightsLoader(self)
+        return loader.load_weights(weights, mapper=self.jina_to_vllm_mapper)
 
     def pooler(
         self,