main
py 78 lines 2.39 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import pytest
5
6 from tests.models.utils import EmbedModelInfo
7 from tests.quantization.utils import is_quant_method_supported
8 from vllm.platforms import current_platform
9
10 from .mteb_embed_utils import mteb_test_embed_models
11
12
13 def _fp8_scaled_mm_unsupported() -> bool:
14 """Whether the ModernBERT online FP8 test should be skipped.
15
16 ``is_quant_method_supported("fp8")`` returns True on MI250 (gfx90a), but
17 serving a per-tensor FP8 linear layer on ROCm requires CDNA3+ or RDNA4,
18 so treat it as unsupported here.
19 """
20 if not is_quant_method_supported("fp8"):
21 return True
22 if current_platform.is_rocm():
23 from vllm.platforms.rocm import get_cdna_version, on_gfx12x
24
25 return get_cdna_version() <= 2 and not on_gfx12x()
26 return False
27
28
29 MODEL_INFO = EmbedModelInfo(
30 "Alibaba-NLP/gte-modernbert-base",
31 mteb_score=0.748193353,
32 architecture="ModernBertModel",
33 seq_pooling_type="CLS",
34 attn_type="encoder_only",
35 is_prefix_caching_supported=False,
36 is_chunked_prefill_supported=False,
37 enable_test=True,
38 )
39
40
41 def _assert_modernbert_online_fp8(model) -> None:
42 from vllm.model_executor.layers.quantization.online.fp8 import (
43 Fp8PerTensorOnlineLinearMethod,
44 )
45
46 layer = model.encoder_layer.layers[0]
47 linears = {
48 "attn.Wqkv": layer.attn.Wqkv,
49 "attn.Wo": layer.attn.Wo,
50 "mlp.Wi": layer.mlp.Wi,
51 "mlp.Wo": layer.mlp.Wo,
52 }
53 for name, linear in linears.items():
54 assert isinstance(linear.quant_method, Fp8PerTensorOnlineLinearMethod), (
55 f"{name} was not quantized with online FP8"
56 )
57
58
59 @pytest.mark.flaky(reruns=2)
60 @pytest.mark.skipif(
61 _fp8_scaled_mm_unsupported(),
62 reason="No FP8 ScaledMM kernel is available on this GPU type.",
63 )
64 def test_modernbert_online_fp8_mteb(hf_runner, vllm_runner, monkeypatch) -> None:
65 monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
66
67 def check_online_fp8(vllm_model) -> None:
68 vllm_model.apply_model(_assert_modernbert_online_fp8)
69
70 mteb_test_embed_models(
71 hf_runner,
72 vllm_runner,
73 MODEL_INFO,
74 vllm_extra_kwargs={"quantization": "fp8_per_tensor"},
75 # Account for Online FP8 kernel variance across GPU architectures.
76 atol=2e-3,
77 vllm_model_callback=check_online_fp8,
78 )