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 11, 11:15 UTC
[Model] Automatic conversion of TokenClassification model (#30666)
Signed-off-by: wang.yuqi <yuqi.wang@daocloud.io>
wang.yuqi committed
Dec 15, 2025 at 16:13 UTC
4429d934de3c5cc327b0d7aec8e473aeba38db90
4 files changed
+45
tests/models/language/pooling/test_token_classification.py
+31
@@ -68,3 +68,34 @@ def test_modernbert_models(
68
hf_output = torch.tensor(hf_output).cpu().float()
69
vllm_output = torch.tensor(vllm_output).cpu().float()
70
assert torch.allclose(hf_output, vllm_output, atol=1e-2)
71
+
72
+
73
+@pytest.mark.parametrize("model", ["bd2lcco/Qwen3-0.6B-finetuned"])
74
+@pytest.mark.parametrize("dtype", ["float"])
75
+@torch.inference_mode
76
+def test_auto_conversion(
77
+ hf_runner,
78
+ vllm_runner,
79
+ example_prompts,
80
+ model: str,
81
+ dtype: str,
82
+) -> None:
83
+ with vllm_runner(model, max_model_len=1024, dtype=dtype) as vllm_model:
84
+ vllm_outputs = vllm_model.token_classify(example_prompts)
85
+
86
+ with hf_runner(
87
+ model, dtype=dtype, auto_cls=AutoModelForTokenClassification
88
+ ) as hf_model:
89
+ tokenizer = hf_model.tokenizer
90
+ hf_outputs = []
91
+ for prompt in example_prompts:
92
+ inputs = tokenizer([prompt], return_tensors="pt")
93
+ inputs = hf_model.wrap_device(inputs)
94
+ output = hf_model.model(**inputs)
95
+ hf_outputs.append(softmax(output.logits[0]))
96
+
97
+ # check logits difference
98
+ for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
99
+ hf_output = torch.tensor(hf_output).cpu().float()
100
+ vllm_output = torch.tensor(vllm_output).cpu().float()
101
+ assert torch.allclose(hf_output, vllm_output, atol=1e-2)
tests/models/registry.py
+1
@@ -573,6 +573,7 @@ _AUTOMATIC_CONVERTED_MODELS = {
573
"Qwen3ForSequenceClassification": _HfExamplesInfo(
574
"tomaarsen/Qwen3-Reranker-0.6B-seq-cls"
575
),
576
+ "Qwen3ForTokenClassification": _HfExamplesInfo("bd2lcco/Qwen3-0.6B-finetuned"),
577
}
578
579
_MULTIMODAL_EXAMPLE_MODELS = {
vllm/config/model.py
+1
@@ -1796,6 +1796,7 @@ _SUFFIX_TO_DEFAULTS: list[tuple[str, tuple[RunnerType, ConvertType]]] = [
1796
("ForTextEncoding", ("pooling", "embed")),
1797
("EmbeddingModel", ("pooling", "embed")),
1798
("ForSequenceClassification", ("pooling", "classify")),
1799
+ ("ForTokenClassification", ("pooling", "classify")),
1800
("ForAudioClassification", ("pooling", "classify")),
1801
("ForImageClassification", ("pooling", "classify")),
1802
("ForVideoClassification", ("pooling", "classify")),
vllm/model_executor/models/adapters.py
+12
@@ -337,6 +337,18 @@ def as_seq_cls_model(cls: _T) -> _T:
337
tokens = getattr(text_config, "classifier_from_token", None)
338
method = getattr(text_config, "method", None)
339
340
+ def auto_set_score_bias(weights):
341
+ for name, weight in weights:
342
+ if name == "score.bias":
343
+ device = self.score.weight.device
344
+ dtype = self.score.weight.dtype
345
+ bias = weight.to(device).to(dtype)
346
+ self.score.bias = torch.nn.Parameter(bias)
347
+ self.score.skip_bias_add = False
348
+ else:
349
+ yield name, weight
350
+
351
+ weights = auto_set_score_bias(weights)
352
if tokens is None and method is None:
353
return super().load_weights(weights)
354
else: