2fd6268e6d · tduncan/Nvidia-Optimization-Stack-VLLM

tests/models/language/pooling_mteb_test/mteb_embed_utils.py

+10 -2

@@ -161,10 +161,14 @@ def mteb_test_embed_models(

161 model_info: EmbedModelInfo,

162 vllm_extra_kwargs=None,

163 hf_model_callback=None,

164 - atol=MTEB_EMBED_TOL,

164 + atol: float | None = None,

165 prompt_prefix: str | None = None,

166 vllm_model_callback=None,

167 ):

168 + if atol is None:

169 + atol = (

170 + model_info.mteb_tol if model_info.mteb_tol is not None else MTEB_EMBED_TOL

171 + )

172 vllm_extra_kwargs = get_vllm_extra_kwargs(model_info, vllm_extra_kwargs)

173

174 # Test embed_dims, isnan and whether to use normalize

@@ -268,4 +272,8 @@ def mteb_test_embed_models(

272

273 # We are not concerned that the vllm mteb results are better

274 # than SentenceTransformers, so we only perform one-sided testing.

271 - assert st_main_score - vllm_main_score < atol

275 + diff = st_main_score - vllm_main_score

276 + assert diff < atol, (

277 + f"diff={diff:.6g} tol={atol} model={model_info.name} "

278 + f"(st={st_main_score}, vllm={vllm_main_score})"

279 + )

tests/models/language/pooling_mteb_test/mteb_score_utils.py

+10 -2

@@ -230,8 +230,12 @@ def mteb_test_rerank_models(

230 hf_runner=HFMtebCrossEncoder,

231 vllm_extra_kwargs=None,

232 vllm_mteb_encoder=VllmMtebCrossEncoder,

233 - atol=MTEB_RERANK_TOL,

233 + atol: float | None = None,

234 ):

235 + if atol is None:

236 + atol = (

237 + model_info.mteb_tol if model_info.mteb_tol is not None else MTEB_RERANK_TOL

238 + )

239 vllm_extra_kwargs = get_vllm_extra_kwargs(model_info, vllm_extra_kwargs)

240

241 # Maybe load chat_template.

@@ -308,4 +312,8 @@ def mteb_test_rerank_models(

312

313 # We are not concerned that the vllm mteb results are better

314 # than SentenceTransformers, so we only perform one-sided testing.

311 - assert st_main_score - vllm_main_score < atol

315 + diff = st_main_score - vllm_main_score

316 + assert diff < atol, (

317 + f"diff={diff:.6g} tol={atol} model={model_info.name} "

318 + f"(st={st_main_score}, vllm={vllm_main_score})"

319 + )

tests/models/language/pooling_mteb_test/test_baai.py

+5

@@ -17,6 +17,7 @@ MODELS = [

17 "BAAI/bge-base-en",

18 architecture="BertModel",

19 mteb_score=0.779336792,

20 + mteb_tol=2e-3,

21 seq_pooling_type="CLS",

22 attn_type="encoder_only",

23 is_prefix_caching_supported=False,

@@ -54,6 +55,7 @@ MODELS = [

55 "BAAI/bge-m3",

56 architecture="XLMRobertaModel",

57 mteb_score=0.787343078,

58 + mteb_tol=2e-3,

59 seq_pooling_type="CLS",

60 attn_type="encoder_only",

61 is_prefix_caching_supported=False,

@@ -82,6 +84,7 @@ RERANK_MODELS = [

84 "BAAI/bge-reranker-base",

85 architecture="XLMRobertaForSequenceClassification",

86 mteb_score=0.32398,

87 + mteb_tol=1e-2,

88 seq_pooling_type="CLS",

89 attn_type="encoder_only",

90 is_prefix_caching_supported=False,

@@ -101,6 +104,7 @@ RERANK_MODELS = [

104 ]

105

106

107 +@pytest.mark.flaky(reruns=2)

108 @pytest.mark.parametrize("model_info", MODELS)

109 def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:

110 mteb_test_embed_models(hf_runner, vllm_runner, model_info)

@@ -113,6 +117,7 @@ def test_embed_models_correctness(

117 correctness_test_embed_models(hf_runner, vllm_runner, model_info, example_prompts)

118

119

120 +@pytest.mark.flaky(reruns=2)

121 @pytest.mark.parametrize("model_info", RERANK_MODELS)

122 def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:

123 mteb_test_rerank_models(vllm_runner, model_info)

tests/models/language/pooling_mteb_test/test_bge_reranker_v2_gemma.py

+2

@@ -26,6 +26,7 @@ RERANK_MODELS = [

26 "method": "no_post_processing",

27 },

28 mteb_score=0.33757,

29 + mteb_tol=1e-2,

30 seq_pooling_type="LAST",

31 attn_type="decoder",

32 is_prefix_caching_supported=True,

@@ -127,6 +128,7 @@ class GemmaRerankerHfRunner(MtebCrossEncoderMixin, HfRunner):

128 return torch.Tensor(scores)

129

130

131 +@pytest.mark.flaky(reruns=2)

132 @pytest.mark.parametrize("model_info", RERANK_MODELS)

133 def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:

134 mteb_test_rerank_models(

tests/models/language/pooling_mteb_test/test_cross_encoder.py

+3

@@ -17,6 +17,7 @@ RERANK_MODELS = [

17 is_prefix_caching_supported=False,

18 is_chunked_prefill_supported=False,

19 mteb_score=0.32898,

20 + mteb_tol=1e-2,

21 ),

22 RerankModelInfo(

23 "tomaarsen/Qwen3-Reranker-0.6B-seq-cls",

@@ -27,10 +28,12 @@ RERANK_MODELS = [

28 is_chunked_prefill_supported=True,

29 chat_template_name="qwen3_reranker.jinja",

30 mteb_score=0.33459,

31 + mteb_tol=1e-2,

32 ),

33 ]

34

35

36 +@pytest.mark.flaky(reruns=2)

37 @pytest.mark.parametrize("model_info", RERANK_MODELS)

38 def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:

39 mteb_test_rerank_models(vllm_runner, model_info)

tests/models/language/pooling_mteb_test/test_gte.py

+8

@@ -18,6 +18,7 @@ MODELS = [

18 EmbedModelInfo(

19 "thenlper/gte-large",

20 mteb_score=0.76807651,

21 + mteb_tol=2e-3,

22 architecture="BertModel",

23 seq_pooling_type="MEAN",

24 attn_type="encoder_only",

@@ -44,6 +45,7 @@ MODELS = [

45 "Alibaba-NLP/gte-multilingual-base",

46 architecture="GteNewModel",

47 mteb_score=0.775074696,

48 + mteb_tol=2e-3,

49 hf_overrides={"architectures": ["GteNewModel"]},

50 seq_pooling_type="CLS",

51 attn_type="encoder_only",

@@ -79,6 +81,7 @@ MODELS = [

81 EmbedModelInfo(

82 "Alibaba-NLP/gte-modernbert-base",

83 mteb_score=0.748193353,

84 + mteb_tol=2e-3,

85 architecture="ModernBertModel",

86 seq_pooling_type="CLS",

87 attn_type="encoder_only",

@@ -90,6 +93,7 @@ MODELS = [

93 EmbedModelInfo(

94 "Qwen/Qwen3-Embedding-0.6B",

95 mteb_score=0.771163695,

96 + mteb_tol=2e-3,

97 architecture="Qwen3ForCausalLM",

98 seq_pooling_type="LAST",

99 attn_type="decoder",

@@ -109,6 +113,7 @@ RERANK_MODELS = [

113 # classifier_pooling: mean

114 "Alibaba-NLP/gte-reranker-modernbert-base",

115 mteb_score=0.33386,

116 + mteb_tol=1e-2,

117 architecture="ModernBertForSequenceClassification",

118 seq_pooling_type="CLS",

119 attn_type="encoder_only",

@@ -119,6 +124,7 @@ RERANK_MODELS = [

124 RerankModelInfo(

125 "Alibaba-NLP/gte-multilingual-reranker-base",

126 mteb_score=0.33062,

127 + mteb_tol=1e-2,

128 architecture="GteNewForSequenceClassification",

129 hf_overrides={"architectures": ["GteNewForSequenceClassification"]},

130 seq_pooling_type="CLS",

@@ -130,6 +136,7 @@ RERANK_MODELS = [

136 ]

137

138

139 +@pytest.mark.flaky(reruns=2)

140 @pytest.mark.parametrize("model_info", MODELS)

141 def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:

142 mteb_test_embed_models(hf_runner, vllm_runner, model_info)

@@ -142,6 +149,7 @@ def test_embed_models_correctness(

149 correctness_test_embed_models(hf_runner, vllm_runner, model_info, example_prompts)

150

151

152 +@pytest.mark.flaky(reruns=2)

153 @pytest.mark.parametrize("model_info", RERANK_MODELS)

154 def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:

155 vllm_extra_kwargs = {}

tests/models/language/pooling_mteb_test/test_intfloat.py

+3

@@ -13,6 +13,7 @@ MODELS = [

13 "intfloat/e5-small",

14 architecture="BertModel",

15 mteb_score=0.742285423,

16 + mteb_tol=2e-3,

17 seq_pooling_type="MEAN",

18 attn_type="encoder_only",

19 is_prefix_caching_supported=False,

@@ -29,6 +30,7 @@ MODELS = [

30 "intfloat/multilingual-e5-base",

31 architecture="XLMRobertaModel",

32 mteb_score=0.779325955,

33 + mteb_tol=2e-3,

34 seq_pooling_type="MEAN",

35 attn_type="encoder_only",

36 is_prefix_caching_supported=False,

@@ -48,6 +50,7 @@ MODELS = [

50 ]

51

52

53 +@pytest.mark.flaky(reruns=2)

54 @pytest.mark.parametrize("model_info", MODELS)

55 def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:

56 mteb_test_embed_models(hf_runner, vllm_runner, model_info)

tests/models/language/pooling_mteb_test/test_jina.py

+6

@@ -23,6 +23,7 @@ EMBEDDING_MODELS = [

23 EmbedModelInfo(

24 "jinaai/jina-embeddings-v3",

25 mteb_score=0.824413164,

26 + mteb_tol=2e-3,

27 architecture="XLMRobertaModel",

28 is_matryoshka=True,

29 seq_pooling_type="MEAN",

@@ -33,6 +34,7 @@ EMBEDDING_MODELS = [

34 EmbedModelInfo(

35 "jinaai/jina-embeddings-v5-text-small",

36 mteb_score=0.794535707854956,

37 + mteb_tol=2e-3,

38 architecture="JinaEmbeddingsV5Model",

39 seq_pooling_type="LAST",

40 attn_type="decoder",

@@ -42,6 +44,7 @@ EMBEDDING_MODELS = [

44 EmbedModelInfo(

45 "jinaai/jina-embeddings-v5-text-nano",

46 architecture="JinaEmbeddingsV5Model",

47 + mteb_tol=2e-3,

48 dtype="bfloat16" if current_platform.is_rocm() else "auto",

49 seq_pooling_type="LAST",

50 attn_type="encoder_only",

@@ -54,6 +57,7 @@ RERANK_MODELS = [

57 RerankModelInfo(

58 "jinaai/jina-reranker-v2-base-multilingual",

59 mteb_score=0.33643,

60 + mteb_tol=1e-2,

61 architecture="XLMRobertaForSequenceClassification",

62 seq_pooling_type="CLS",

63 attn_type="encoder_only",

@@ -63,6 +67,7 @@ RERANK_MODELS = [

67 ]

68

69

70 +@pytest.mark.flaky(reruns=2)

71 @pytest.mark.parametrize("model_info", EMBEDDING_MODELS)

72 def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:

73 task = "retrieval" if "v5" in model_info.name else "text-matching"

@@ -98,6 +103,7 @@ def test_embed_models_correctness(

103 )

104

105

106 +@pytest.mark.flaky(reruns=2)

107 @pytest.mark.parametrize("model_info", RERANK_MODELS)

108 def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:

109 mteb_test_rerank_models(vllm_runner, model_info)

tests/models/language/pooling_mteb_test/test_modernbert_fp8.py

+1

@@ -56,6 +56,7 @@ def _assert_modernbert_online_fp8(model) -> None:

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.",

tests/models/language/pooling_mteb_test/test_mxbai_rerank.py

+2

@@ -30,6 +30,7 @@ RERANK_MODELS = [

30 is_chunked_prefill_supported=True,

31 chat_template_name="mxbai_rerank_v2.jinja",

32 mteb_score=0.33651,

33 + mteb_tol=1e-2,

34 enable_test=True,

35 ),

36 RerankModelInfo(

@@ -104,6 +105,7 @@ class MxbaiRerankerHfRunner(MtebCrossEncoderMixin, HfRunner):

105 return torch.Tensor(scores)

106

107

108 +@pytest.mark.flaky(reruns=2)

109 @pytest.mark.parametrize("model_info", RERANK_MODELS)

110 def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:

111 mteb_test_rerank_models(vllm_runner, model_info, hf_runner=MxbaiRerankerHfRunner)

tests/models/language/pooling_mteb_test/test_nemotron.py

+4

@@ -19,6 +19,7 @@ EMBEDDING_MODELS = [

19 "nvidia/llama-nemotron-embed-1b-v2",

20 architecture="LlamaBidirectionalModel",

21 mteb_score=0.689164662128673,

22 + mteb_tol=2e-3,

23 seq_pooling_type="MEAN",

24 attn_type="encoder_only",

25 is_prefix_caching_supported=False,

@@ -32,6 +33,7 @@ RERANK_MODELS = [

33 architecture="LlamaBidirectionalForSequenceClassification",

34 chat_template_name="nemotron-rerank.jinja",

35 mteb_score=0.33994,

36 + mteb_tol=1e-2,

37 seq_pooling_type="MEAN",

38 attn_type="encoder_only",

39 is_prefix_caching_supported=False,

@@ -40,11 +42,13 @@ RERANK_MODELS = [

42 ]

43

44

45 +@pytest.mark.flaky(reruns=2)

46 @pytest.mark.parametrize("model_info", EMBEDDING_MODELS)

47 def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:

48 mteb_test_embed_models(hf_runner, vllm_runner, model_info)

49

50

51 +@pytest.mark.flaky(reruns=2)

52 @pytest.mark.parametrize("model_info", RERANK_MODELS)

53 def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:

54 mteb_test_rerank_models(vllm_runner, model_info)

tests/models/language/pooling_mteb_test/test_nomic.py

+3

@@ -17,6 +17,7 @@ MODELS = [

17 # HF version and remove revision set.

18 revision="720244025c1a7e15661a174c63cce63c8218e52b",

19 mteb_score=0.737568559,

20 + mteb_tol=2e-3,

21 enable_test=True,

22 seq_pooling_type="MEAN",

23 attn_type="encoder_only",

@@ -35,6 +36,7 @@ MODELS = [

36 "nomic-ai/nomic-embed-text-v2-moe",

37 architecture="NomicBertModel",

38 mteb_score=0.715488912,

39 + mteb_tol=2e-3,

40 enable_test=True,

41 seq_pooling_type="MEAN",

42 attn_type="encoder_only",

@@ -44,6 +46,7 @@ MODELS = [

46 ]

47

48

49 +@pytest.mark.flaky(reruns=2)

50 @pytest.mark.parametrize("model_info", MODELS)

51 def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:

52 mteb_test_embed_models(hf_runner, vllm_runner, model_info)

tests/models/language/pooling_mteb_test/test_qwen3_reranker.py

+3

@@ -32,6 +32,7 @@ RERANK_MODELS = [

32 is_prefix_caching_supported=True,

33 is_chunked_prefill_supported=True,

34 mteb_score=0.33459,

35 + mteb_tol=1e-2,

36 enable_test=True,

37 ),

38 RerankModelInfo(

@@ -108,11 +109,13 @@ class Qwen3RerankerHfRunner(MtebCrossEncoderMixin, HfRunner):

109 return torch.Tensor(scores)

110

111

112 +@pytest.mark.flaky(reruns=2)

113 @pytest.mark.parametrize("model_info", RERANK_MODELS)

114 def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:

115 mteb_test_rerank_models(vllm_runner, model_info, hf_runner=Qwen3RerankerHfRunner)

116

117

118 +@pytest.mark.flaky(reruns=2)

119 @pytest.mark.parametrize("model_info", RERANK_MODELS)

120 @multi_gpu_test(num_gpus=2)

121 def test_rerank_models_mteb_tp(vllm_runner, model_info: RerankModelInfo) -> None:

tests/models/language/pooling_mteb_test/test_snowflake_arctic_embed.py

+6

@@ -14,6 +14,7 @@ MODELS = [

14 is_matryoshka=False,

15 architecture="BertModel",

16 mteb_score=0.714927797,

17 + mteb_tol=2e-3,

18 seq_pooling_type="CLS",

19 attn_type="encoder_only",

20 is_prefix_caching_supported=False,

@@ -37,6 +38,7 @@ MODELS = [

38 is_matryoshka=False,

39 architecture="NomicBertModel",

40 mteb_score=0.681146831,

41 + mteb_tol=2e-3,

42 seq_pooling_type="CLS",

43 attn_type="encoder_only",

44 is_prefix_caching_supported=False,

@@ -54,6 +56,7 @@ MODELS = [

56 is_matryoshka=True,

57 architecture="BertModel",

58 mteb_score=0.649088363,

59 + mteb_tol=2e-3,

60 seq_pooling_type="CLS",

61 attn_type="encoder_only",

62 is_prefix_caching_supported=False,

@@ -65,6 +68,7 @@ MODELS = [

68 is_matryoshka=True,

69 architecture="XLMRobertaModel",

70 mteb_score=0.712258299,

71 + mteb_tol=2e-3,

72 seq_pooling_type="CLS",

73 attn_type="encoder_only",

74 is_prefix_caching_supported=False,

@@ -76,6 +80,7 @@ MODELS = [

80 is_matryoshka=True,

81 architecture="GteModel",

82 mteb_score=0.706622444,

83 + mteb_tol=2e-3,

84 seq_pooling_type="CLS",

85 attn_type="encoder_only",

86 is_prefix_caching_supported=False,

@@ -85,6 +90,7 @@ MODELS = [

90 ]

91

92

93 +@pytest.mark.flaky(reruns=2)

94 @pytest.mark.parametrize("model_info", MODELS)

95 def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:

96 mteb_test_embed_models(hf_runner, vllm_runner, model_info)

tests/models/language/pooling_mteb_test/test_st_projector.py

+3

@@ -14,6 +14,7 @@ ST_PROJECTOR_MODELS = [

14 "TencentBAC/Conan-embedding-v1",

15 architecture="BertModel",

16 mteb_score=0.688611955,

17 + mteb_tol=2e-3,

18 seq_pooling_type="MEAN",

19 attn_type="encoder_only",

20 is_prefix_caching_supported=False,

@@ -24,6 +25,7 @@ ST_PROJECTOR_MODELS = [

25 "google/embeddinggemma-300m",

26 architecture="Gemma3TextModel",

27 mteb_score=0.7473819294684156,

28 + mteb_tol=2e-3,

29 seq_pooling_type="MEAN",

30 attn_type="encoder_only",

31 is_prefix_caching_supported=False,

@@ -33,6 +35,7 @@ ST_PROJECTOR_MODELS = [

35 ]

36

37

38 +@pytest.mark.flaky(reruns=2)

39 @pytest.mark.parametrize("model_info", ST_PROJECTOR_MODELS)

40 def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:

41 mteb_test_embed_models(hf_runner, vllm_runner, model_info)

tests/models/language/pooling_mteb_test/test_voyage.py

+2

@@ -22,6 +22,7 @@ MODELS = [

22 "num_labels": 2048,

23 },

24 mteb_score=0.7054,

25 + mteb_tol=2e-3,

26 # === MTEB Results ===

27 # STS12: 0.6613

28 # STS13: 0.6906

@@ -34,6 +35,7 @@ MODELS = [

35 ]

36

37

38 +@pytest.mark.flaky(reruns=2)

39 @pytest.mark.parametrize("model_info", MODELS)

40 def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:

41 # Encoder-only attention models need enforce_eager=True to avoid

tests/models/utils.py

+4

@@ -396,6 +396,8 @@ class ModelInfo:

396 @dataclass

397 class EmbedModelInfo(ModelInfo):

398 mteb_score: float | None = None

399 + # Per-model MTEB score tolerance; falls back to MTEB_EMBED_TOL.

400 + mteb_tol: float | None = None

401 is_matryoshka: bool = False

402 matryoshka_dimensions: list[int] | None = None

403

@@ -403,6 +405,8 @@ class EmbedModelInfo(ModelInfo):

405 @dataclass

406 class RerankModelInfo(ModelInfo):

407 mteb_score: float | None = None

408 + # Per-model MTEB score tolerance; falls back to MTEB_RERANK_TOL.

409 + mteb_tol: float | None = None

410 chat_template_name: str | None = None

411

412

[CI] Deflake MTEB score tests with per-model mteb_tol and reruns (#57364)