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 14, 05:30 UTC
[Model] Add support for ModernBertForTokenClassification (#26340)
Signed-off-by: Antoine Recanati Le Goat <antoine.recanati@sancare.fr> Signed-off-by: antrec <antoine.recanati@gmail.com> Co-authored-by: Antoine Recanati Le Goat <antoine.recanati@sancare.fr> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>
antrec committed
Oct 7, 2025 at 16:29 UTC
6f59beaf0b1f2ec922c027bcee5bf52042af7430
5 files changed
+112
-2
docs/models/supported_models.md
+1
@@ -576,6 +576,7 @@ These models primarily support the [`LLM.encode`](./pooling_models.md#llmencode)
576
| Architecture | Models | Example HF Models | [LoRA](../features/lora.md) | [PP](../serving/parallelism_scaling.md) | [V1](gh-issue:8779) |
577
|--------------|--------|-------------------|-----------------------------|-----------------------------------------|---------------------|
578
| `BertForTokenClassification` | bert-based | `boltuix/NeuroBERT-NER` (see note), etc. | | | ✅︎ |
579
+| `ModernBertForTokenClassification` | ModernBERT-based | `disham993/electrical-ner-ModernBERT-base` | | | ✅︎ |
580
581
!!! note
582
Named Entity Recognition (NER) usage, please refer to <gh-file:examples/offline_inference/pooling/ner.py>, <gh-file:examples/online_serving/pooling/ner.py>.
tests/models/language/pooling/test_token_classification.py
+32
-1
@@ -11,7 +11,38 @@ from tests.models.utils import softmax
11
# The float32 is required for this tiny model to pass the test.
12
@pytest.mark.parametrize("dtype", ["float"])
13
@torch.inference_mode
14
-def test_models(
14
+def test_bert_models(
15
+ hf_runner,
16
+ vllm_runner,
17
+ example_prompts,
18
+ model: str,
19
+ dtype: str,
20
+) -> None:
21
+ with vllm_runner(model, max_model_len=None, dtype=dtype) as vllm_model:
22
+ vllm_outputs = vllm_model.encode(example_prompts)
23
+
24
+ with hf_runner(
25
+ model, dtype=dtype, auto_cls=AutoModelForTokenClassification
26
+ ) as hf_model:
27
+ tokenizer = hf_model.tokenizer
28
+ hf_outputs = []
29
+ for prompt in example_prompts:
30
+ inputs = tokenizer([prompt], return_tensors="pt")
31
+ inputs = hf_model.wrap_device(inputs)
32
+ output = hf_model.model(**inputs)
33
+ hf_outputs.append(softmax(output.logits[0]))
34
+
35
+ # check logits difference
36
+ for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
37
+ hf_output = torch.tensor(hf_output).cpu().float()
38
+ vllm_output = torch.tensor(vllm_output).cpu().float()
39
+ assert torch.allclose(hf_output, vllm_output, 1e-2)
40
+
41
+
42
+@pytest.mark.parametrize("model", ["disham993/electrical-ner-ModernBERT-base"])
43
+@pytest.mark.parametrize("dtype", ["float"])
44
+@torch.inference_mode
45
+def test_modernbert_models(
46
hf_runner,
47
vllm_runner,
48
example_prompts,
tests/models/registry.py
+3
@@ -527,6 +527,9 @@ _SEQUENCE_CLASSIFICATION_EXAMPLE_MODELS = {
527
"ModernBertForSequenceClassification": _HfExamplesInfo(
528
"Alibaba-NLP/gte-reranker-modernbert-base"
529
),
530
+ "ModernBertForTokenClassification": _HfExamplesInfo(
531
+ "disham993/electrical-ner-ModernBERT-base"
532
+ ),
533
"RobertaForSequenceClassification": _HfExamplesInfo(
534
"cross-encoder/quora-roberta-base"
535
),
vllm/model_executor/models/modernbert.py
+72
-1
@@ -6,6 +6,7 @@ from typing import Optional, Union
6
import torch
7
from torch import nn
8
from transformers import ModernBertConfig
9
+from transformers.activations import ACT2FN
10
11
from vllm.attention.layers.encoder_only_attention import EncoderOnlyAttention
12
from vllm.compilation.decorators import support_torch_compile
@@ -29,7 +30,7 @@ from vllm.v1.pool.metadata import PoolingMetadata
30
31
from .interfaces import SupportsCrossEncoding
32
from .interfaces_base import default_pooling_type
32
-from .utils import WeightsMapper, maybe_prefix
33
+from .utils import AutoWeightsLoader, WeightsMapper, maybe_prefix
34
35
36
class ModernBertEmbeddings(nn.Module):
@@ -379,3 +380,73 @@ class ModernBertForSequenceClassification(nn.Module, SupportsCrossEncoding):
380
inputs_embeds=inputs_embeds,
381
positions=positions,
382
)
383
+
384
+
385
+class ModernBertPredictionHead(nn.Module):
386
+ def __init__(self, config):
387
+ super().__init__()
388
+ self.config = config
389
+ self.dense = nn.Linear(
390
+ config.hidden_size, config.hidden_size, bias=config.classifier_bias
391
+ )
392
+ self.act = ACT2FN[config.classifier_activation]
393
+ self.norm = nn.LayerNorm(
394
+ config.hidden_size,
395
+ eps=getattr(config, "norm_eps", 1e-5),
396
+ bias=getattr(config, "norm_bias", True),
397
+ )
398
+
399
+ def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
400
+ return self.norm(self.act(self.dense(hidden_states)))
401
+
402
+
403
+@default_pooling_type("ALL")
404
+class ModernBertForTokenClassification(nn.Module):
405
+ is_pooling_model = True
406
+
407
+ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
408
+ super().__init__()
409
+ config = vllm_config.model_config.hf_config
410
+ self.head_dtype = vllm_config.model_config.head_dtype
411
+ self.num_labels = config.num_labels
412
+ self.model = ModernBertModel(
413
+ vllm_config=vllm_config, prefix=maybe_prefix(prefix, "modernbert")
414
+ )
415
+ self.head = ModernBertPredictionHead(config)
416
+ self.classifier = nn.Linear(
417
+ config.hidden_size, config.num_labels, dtype=self.head_dtype
418
+ )
419
+
420
+ pooler_config = vllm_config.model_config.pooler_config
421
+ assert pooler_config is not None
422
+
423
+ self.pooler = DispatchPooler(
424
+ {
425
+ "encode": Pooler.for_encode(pooler_config),
426
+ }
427
+ )
428
+
429
+ def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
430
+ return self.model.get_input_embeddings(input_ids)
431
+
432
+ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
433
+ loader = AutoWeightsLoader(self, skip_prefixes=["drop"])
434
+ loaded_params = loader.load_weights(weights)
435
+ return loaded_params
436
+
437
+ def forward(
438
+ self,
439
+ input_ids: Optional[torch.Tensor],
440
+ positions: torch.Tensor,
441
+ intermediate_tensors: Optional[IntermediateTensors] = None,
442
+ inputs_embeds: Optional[torch.Tensor] = None,
443
+ ) -> torch.Tensor:
444
+ hidden_states = self.model(
445
+ input_ids=input_ids,
446
+ positions=positions,
447
+ inputs_embeds=inputs_embeds,
448
+ intermediate_tensors=intermediate_tensors,
449
+ )
450
+ hidden_states = self.head(hidden_states)
451
+ hidden_states = hidden_states.to(self.head_dtype)
452
+ return self.classifier(hidden_states)
vllm/model_executor/models/registry.py
+4
@@ -225,6 +225,10 @@ _CROSS_ENCODER_MODELS = {
225
"modernbert",
226
"ModernBertForSequenceClassification",
227
),
228
+ "ModernBertForTokenClassification": (
229
+ "modernbert",
230
+ "ModernBertForTokenClassification",
231
+ ),
232
"RobertaForSequenceClassification": ("roberta", "RobertaForSequenceClassification"),
233
"XLMRobertaForSequenceClassification": (
234
"roberta",