@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 7738ef35b8

[Feat] Add Support for BertForMaskedLM to vLLM (#48463)

Signed-off-by: atalhens <sneh.lata@nutanix.com>

Snehlata committed Jul 14, 2026 at 02:26 UTC 7738ef35b8c62070773ef3cfd53f2a59d94626b1

4 files changed +158 -2

tests/models/language/pooling/test_token_classification.py

+54 -1

@@ -3,7 +3,7 @@

4 import pytest

5 import torch

6 -from transformers import AutoModelForTokenClassification

6 +from transformers import AutoModelForMaskedLM, AutoModelForTokenClassification

8 from tests.models.registry import HF_EXAMPLE_MODELS

9 from tests.models.utils import softmax

@@ -200,3 +200,56 @@ def test_auto_conversion(

200 hf_output = hf_output.detach().clone().cpu().float()

201 vllm_output = vllm_output.detach().clone().cpu().float()

202 assert torch.allclose(hf_output, vllm_output, atol=1e-2)

203 +

204 +

205 +@pytest.mark.parametrize(

206 + "model",

207 + [

208 + # Original Google checkpoint: legacy `gamma`/`beta` LayerNorm names, an

209 + # NSP head (`cls.seq_relationship.*`) and a decoder tied to the input

210 + # embeddings (no explicit decoder weight in the checkpoint).

211 + "google-bert/bert-base-uncased",

212 + ],

213 +)

214 +@pytest.mark.parametrize("dtype", ["float"])

215 +@torch.inference_mode

216 +def test_bert_for_masked_lm(

217 + hf_runner,

218 + vllm_runner,

219 + example_prompts,

220 + model: str,

221 + dtype: str,

222 +) -> None:

223 + # BertForMaskedLM exposes its MLM head as a token-level pooling task; the

224 + # head applies softmax over the vocabulary, so each output row is a

225 + # distribution (matching HF's softmax(logits) below).

226 + with vllm_runner(model, max_model_len=None, dtype=dtype) as vllm_model:

227 + vllm_outputs = vllm_model.token_classify(example_prompts)

228 +

229 + # Use eager attention on ROCm to avoid HF Transformers flash attention

230 + # accuracy issues: https://github.com/vllm-project/vllm/issues/30167

231 + hf_model_kwargs = {}

232 + if current_platform.is_rocm():

233 + hf_model_kwargs["attn_implementation"] = "eager"

234 +

235 + with hf_runner(

236 + model,

237 + dtype=dtype,

238 + auto_cls=AutoModelForMaskedLM,

239 + model_kwargs=hf_model_kwargs,

240 + ) as hf_model:

241 + tokenizer = hf_model.tokenizer

242 + hf_outputs = []

243 + for prompt in example_prompts:

244 + inputs = tokenizer([prompt], return_tensors="pt")

245 + inputs = hf_model.wrap_device(inputs)

246 + output = hf_model.model(**inputs)

247 + hf_outputs.append(softmax(output.logits[0]))

248 +

249 + # Compare the per-token vocabulary distributions position by position.

250 + for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):

251 + hf_output = hf_output.detach().clone().cpu().float()

252 + vllm_output = vllm_output.detach().clone().cpu().float()

253 + assert hf_output.shape == vllm_output.shape

254 + assert torch.equal(hf_output.argmax(dim=-1), vllm_output.argmax(dim=-1))

255 + torch.testing.assert_close(hf_output, vllm_output, atol=3.2e-2, rtol=1e-3)

tests/models/registry.py

@@ -576,6 +576,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {

576 _EMBEDDING_EXAMPLE_MODELS = {

577 # [Text-only]

578 "BertModel": _HfExamplesInfo("BAAI/bge-base-en-v1.5"),

579 + "BertForMaskedLM": _HfExamplesInfo("google-bert/bert-base-uncased"),

580 "BertSpladeSparseEmbeddingModel": _HfExamplesInfo(

581 "naver/splade-v3",

582 hf_overrides={"architectures": ["BertSpladeSparseEmbeddingModel"]},

vllm/model_executor/models/bert.py

+102 -1

@@ -370,11 +370,17 @@ class BertModel(nn.Module, SupportsQuant):

370 packed_modules_mapping = {"qkv_proj": ["query", "key", "value"]}

371

372 hf_to_vllm_mapper = WeightsMapper(

373 + # Original google-bert checkpoints use the legacy `gamma`/`beta`

374 + # LayerNorm names; rename to vLLM's `weight`/`bias`.

375 + orig_to_new_substr={

376 + "LayerNorm.gamma": "LayerNorm.weight",

377 + "LayerNorm.beta": "LayerNorm.bias",

378 + },

379 orig_to_new_stacked={

380 ".self.query": (".self.qkv_proj", "q"),

381 ".self.key": (".self.qkv_proj", "k"),

382 ".self.value": (".self.qkv_proj", "v"),

377 - }

383 + },

384 )

385

386 def __init__(

@@ -877,3 +883,98 @@ class BertForTokenClassification(nn.Module):

883

884 hidden_states = hidden_states.to(self.head_dtype)

885 return self.classifier(hidden_states)

886 +

887 +

888 +@attn_type("encoder_only")

889 +@default_pooling_type(tok_pooling_type="ALL")

890 +class BertForMaskedLM(nn.Module):

891 + """Bert with a masked-language-modeling head on top of ``BertModel``.

892 +

893 + Produces per-token logits over the vocabulary. In vLLM terms this is a

894 + token-level pooling model (``tok_pooling_type="ALL"``): the encoder output

895 + is projected by the MLM head to ``vocab_size`` logits for every position,

896 + and the token pooler returns one vector per token.

897 + """

898 +

899 + is_pooling_model = True

900 +

901 + # Map the HF ``cls.predictions.*`` checkpoint names onto our ``mlm_head.*``

902 + # submodule. Order matters: the ``None`` (drop) rules and the more specific

903 + # names are listed before the broader ``cls.predictions.decoder`` rule so

904 + # that substring replacement doesn't rewrite them first.

905 + hf_to_vllm_mapper = WeightsMapper(

906 + orig_to_new_substr={

907 + # Next-sentence-prediction head: not part of masked LM.

908 + "cls.seq_relationship": None,

909 + # Some checkpoints ship an explicit (tied) decoder bias; we load the

910 + # canonical ``cls.predictions.bias`` instead, so drop the duplicate.

911 + "cls.predictions.decoder.bias": None,

912 + # Legacy LayerNorm affine names in the MLM head transform.

913 + "cls.predictions.transform.LayerNorm.gamma": "mlm_head.layer_norm.weight",

914 + "cls.predictions.transform.LayerNorm.beta": "mlm_head.layer_norm.bias",

915 + "cls.predictions.transform.LayerNorm": "mlm_head.layer_norm",

916 + "cls.predictions.transform.dense": "mlm_head.dense",

917 + "cls.predictions.decoder": "mlm_head.decoder",

918 + # In HF ``cls.predictions.bias`` *is* the decoder bias.

919 + "cls.predictions.bias": "mlm_head.decoder.bias",

920 + }

921 + )

922 +

923 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

924 + super().__init__()

925 + config = vllm_config.model_config.hf_config

926 + self.bert = BertModel(

927 + vllm_config=vllm_config,

928 + prefix=maybe_prefix(prefix, "bert"),

929 + embedding_class=BertEmbedding,

930 + )

931 + self.mlm_head = BertMLMHead(

932 + hidden_size=config.hidden_size,

933 + vocab_size=config.vocab_size,

934 + layer_norm_eps=getattr(config, "layer_norm_eps", 1e-12),

935 + )

936 +

937 + pooler_config = vllm_config.model_config.pooler_config

938 + assert pooler_config is not None

939 +

940 + self.pooler = pooler_for_token_classify(pooler_config)

941 +

942 + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:

943 + return self.bert.embed_input_ids(input_ids)

944 +

945 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):

946 + loader = AutoWeightsLoader(self)

947 + loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

948 +

949 + # The MLM decoder shares its weight with the input embeddings. When the

950 + # checkpoint relies on `tie_word_embeddings` (e.g. google-bert/*) it

951 + # doesn't ship an explicit decoder weight, so tie it here. Copy the

952 + # first `vocab_size` rows since VocabParallelEmbedding may pad the vocab.

953 + if "mlm_head.decoder.weight" not in loaded:

954 + emb = self.bert.embeddings.word_embeddings.weight

955 + decoder = self.mlm_head.decoder.weight

956 + decoder.data.copy_(emb.data[: decoder.shape[0]])

957 + loaded.add("mlm_head.decoder.weight")

958 + return loaded

959 +

960 + def forward(

961 + self,

962 + input_ids: torch.Tensor | None,

963 + positions: torch.Tensor,

964 + intermediate_tensors: IntermediateTensors | None = None,

965 + inputs_embeds: torch.Tensor | None = None,

966 + token_type_ids: torch.Tensor | None = None,

967 + ) -> torch.Tensor:

968 + if token_type_ids is not None:

969 + assert self.bert.config.vocab_size < (1 << TOKEN_TYPE_SHIFT)

970 + assert input_ids is not None

971 + _encode_token_type_ids(input_ids, token_type_ids)

972 +

973 + hidden_states = self.bert(

974 + input_ids=input_ids,

975 + positions=positions,

976 + inputs_embeds=inputs_embeds,

977 + intermediate_tensors=intermediate_tensors,

978 + )

979 +

980 + return self.mlm_head(hidden_states)

vllm/model_executor/models/registry.py

@@ -210,6 +210,7 @@ _TEXT_GENERATION_MODELS = {

210 _EMBEDDING_MODELS = {

211 # [Text-only]

212 "BertModel": ("bert", "BertEmbeddingModel"),

213 + "BertForMaskedLM": ("bert", "BertForMaskedLM"),

214 "BertSpladeSparseEmbeddingModel": ("bert", "BertSpladeSparseEmbeddingModel"),

215 "BgeM3EmbeddingModel": ("roberta", "BgeM3EmbeddingModel"),

216 "DeciLMForCausalLM": ("nemotron_nas", "DeciLMForCausalLM"),