Merge pull request #771 from superlinked/fix/bge-flash-default-fp16

fix(server): align BGE Flash profiles with FP16 default

Author: huronat

packages/sie_server/models/BAAI__bge-m3.yaml

@@ -20,7 +20,7 @@ max_sequence_length: 8192
 profiles:
   default:
     max_batch_tokens: 16384
-    compute_precision: bfloat16
+    compute_precision: float16
     adapter_path: sie_server.adapters.bge_m3_flash:BGEM3FlashAdapter
     adapter_options:
       loadtime: {}
@@ -53,13 +53,15 @@ profiles:
         - multivector
   candle:
     extends: default
+    compute_precision: bfloat16
     adapter_path: sie_server_rust.adapters.candle:CandleEmbeddingAdapter
     adapter_options:
       runtime:
         pooling: cls
         normalize: true
   banking:
     extends: default
+    compute_precision: bfloat16
     adapter_options:
       runtime:
         pooling: cls
@@ -68,6 +70,7 @@ profiles:
         instruction: Classify banking intent
   medical-vn:
     extends: default
+    compute_precision: bfloat16
     adapter_options:
       runtime:
         pooling: cls

packages/sie_server/tests/adapters/test_bge_m3_flash.py

@@ -7,10 +7,13 @@
 import numpy as np
 import pytest
 import torch
+import transformers
 import yaml
+from sie_server.adapters._types import ComputePrecision
 from sie_server.adapters.bge_m3_flash import BGEM3FlashAdapter
 from sie_server.config.model import ModelConfig
 from sie_server.core.inference_output import EncodeOutput, SparseVector
+from sie_server.core.loader import _build_adapter_kwargs
 from sie_server.types.inputs import Item
 from transformers import XLMRobertaConfig
 from transformers.models.xlm_roberta.modeling_xlm_roberta import XLMRobertaEmbeddings
@@ -93,6 +96,68 @@ def test_score_task_enabled(self, config: ModelConfig) -> None:
     def test_score_in_outputs(self, config: ModelConfig) -> None:
         assert "score" in config.outputs
 
+    def test_resolved_profile_precisions(self, config: ModelConfig) -> None:
+        assert {name: config.resolve_profile(name).compute_precision for name in config.profiles} == {
+            "default": "float16",
+            "dense": "float16",
+            "sparse": "float16",
+            "multivector": "float16",
+            "candle": "bfloat16",
+            "banking": "bfloat16",
+            "medical-vn": "bfloat16",
+            "bge_m3_flag": None,
+        }
+
+    def test_default_precision_reaches_adapter_loading(
+        self,
+        config: ModelConfig,
+        monkeypatch: pytest.MonkeyPatch,
+    ) -> None:
+        kwargs = _build_adapter_kwargs(config, "bfloat16")
+        assert kwargs["compute_precision"] == "float16"
+        assert kwargs["revision"] == "5617a9f61b028005a4858fdac845db406aefb181"
+        model_loader = MagicMock()
+        tokenizer_loader = MagicMock()
+        monkeypatch.setattr(transformers.AutoModel, "from_pretrained", model_loader)
+        monkeypatch.setattr(transformers.AutoTokenizer, "from_pretrained", tokenizer_loader)
+        adapter = BGEM3FlashAdapter(**kwargs)
+        head_loader = MagicMock()
+        monkeypatch.setattr(adapter, "_load_linear_layers", head_loader)
+
+        adapter.load("cuda:0")
+
+        model_loader.assert_called_once_with(
+            "BAAI/bge-m3",
+            torch_dtype=torch.float16,
+            attn_implementation="eager",
+            revision=config.hf_revision,
+        )
+        tokenizer_loader.assert_called_once_with("BAAI/bge-m3", revision=config.hf_revision)
+        head_loader.assert_called_once_with("BAAI/bge-m3", torch.float16, "cuda:0")
+
+    @pytest.mark.parametrize("name", ["default", "dense", "sparse", "multivector"])
+    def test_flash_profiles_preserve_recipe(self, config: ModelConfig, name: str) -> None:
+        profile = config.resolve_profile(name)
+        assert profile.adapter_path == "sie_server.adapters.bge_m3_flash:BGEM3FlashAdapter"
+        assert profile.max_batch_tokens == 16384
+        assert profile.loadtime == {}
+        expected_runtime: dict[str, Any] = {"pooling": "cls", "normalize": True}
+        if name != "default":
+            expected_runtime["output_types"] = [name]
+        assert profile.runtime == expected_runtime
+        assert config.max_sequence_length == 8192
+        assert config.dims == {"dense": 1024, "sparse": 250002, "multivector": 1024}
+
+    @pytest.mark.parametrize(
+        ("precision", "expected_dtype"),
+        [("float16", torch.float16), ("bfloat16", torch.bfloat16)],
+    )
+    def test_explicit_flash_precision_remains_supported(
+        self, precision: ComputePrecision, expected_dtype: torch.dtype
+    ) -> None:
+        adapter = BGEM3FlashAdapter(compute_precision=precision)
+        assert adapter._resolve_dtype() == expected_dtype
+
 
 class TestBGEM3EmbeddingParity:
     @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])

packages/sie_server/tests/test_all_models.py

@@ -451,27 +451,26 @@ def test_openai_clip_vit_large_patch14_image_dense() -> None:
 # =============================================================================
 
 
-_BGE_M3_CPU_BFLOAT16 = pytest.mark.xfail(
+_BGE_M3_CPU_FLOAT16 = pytest.mark.xfail(
     reason=(
-        "The default profile's compute_precision: bfloat16 also applies on CPU, so BGEM3Adapter runs in "
-        "bfloat16 and drifts from the float32 FlagEmbedding output pinned here (bge_m3_flag still matches it)"
+        "The default profile's compute_precision: float16 also applies on CPU, so BGEM3Adapter's "
+        "dense and multivector outputs differ from the float32 FlagEmbedding reference pinned here"
     ),
     raises=_NumericalReferenceMismatchError,
     strict=True,
 )
 
 
-@_BGE_M3_CPU_BFLOAT16
+@_BGE_M3_CPU_FLOAT16
 def test_baai_bge_m3_dense() -> None:
     _assert_dense("BAAI/bge-m3", 1024, [0.001247406005859375, 0.0228271484375, -0.02349853515625])
 
 
-@_BGE_M3_CPU_BFLOAT16
 def test_baai_bge_m3_sparse() -> None:
     _assert_sparse("BAAI/bge-m3", [3034], [0.342041015625])
 
 
-@_BGE_M3_CPU_BFLOAT16
+@_BGE_M3_CPU_FLOAT16
 def test_baai_bge_m3_multivector() -> None:
     _assert_multivector("BAAI/bge-m3", 1024, [0.01641845703125, 0.0304718017578125, -0.036468505859375])
 

packages/sie_server/tests/test_all_models_assertions.py

@@ -9,6 +9,18 @@
 from packages.sie_server.tests import test_all_models as models
 
 
+@pytest.mark.parametrize("kind", ["dense", "sparse", "multivector"])
+def test_bge_m3_cpu_numeric_xfail_applies_only_to_dense_and_multivector(kind: str) -> None:
+    test = getattr(models, f"test_baai_bge_m3_{kind}")
+    markers = [mark for mark in getattr(test, "pytestmark", []) if mark.name == "xfail"]
+    if kind == "sparse":
+        assert markers == []
+    else:
+        assert len(markers) == 1
+        assert markers[0].kwargs["strict"] is True
+        assert markers[0].kwargs["raises"] is models._NumericalReferenceMismatchError
+
+
 @pytest.mark.parametrize("kind", ["dense", "sparse", "multivector"])
 @pytest.mark.parametrize("defect", [None, "absent", "shape", "nonfinite", "numerical"])
 def test_bge_m3_xfail_preserves_output_structure(
@@ -37,9 +49,9 @@ def test_bge_m3_xfail_preserves_output_structure(
     output = SimpleNamespace(**{kind: payload})
     monkeypatch.setattr(models, "_get_adapter", lambda _: SimpleNamespace(encode=lambda *_args, **_kwargs: output))
 
-    expected_failure = models._BGE_M3_CPU_BFLOAT16.kwargs["raises"]
+    expected_failure = models._BGE_M3_CPU_FLOAT16.kwargs["raises"]
     assert expected_failure is models._NumericalReferenceMismatchError
-    assert models._BGE_M3_CPU_BFLOAT16.kwargs["strict"] is True
+    assert models._BGE_M3_CPU_FLOAT16.kwargs["strict"] is True
     if defect is None:
         check()
     elif defect == "numerical":