| --- |
| base_model: google/embeddinggemma-2 |
| base_model_relation: quantized |
| language: |
| - multilingual |
| library_name: mlx |
| license: apache-2.0 |
| pipeline_tag: feature-extraction |
| tags: |
| - mlx |
| - mlx-vlm |
| - embedding |
| - feature-extraction |
| - sentence-similarity |
| - multimodal |
| - image-feature-extraction |
| - audio-feature-extraction |
| - video-feature-extraction |
| --- |
| |
| # embeddinggemma-2-bf16 |
|
|
| The weights and configuration are those of [`mlx-community/embeddinggemma-2-bf16`](https://huggingface.co/mlx-community/embeddinggemma-2-bf16) at revision [`1a4ffdd`](https://huggingface.co/mlx-community/embeddinggemma-2-bf16/tree/1a4ffddb7905d3f63486748deabe091a01fb6201), an MLX conversion of [`google/embeddinggemma-2`](https://huggingface.co/google/embeddinggemma-2) in BF16. Refer to the [original model card](https://huggingface.co/google/embeddinggemma-2) 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](https://github.com/Blaizzy/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: |
|
|
| ```python |
| 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`](https://huggingface.co/google/embeddinggemma-2), created by Google and converted to MLX by [mlx-community](https://huggingface.co/mlx-community). They are unchanged from [`mlx-community/embeddinggemma-2-bf16`](https://huggingface.co/mlx-community/embeddinggemma-2-bf16). |
|
|
| Licensed under [Apache 2.0](https://www.apache.org/licenses/LICENSE-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. |
|
|
