float16 gives NaN or bad vectors

What the model card says

From the EmbeddingGemma 2 model card: run inference in bfloat16 or float32, not float16. The activations are larger than float16 can hold. In float16 the model returns NaN or quietly worse embeddings instead of an error. bfloat16 is safe because it has the same exponent range as float32. Use float32 on most CPUs.

This matters most when you run the model yourself, for example with PyTorch, or when a mobile runtime or GPU delegate picks float16 for speed. Wrong vectors with no error happen in other stacks too, for other reasons: immich #32257 reports wrong CLIP embeddings on Apple Silicon with nothing in the logs. That is why the check below runs on every vector, whatever the runtime.

In Python with PyTorch

import torch

if torch.cuda.is_available() and torch.cuda.is_bf16_supported():
    dtype = torch.bfloat16
elif torch.backends.mps.is_available():
    dtype = torch.float32   # safe default on Apple Silicon
else:
    dtype = torch.float32
# Never: torch.float16

This is the pattern from the model card. We did not run the PyTorch path on the test box, so check the card for the current loading code before you use it.

Check every vector

Whatever runtime you use, check vectors before you store them. A NaN vector looks like any other row in a database until search results go wrong.

"""Store how an index was built, and refuse to mix vectors from another setup."""
import json
from pathlib import Path

import numpy as np

SETUP = {
    "model": "ggml-org/embeddinggemma-2-GGUF:Q8_0",
    "query_prefix": "task: search result | query: ",
    "doc_prefix": "title: none | text: ",
    "dim": 256,
    "normalized": True,
}


def check_vectors(vectors: np.ndarray, dim: int) -> None:
    if not np.all(np.isfinite(vectors)):
        raise ValueError("vectors contain NaN or inf (float16 problem?)")
    if vectors.shape[1] != dim:
        raise ValueError(f"expected {dim} dimensions, got {vectors.shape[1]}")
    norms = np.linalg.norm(vectors, axis=1)
    if not np.allclose(norms, 1.0, atol=1e-3):
        raise ValueError(f"vectors are not unit length (min {norms.min():.3f}); renormalize after truncating")


def save_index(folder: str, vectors: np.ndarray, ids: list[str]) -> None:
    check_vectors(vectors, SETUP["dim"])
    path = Path(folder)
    path.mkdir(parents=True, exist_ok=True)
    np.save(path / "vectors.npy", vectors)
    (path / "index.json").write_text(json.dumps({"setup": SETUP, "ids": ids}, indent=2))


def load_index(folder: str):
    path = Path(folder)
    meta = json.loads((path / "index.json").read_text())
    if meta["setup"] != SETUP:
        raise ValueError(f"index was built with {meta['setup']}, app uses {SETUP}. Re-embed everything.")
    return np.load(path / "vectors.npy"), meta["ids"]

On the test box, check_vectors caught both problems:

caught: vectors are not unit length (min 0.539); renormalize after truncating
caught: vectors contain NaN or inf (float16 problem?)

Lesson 03 uses the rest of this file: the stored setup and load_index.

Sources

• Model card: google/embeddinggemma-2