From 4893270bab18c1e1262098ae57609ac2e694c5b9 Mon Sep 17 00:00:00 2001
From: Ahmet Selim FEDAKAR <122402435+selimfedakar@users.noreply.github.com>
Date: Mon, 2 Nov 2026 19:22:47 +0300
Subject: [PATCH] fix: quantize int8 in float32 to avoid float16 rounding and
 overflow

The int8 path in quantize_embeddings scaled the embeddings through a
float16 buffer before rounding. float16 keeps only ~3 significant digits,
so the scaled values round differently than the exact quotient and a
non-trivial fraction of codes flip by +/-1 even for ordinary,
small-magnitude embeddings (about 0.65% of codes, touching ~80% of rows
on a normal 2000x256 matrix at std 0.1). Any embedding value above the
float16 maximum (65504) also overflowed to inf before scaling.

Compute the scaled values in float32 instead. This matches the reference
int8 codes exactly and removes the overflow, at the cost of a transient
float32 buffer instead of float16.

Adds regression tests: the precision test fails on the previous float16
buffer, and an overflow test covers values above the float16 maximum.
---
 model2vec/quantization.py  |  9 ++++++---
 tests/test_quantization.py | 31 +++++++++++++++++++++++++++++++
 2 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/model2vec/quantization.py b/model2vec/quantization.py
index f06bcd6f..de0216ac 100644
--- a/model2vec/quantization.py
+++ b/model2vec/quantization.py
@@ -43,9 +43,12 @@ def quantize_embeddings(embeddings: np.ndarray, quantize_to: DType) -> np.ndarra
         # Normalize to [-128, 127] range for int8
         # We normalize to -127 to 127 to keep symmetry.
         scale = np.max(np.abs(embeddings)) / 127.0
-        # Turn into float16 to minimize memory usage during computation
-        # we copy once.
-        buf = embeddings.astype(np.float16, copy=True)
+        # Compute the quantization in float32. We copy once.
+        # float16 cannot be used here: its ~3 significant digits round the
+        # scaled values differently than the true quotient, which flips a
+        # non-trivial fraction of codes by +/-1, and any embedding value above
+        # the float16 maximum (65504) would overflow to inf before scaling.
+        buf = embeddings.astype(np.float32, copy=True)
         # Divide by the scale
         np.divide(buf, scale, out=buf)
         # Round to int, copy to the buffer
diff --git a/tests/test_quantization.py b/tests/test_quantization.py
index c2fbb260..5161b992 100644
--- a/tests/test_quantization.py
+++ b/tests/test_quantization.py
@@ -31,6 +31,37 @@ def test_quantize_embeddings(input_dtype: DType, target_dtype: DType, expected_d
         assert np.allclose(quantized, embeddings.astype(expected_dtype))
 
 
+def test_quantize_embeddings_int8_matches_float32_reference() -> None:
+    """Int8 quantization must round in full precision, not through a float16 buffer.
+
+    Scaling in float16 (~3 significant digits) rounds the scaled values differently
+    than the exact quotient, flipping a fraction of codes by +/-1 even for ordinary,
+    small-magnitude embeddings. The quantized codes must match the reference computed
+    directly in float32.
+    """
+    rng = np.random.default_rng(0)
+    embeddings = (rng.standard_normal((2000, 256)) * 0.1).astype(np.float32)
+
+    quantized = quantize_embeddings(embeddings, DType.Int8)
+
+    scale = np.max(np.abs(embeddings)) / 127.0
+    reference = np.clip(np.rint(embeddings.astype(np.float32) / scale), -127, 127).astype(np.int8)
+
+    assert np.array_equal(quantized, reference)
+
+
+def test_quantize_embeddings_int8_handles_values_above_float16_max() -> None:
+    """Values above the float16 maximum (65504) must not overflow to inf before scaling."""
+    embeddings = np.array([[70000.0, 1.0, 2.0, 3.0]], dtype=np.float32)
+
+    quantized = quantize_embeddings(embeddings, DType.Int8)
+
+    # The dominant value maps to the clip bound, the rest map to 0 at this scale.
+    assert quantized.dtype == np.int8
+    assert quantized[0, 0] == 127
+    assert np.all(quantized[0, 1:] == 0)
+
+
 def test_quantize_and_reduce_dim() -> None:
     """Test quantization and dimensionality reduction."""
     embeddings = np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]], dtype=np.float32)