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

Signed-off-by: khluu <khluu000@gmail.com>

tests/models/language/pooling_mteb_test/mteb_embed_utils.py

@@ -161,10 +161,14 @@ def mteb_test_embed_models(
     model_info: EmbedModelInfo,
     vllm_extra_kwargs=None,
     hf_model_callback=None,
-    atol=MTEB_EMBED_TOL,
+    atol: float | None = None,
     prompt_prefix: str | None = None,
     vllm_model_callback=None,
 ):
+    if atol is None:
+        atol = (
+            model_info.mteb_tol if model_info.mteb_tol is not None else MTEB_EMBED_TOL
+        )
     vllm_extra_kwargs = get_vllm_extra_kwargs(model_info, vllm_extra_kwargs)
 
     # Test embed_dims, isnan and whether to use normalize
@@ -268,4 +272,8 @@ def mteb_test_embed_models(
 
     # We are not concerned that the vllm mteb results are better
     # than SentenceTransformers, so we only perform one-sided testing.
-    assert st_main_score - vllm_main_score < atol
+    diff = st_main_score - vllm_main_score
+    assert diff < atol, (
+        f"diff={diff:.6g} tol={atol} model={model_info.name} "
+        f"(st={st_main_score}, vllm={vllm_main_score})"
+    )

tests/models/language/pooling_mteb_test/mteb_score_utils.py

@@ -230,8 +230,12 @@ def mteb_test_rerank_models(
     hf_runner=HFMtebCrossEncoder,
     vllm_extra_kwargs=None,
     vllm_mteb_encoder=VllmMtebCrossEncoder,
-    atol=MTEB_RERANK_TOL,
+    atol: float | None = None,
 ):
+    if atol is None:
+        atol = (
+            model_info.mteb_tol if model_info.mteb_tol is not None else MTEB_RERANK_TOL
+        )
     vllm_extra_kwargs = get_vllm_extra_kwargs(model_info, vllm_extra_kwargs)
 
     # Maybe load chat_template.
@@ -308,4 +312,8 @@ def mteb_test_rerank_models(
 
     # We are not concerned that the vllm mteb results are better
     # than SentenceTransformers, so we only perform one-sided testing.
-    assert st_main_score - vllm_main_score < atol
+    diff = st_main_score - vllm_main_score
+    assert diff < atol, (
+        f"diff={diff:.6g} tol={atol} model={model_info.name} "
+        f"(st={st_main_score}, vllm={vllm_main_score})"
+    )

tests/models/language/pooling_mteb_test/test_baai.py

@@ -17,6 +17,7 @@ MODELS = [
         "BAAI/bge-base-en",
         architecture="BertModel",
         mteb_score=0.779336792,
+        mteb_tol=2e-3,
         seq_pooling_type="CLS",
         attn_type="encoder_only",
         is_prefix_caching_supported=False,
@@ -54,6 +55,7 @@ MODELS = [
         "BAAI/bge-m3",
         architecture="XLMRobertaModel",
         mteb_score=0.787343078,
+        mteb_tol=2e-3,
         seq_pooling_type="CLS",
         attn_type="encoder_only",
         is_prefix_caching_supported=False,
@@ -82,6 +84,7 @@ RERANK_MODELS = [
         "BAAI/bge-reranker-base",
         architecture="XLMRobertaForSequenceClassification",
         mteb_score=0.32398,
+        mteb_tol=1e-2,
         seq_pooling_type="CLS",
         attn_type="encoder_only",
         is_prefix_caching_supported=False,
@@ -101,6 +104,7 @@ RERANK_MODELS = [
 ]
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", MODELS)
 def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
     mteb_test_embed_models(hf_runner, vllm_runner, model_info)
@@ -113,6 +117,7 @@ def test_embed_models_correctness(
     correctness_test_embed_models(hf_runner, vllm_runner, model_info, example_prompts)
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", RERANK_MODELS)
 def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
     mteb_test_rerank_models(vllm_runner, model_info)

tests/models/language/pooling_mteb_test/test_bge_reranker_v2_gemma.py

@@ -26,6 +26,7 @@ RERANK_MODELS = [
             "method": "no_post_processing",
         },
         mteb_score=0.33757,
+        mteb_tol=1e-2,
         seq_pooling_type="LAST",
         attn_type="decoder",
         is_prefix_caching_supported=True,
@@ -127,6 +128,7 @@ class GemmaRerankerHfRunner(MtebCrossEncoderMixin, HfRunner):
         return torch.Tensor(scores)
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", RERANK_MODELS)
 def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
     mteb_test_rerank_models(

tests/models/language/pooling_mteb_test/test_cross_encoder.py

@@ -17,6 +17,7 @@ RERANK_MODELS = [
         is_prefix_caching_supported=False,
         is_chunked_prefill_supported=False,
         mteb_score=0.32898,
+        mteb_tol=1e-2,
     ),
     RerankModelInfo(
         "tomaarsen/Qwen3-Reranker-0.6B-seq-cls",
@@ -27,10 +28,12 @@ RERANK_MODELS = [
         is_chunked_prefill_supported=True,
         chat_template_name="qwen3_reranker.jinja",
         mteb_score=0.33459,
+        mteb_tol=1e-2,
     ),
 ]
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", RERANK_MODELS)
 def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
     mteb_test_rerank_models(vllm_runner, model_info)

tests/models/language/pooling_mteb_test/test_gte.py

@@ -18,6 +18,7 @@ MODELS = [
     EmbedModelInfo(
         "thenlper/gte-large",
         mteb_score=0.76807651,
+        mteb_tol=2e-3,
         architecture="BertModel",
         seq_pooling_type="MEAN",
         attn_type="encoder_only",
@@ -44,6 +45,7 @@ MODELS = [
         "Alibaba-NLP/gte-multilingual-base",
         architecture="GteNewModel",
         mteb_score=0.775074696,
+        mteb_tol=2e-3,
         hf_overrides={"architectures": ["GteNewModel"]},
         seq_pooling_type="CLS",
         attn_type="encoder_only",
@@ -79,6 +81,7 @@ MODELS = [
     EmbedModelInfo(
         "Alibaba-NLP/gte-modernbert-base",
         mteb_score=0.748193353,
+        mteb_tol=2e-3,
         architecture="ModernBertModel",
         seq_pooling_type="CLS",
         attn_type="encoder_only",
@@ -90,6 +93,7 @@ MODELS = [
     EmbedModelInfo(
         "Qwen/Qwen3-Embedding-0.6B",
         mteb_score=0.771163695,
+        mteb_tol=2e-3,
         architecture="Qwen3ForCausalLM",
         seq_pooling_type="LAST",
         attn_type="decoder",
@@ -109,6 +113,7 @@ RERANK_MODELS = [
         # classifier_pooling: mean
         "Alibaba-NLP/gte-reranker-modernbert-base",
         mteb_score=0.33386,
+        mteb_tol=1e-2,
         architecture="ModernBertForSequenceClassification",
         seq_pooling_type="CLS",
         attn_type="encoder_only",
@@ -119,6 +124,7 @@ RERANK_MODELS = [
     RerankModelInfo(
         "Alibaba-NLP/gte-multilingual-reranker-base",
         mteb_score=0.33062,
+        mteb_tol=1e-2,
         architecture="GteNewForSequenceClassification",
         hf_overrides={"architectures": ["GteNewForSequenceClassification"]},
         seq_pooling_type="CLS",
@@ -130,6 +136,7 @@ RERANK_MODELS = [
 ]
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", MODELS)
 def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
     mteb_test_embed_models(hf_runner, vllm_runner, model_info)
@@ -142,6 +149,7 @@ def test_embed_models_correctness(
     correctness_test_embed_models(hf_runner, vllm_runner, model_info, example_prompts)
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", RERANK_MODELS)
 def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
     vllm_extra_kwargs = {}

tests/models/language/pooling_mteb_test/test_intfloat.py

@@ -13,6 +13,7 @@ MODELS = [
         "intfloat/e5-small",
         architecture="BertModel",
         mteb_score=0.742285423,
+        mteb_tol=2e-3,
         seq_pooling_type="MEAN",
         attn_type="encoder_only",
         is_prefix_caching_supported=False,
@@ -29,6 +30,7 @@ MODELS = [
         "intfloat/multilingual-e5-base",
         architecture="XLMRobertaModel",
         mteb_score=0.779325955,
+        mteb_tol=2e-3,
         seq_pooling_type="MEAN",
         attn_type="encoder_only",
         is_prefix_caching_supported=False,
@@ -48,6 +50,7 @@ MODELS = [
 ]
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", MODELS)
 def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
     mteb_test_embed_models(hf_runner, vllm_runner, model_info)

tests/models/language/pooling_mteb_test/test_jina.py

@@ -23,6 +23,7 @@ EMBEDDING_MODELS = [
     EmbedModelInfo(
         "jinaai/jina-embeddings-v3",
         mteb_score=0.824413164,
+        mteb_tol=2e-3,
         architecture="XLMRobertaModel",
         is_matryoshka=True,
         seq_pooling_type="MEAN",
@@ -33,6 +34,7 @@ EMBEDDING_MODELS = [
     EmbedModelInfo(
         "jinaai/jina-embeddings-v5-text-small",
         mteb_score=0.794535707854956,
+        mteb_tol=2e-3,
         architecture="JinaEmbeddingsV5Model",
         seq_pooling_type="LAST",
         attn_type="decoder",
@@ -42,6 +44,7 @@ EMBEDDING_MODELS = [
     EmbedModelInfo(
         "jinaai/jina-embeddings-v5-text-nano",
         architecture="JinaEmbeddingsV5Model",
+        mteb_tol=2e-3,
         dtype="bfloat16" if current_platform.is_rocm() else "auto",
         seq_pooling_type="LAST",
         attn_type="encoder_only",
@@ -54,6 +57,7 @@ RERANK_MODELS = [
     RerankModelInfo(
         "jinaai/jina-reranker-v2-base-multilingual",
         mteb_score=0.33643,
+        mteb_tol=1e-2,
         architecture="XLMRobertaForSequenceClassification",
         seq_pooling_type="CLS",
         attn_type="encoder_only",
@@ -63,6 +67,7 @@ RERANK_MODELS = [
 ]
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", EMBEDDING_MODELS)
 def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
     task = "retrieval" if "v5" in model_info.name else "text-matching"
@@ -98,6 +103,7 @@ def test_embed_models_correctness(
     )
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", RERANK_MODELS)
 def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
     mteb_test_rerank_models(vllm_runner, model_info)

tests/models/language/pooling_mteb_test/test_modernbert_fp8.py

@@ -56,6 +56,7 @@ def _assert_modernbert_online_fp8(model) -> None:
         )
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.skipif(
     _fp8_scaled_mm_unsupported(),
     reason="No FP8 ScaledMM kernel is available on this GPU type.",

tests/models/language/pooling_mteb_test/test_mxbai_rerank.py

@@ -30,6 +30,7 @@ RERANK_MODELS = [
         is_chunked_prefill_supported=True,
         chat_template_name="mxbai_rerank_v2.jinja",
         mteb_score=0.33651,
+        mteb_tol=1e-2,
         enable_test=True,
     ),
     RerankModelInfo(
@@ -104,6 +105,7 @@ class MxbaiRerankerHfRunner(MtebCrossEncoderMixin, HfRunner):
         return torch.Tensor(scores)
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", RERANK_MODELS)
 def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
     mteb_test_rerank_models(vllm_runner, model_info, hf_runner=MxbaiRerankerHfRunner)

tests/models/language/pooling_mteb_test/test_nemotron.py

@@ -19,6 +19,7 @@ EMBEDDING_MODELS = [
         "nvidia/llama-nemotron-embed-1b-v2",
         architecture="LlamaBidirectionalModel",
         mteb_score=0.689164662128673,
+        mteb_tol=2e-3,
         seq_pooling_type="MEAN",
         attn_type="encoder_only",
         is_prefix_caching_supported=False,
@@ -32,6 +33,7 @@ RERANK_MODELS = [
         architecture="LlamaBidirectionalForSequenceClassification",
         chat_template_name="nemotron-rerank.jinja",
         mteb_score=0.33994,
+        mteb_tol=1e-2,
         seq_pooling_type="MEAN",
         attn_type="encoder_only",
         is_prefix_caching_supported=False,
@@ -40,11 +42,13 @@ RERANK_MODELS = [
 ]
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", EMBEDDING_MODELS)
 def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
     mteb_test_embed_models(hf_runner, vllm_runner, model_info)
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", RERANK_MODELS)
 def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
     mteb_test_rerank_models(vllm_runner, model_info)

tests/models/language/pooling_mteb_test/test_nomic.py

@@ -17,6 +17,7 @@ MODELS = [
         #  HF version and remove revision set.
         revision="720244025c1a7e15661a174c63cce63c8218e52b",
         mteb_score=0.737568559,
+        mteb_tol=2e-3,
         enable_test=True,
         seq_pooling_type="MEAN",
         attn_type="encoder_only",
@@ -35,6 +36,7 @@ MODELS = [
         "nomic-ai/nomic-embed-text-v2-moe",
         architecture="NomicBertModel",
         mteb_score=0.715488912,
+        mteb_tol=2e-3,
         enable_test=True,
         seq_pooling_type="MEAN",
         attn_type="encoder_only",
@@ -44,6 +46,7 @@ MODELS = [
 ]
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", MODELS)
 def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
     mteb_test_embed_models(hf_runner, vllm_runner, model_info)

tests/models/language/pooling_mteb_test/test_qwen3_reranker.py

@@ -32,6 +32,7 @@ RERANK_MODELS = [
         is_prefix_caching_supported=True,
         is_chunked_prefill_supported=True,
         mteb_score=0.33459,
+        mteb_tol=1e-2,
         enable_test=True,
     ),
     RerankModelInfo(
@@ -108,11 +109,13 @@ class Qwen3RerankerHfRunner(MtebCrossEncoderMixin, HfRunner):
         return torch.Tensor(scores)
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", RERANK_MODELS)
 def test_rerank_models_mteb(vllm_runner, model_info: RerankModelInfo) -> None:
     mteb_test_rerank_models(vllm_runner, model_info, hf_runner=Qwen3RerankerHfRunner)
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", RERANK_MODELS)
 @multi_gpu_test(num_gpus=2)
 def test_rerank_models_mteb_tp(vllm_runner, model_info: RerankModelInfo) -> None:

tests/models/language/pooling_mteb_test/test_snowflake_arctic_embed.py

@@ -14,6 +14,7 @@ MODELS = [
         is_matryoshka=False,
         architecture="BertModel",
         mteb_score=0.714927797,
+        mteb_tol=2e-3,
         seq_pooling_type="CLS",
         attn_type="encoder_only",
         is_prefix_caching_supported=False,
@@ -37,6 +38,7 @@ MODELS = [
         is_matryoshka=False,
         architecture="NomicBertModel",
         mteb_score=0.681146831,
+        mteb_tol=2e-3,
         seq_pooling_type="CLS",
         attn_type="encoder_only",
         is_prefix_caching_supported=False,
@@ -54,6 +56,7 @@ MODELS = [
         is_matryoshka=True,
         architecture="BertModel",
         mteb_score=0.649088363,
+        mteb_tol=2e-3,
         seq_pooling_type="CLS",
         attn_type="encoder_only",
         is_prefix_caching_supported=False,
@@ -65,6 +68,7 @@ MODELS = [
         is_matryoshka=True,
         architecture="XLMRobertaModel",
         mteb_score=0.712258299,
+        mteb_tol=2e-3,
         seq_pooling_type="CLS",
         attn_type="encoder_only",
         is_prefix_caching_supported=False,
@@ -76,6 +80,7 @@ MODELS = [
         is_matryoshka=True,
         architecture="GteModel",
         mteb_score=0.706622444,
+        mteb_tol=2e-3,
         seq_pooling_type="CLS",
         attn_type="encoder_only",
         is_prefix_caching_supported=False,
@@ -85,6 +90,7 @@ MODELS = [
 ]
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", MODELS)
 def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
     mteb_test_embed_models(hf_runner, vllm_runner, model_info)

tests/models/language/pooling_mteb_test/test_st_projector.py

@@ -14,6 +14,7 @@ ST_PROJECTOR_MODELS = [
         "TencentBAC/Conan-embedding-v1",
         architecture="BertModel",
         mteb_score=0.688611955,
+        mteb_tol=2e-3,
         seq_pooling_type="MEAN",
         attn_type="encoder_only",
         is_prefix_caching_supported=False,
@@ -24,6 +25,7 @@ ST_PROJECTOR_MODELS = [
         "google/embeddinggemma-300m",
         architecture="Gemma3TextModel",
         mteb_score=0.7473819294684156,
+        mteb_tol=2e-3,
         seq_pooling_type="MEAN",
         attn_type="encoder_only",
         is_prefix_caching_supported=False,
@@ -33,6 +35,7 @@ ST_PROJECTOR_MODELS = [
 ]
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", ST_PROJECTOR_MODELS)
 def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
     mteb_test_embed_models(hf_runner, vllm_runner, model_info)

tests/models/language/pooling_mteb_test/test_voyage.py

@@ -22,6 +22,7 @@ MODELS = [
             "num_labels": 2048,
         },
         mteb_score=0.7054,
+        mteb_tol=2e-3,
         # === MTEB Results ===
         # STS12: 0.6613
         # STS13: 0.6906
@@ -34,6 +35,7 @@ MODELS = [
 ]
 
 
+@pytest.mark.flaky(reruns=2)
 @pytest.mark.parametrize("model_info", MODELS)
 def test_embed_models_mteb(hf_runner, vllm_runner, model_info: EmbedModelInfo) -> None:
     # Encoder-only attention models need enforce_eager=True to avoid

tests/models/utils.py

@@ -396,6 +396,8 @@ class ModelInfo:
 @dataclass
 class EmbedModelInfo(ModelInfo):
     mteb_score: float | None = None
+    # Per-model MTEB score tolerance; falls back to MTEB_EMBED_TOL.
+    mteb_tol: float | None = None
     is_matryoshka: bool = False
     matryoshka_dimensions: list[int] | None = None
 
@@ -403,6 +405,8 @@ class EmbedModelInfo(ModelInfo):
 @dataclass
 class RerankModelInfo(ModelInfo):
     mteb_score: float | None = None
+    # Per-model MTEB score tolerance; falls back to MTEB_RERANK_TOL.
+    mteb_tol: float | None = None
     chat_template_name: str | None = None