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 10, 10:15 UTC
[ROCm][CI] Fix test_token_classification.py::test_bert_models (#31993)
Signed-off-by: Divakar Verma <divakar.verma@amd.com>
Divakar Verma committed
Jan 8, 2026 at 22:04 UTC
a1648c4045f2aba8894efa0e607c5940f6337d85
1 file changed
+12
-4
tests/models/language/pooling/test_token_classification.py
+12
-4
@@ -5,6 +5,7 @@ import torch
5
from transformers import AutoModelForTokenClassification
6
7
from tests.models.utils import softmax
8
+from vllm.platforms import current_platform
9
10
11
@pytest.mark.parametrize("model", ["boltuix/NeuroBERT-NER"])
@@ -21,8 +22,17 @@ def test_bert_models(
22
with vllm_runner(model, max_model_len=None, dtype=dtype) as vllm_model:
23
vllm_outputs = vllm_model.token_classify(example_prompts)
24
25
+ # Use eager attention on ROCm to avoid HF Transformers flash attention
26
+ # accuracy issues: https://github.com/vllm-project/vllm/issues/30167
27
+ hf_model_kwargs = {}
28
+ if current_platform.is_rocm():
29
+ hf_model_kwargs["attn_implementation"] = "eager"
30
+
31
with hf_runner(
25
- model, dtype=dtype, auto_cls=AutoModelForTokenClassification
32
+ model,
33
+ dtype=dtype,
34
+ auto_cls=AutoModelForTokenClassification,
35
+ model_kwargs=hf_model_kwargs,
36
) as hf_model:
37
tokenizer = hf_model.tokenizer
38
hf_outputs = []
@@ -36,7 +46,7 @@ def test_bert_models(
46
for hf_output, vllm_output in zip(hf_outputs, vllm_outputs):
47
hf_output = hf_output.detach().clone().cpu().float()
48
vllm_output = vllm_output.detach().clone().cpu().float()
39
- assert torch.allclose(hf_output, vllm_output, 1e-2)
49
+ torch.testing.assert_close(hf_output, vllm_output, atol=1.2e-2, rtol=1e-3)
50
51
52
@pytest.mark.parametrize("model", ["disham993/electrical-ner-ModernBERT-base"])
@@ -49,8 +59,6 @@ def test_modernbert_models(
59
model: str,
60
dtype: str,
61
) -> None:
52
- from vllm.platforms import current_platform
53
-
62
with vllm_runner(model, max_model_len=None, dtype=dtype) as vllm_model:
63
vllm_outputs = vllm_model.token_classify(example_prompts)
64