embeddinggemma-2-bf16

The weights and configuration are those of mlx-community/embeddinggemma-2-bf16 at revision 1a4ffdd, an MLX conversion of google/embeddinggemma-2 in BF16. Refer to the original model card for general information on the model.

This repository adds mlx_vlm_embeddinggemma2_fixes.py, which patches mlx-vlm's EmbeddingGemma 2 implementation at runtime to compute the same embeddings with a lower peak memory.

The saving grows with the number of tokens in a batch. It matters when embedding many long documents at once, for example to index a collection. For a single query the saving is negligible.

Supported stacks

Stack	Version	Setup
mlx-vlm	0.7.8	Apply the fixes file before loading the model

Usage

mlx-vlm

Apply mlx_vlm_embeddinggemma2_fixes.py before loading the model:

import sys
from mlx_vlm.utils import get_model_path

model_path = get_model_path("stefanschmidt/embeddinggemma-2-bf16")
sys.path.insert(0, str(model_path))
import mlx_vlm_embeddinggemma2_fixes
mlx_vlm_embeddinggemma2_fixes.apply_mlx_vlm_embeddinggemma2_fixes()

import mlx.core as mx
from transformers import AutoTokenizer
from mlx_vlm.embedding_loader import load_embedding_model

model = load_embedding_model(model_path)
tokenizer = AutoTokenizer.from_pretrained(model_path)
inputs = tokenizer([
    "task: search result | query: Which planet is known as the Red Planet?",
    "title: none | text: Mars is known as the Red Planet.",
    "title: none | text: Venus is often called Earth's twin.",
], padding=True, return_tensors="np")
embeddings = model(**{key: mx.array(value) for key, value in inputs.items()}).text_embeds
print(embeddings[:1] @ embeddings[1:].T)

apply_mlx_vlm_embeddinggemma2_fixes() returns True when the patch is applied. Before patching, it compares the mlx-vlm methods it replaces with the versions it was written for. If any of them differs, it logs a warning and leaves mlx-vlm unchanged.

Measured effect

MacBook Pro, Apple M5 Pro, 48 GB, macOS 26.6.2, mlx-vlm 0.7.8. Each batch holds the first 320 or 640 titles and abstracts of ICLR 2025 papers, padded to the longest document at 542 tokens. Peak memory includes the 1.5 GB of weights.

Batch	Peak memory, upstream	Peak memory, patched	Time, upstream	Time, patched
320 documents	11.0 GB	7.0 GB	4.61 s	4.57 s
640 documents	20.6 GB	12.4 GB	9.31 s	9.24 s

Times are the median of three runs after a warm-up run. Each version ran in a fresh process starting from an idle machine. Embeddings were bit-for-bit identical with and without the patch on the first 64 documents.

Upstream computes the per-layer inputs of all 24 layers before the first layer runs and keeps them until the last layer has run. They take 24 × 512 values × 2 bytes = 24,576 bytes per padded token. The patch computes each layer's 512 values when that layer runs, which saves 23 of those 24 parts.

Limitations

• BF16 weights only. The patch does not handle the packed weights of quantized checkpoints. After it has been applied, embedding with a quantized EmbeddingGemma 2 checkpoint in the same process fails with a shape error in matmul. Loading such a checkpoint is not affected.
• Measured on text only. Image, audio and video inputs pass through the same patched text model, so the same saving and identical output are expected, but neither has been measured.
• mlx-vlm only. Servers such as omlx do not load the fixes file and run the unpatched implementation, with correct results and the higher peak memory.

License

The weights are those of google/embeddinggemma-2, created by Google and converted to MLX by mlx-community. They are unchanged from mlx-community/embeddinggemma-2-bf16.

Licensed under Apache 2.0, the same license as the original. mlx_vlm_embeddinggemma2_fixes.py is not part of the original model and is licensed under Apache 2.0 as well.