[Model] Make the BERT/RoBERTa embedding class a class attribute (#59348)
Signed-off-by: Avishek Goswami <avishek.goswami@ibm.com> Signed-off-by: Avishek Goswami <86944690+GOavi101@users.noreply.github.com> Co-authored-by: Avishek Goswami <avishek.goswami@ibm.com>
Avishek Goswami committed
Sep 30, 2026 at 14:27 UTC
fb91712b38a68bb751526ad53552b91226290edd
2 files changed
+17
-7
vllm/model_executor/models/bert.py
+9
-4
@@ -461,6 +461,7 @@ class BertEmbeddingModel(nn.Module, SupportsQuant):
461
"""
462
463
is_pooling_model = True
464
+ embedding_class: type[nn.Module] = BertEmbedding
465
466
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
467
super().__init__()
@@ -504,7 +505,7 @@ class BertEmbeddingModel(nn.Module, SupportsQuant):
505
506
def _build_model(self, vllm_config: VllmConfig, prefix: str = "") -> BertModel:
507
return BertModel(
507
- vllm_config=vllm_config, prefix=prefix, embedding_class=BertEmbedding
508
+ vllm_config=vllm_config, prefix=prefix, embedding_class=self.embedding_class
509
)
510
511
def _build_pooler(self, pooler_config: PoolerConfig) -> Pooler:
@@ -781,6 +782,7 @@ class BertForSequenceClassification(nn.Module, SupportsCrossEncoding, SupportsQu
782
"""
783
784
is_pooling_model = True
785
+ embedding_class: type[nn.Module] = BertEmbedding
786
787
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
788
super().__init__()
@@ -790,7 +792,7 @@ class BertForSequenceClassification(nn.Module, SupportsCrossEncoding, SupportsQu
792
self.bert = BertPoolingModel(
793
vllm_config=vllm_config,
794
prefix=maybe_prefix(prefix, "bert"),
793
- embedding_class=BertEmbedding,
795
+ embedding_class=self.embedding_class,
796
)
797
self.classifier = nn.Linear(
798
config.hidden_size,
@@ -840,6 +842,7 @@ class BertForSequenceClassification(nn.Module, SupportsCrossEncoding, SupportsQu
842
@default_pooling_type(tok_pooling_type="ALL")
843
class BertForTokenClassification(nn.Module):
844
is_pooling_model = True
845
+ embedding_class: type[nn.Module] = BertEmbedding
846
847
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
848
super().__init__()
@@ -849,7 +852,7 @@ class BertForTokenClassification(nn.Module):
852
self.bert = BertModel(
853
vllm_config=vllm_config,
854
prefix=maybe_prefix(prefix, "bert"),
852
- embedding_class=BertEmbedding,
855
+ embedding_class=self.embedding_class,
856
)
857
self.classifier = nn.Linear(
858
config.hidden_size, config.num_labels, dtype=self.head_dtype
@@ -927,13 +930,15 @@ class BertForMaskedLM(nn.Module):
930
}
931
)
932
933
+ embedding_class: type[nn.Module] = BertEmbedding
934
+
935
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
936
super().__init__()
937
config = vllm_config.model_config.hf_config
938
self.bert = BertModel(
939
vllm_config=vllm_config,
940
prefix=maybe_prefix(prefix, "bert"),
936
- embedding_class=BertEmbedding,
941
+ embedding_class=self.embedding_class,
942
)
943
self.mlm_head = BertMLMHead(
944
hidden_size=config.hidden_size,
vllm/model_executor/models/roberta.py
+8
-3
@@ -117,6 +117,8 @@ class RobertaClassificationHead(nn.Module):
117
class RobertaEmbeddingModel(BertEmbeddingModel):
118
"""A model that uses Roberta to provide embedding functionalities."""
119
120
+ embedding_class: type[nn.Module] = RobertaEmbedding
121
+
122
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
123
super().__init__(vllm_config=vllm_config, prefix=prefix)
124
self.padding_idx: int = vllm_config.model_config.hf_config.pad_token_id
@@ -143,7 +145,7 @@ class RobertaEmbeddingModel(BertEmbeddingModel):
145
return BertModel(
146
vllm_config=vllm_config,
147
prefix=prefix,
146
- embedding_class=RobertaEmbedding,
148
+ embedding_class=self.embedding_class,
149
)
150
else:
151
return JinaRobertaModel(vllm_config=vllm_config, prefix=prefix)
@@ -293,6 +295,8 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding, Support
295
}
296
)
297
298
+ embedding_class: type[nn.Module] = RobertaEmbedding
299
+
300
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
301
super().__init__()
302
config = vllm_config.model_config.hf_config
@@ -303,7 +307,7 @@ class RobertaForSequenceClassification(nn.Module, SupportsCrossEncoding, Support
307
self.roberta = BertModel(
308
vllm_config=vllm_config,
309
prefix=maybe_prefix(prefix, "bert"),
306
- embedding_class=RobertaEmbedding,
310
+ embedding_class=self.embedding_class,
311
)
312
self.classifier = RobertaClassificationHead(vllm_config.model_config)
313
@@ -355,6 +359,7 @@ class RobertaForTokenClassification(nn.Module):
359
"""
360
361
is_pooling_model = True
362
+ embedding_class: type[nn.Module] = RobertaEmbedding
363
364
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
365
super().__init__()
@@ -364,7 +369,7 @@ class RobertaForTokenClassification(nn.Module):
369
self.roberta = BertModel(
370
vllm_config=vllm_config,
371
prefix=maybe_prefix(prefix, "roberta"),
367
- embedding_class=RobertaEmbedding,
372
+ embedding_class=self.embedding_class,
373
)
374
self.classifier = nn.Linear(
375
config.hidden_size, config.num_labels, dtype=self.head_dtype