DEV Community

Cover image for Three settings that decide whether EmbeddingGemma 2 search works on your GPU
InferenceKey for InferenceKey

Posted on Originally published at inferencekey.com

Three settings that decide whether EmbeddingGemma 2 search works on your GPU

EmbeddingGemma 2 is an embedding model Google DeepMind released on 6 October under Apache 2.0. An embedding model turns an input into a list of numbers (768 here) so that similar meanings get nearby lists, and search becomes "find the nearest list". This one takes text, images, video and audio into the same space, from a 1.5 GB download of 740M parameters.

We served it on a 16 GB NVIDIA RTX 2000 Ada on our own cloud. Loading it was easy. Three settings decided whether the vectors were correct and latency reasonable, so settle them before you write the indexing job.

The number format that fails without an error

GPUs can store numbers in compact 16-bit formats, float16 and bfloat16, to halve memory. Google's model card (the documentation shipped with the weights) allows only bfloat16 or float32. In float16, some values inside the network get too large to represent, and the card says you get NaN or subtly degraded embeddings with no exception. In a search index that means no crash, just worse rankings.

The card's suggested selection fits in one line:

# never float16: NaN or quietly wrong vectors, no exception
dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float32
Enter fullscreen mode Exit fullscreen mode

On a GPU without native bfloat16, the float32 fallback needs twice the memory. All our runs used a safe format, so the float16 failure is Google's warning, not something we reproduced. That run is our next test. Log which dtype your stack actually loads.

Photos were slow because the container saw 64 cores

Our first backend handled images two to five times slower than it should have. A profiler showed about 72 ms of GPU work per image against 0.33 to 0.44 seconds of CPU preprocessing. The container's CPU quota (the share of the host's processors it may use, set through Linux cgroups) was about 7 cores, but PyTorch counted the host's 64 and started 64 threads to fight over 7.

Limiting threads to the quota took a single image from 531 to 111 ms, eight images from 4,333 to 867 ms, and a 6-second video clip from 754 to 307 ms. Text and audio stayed almost flat, so a text-only benchmark never shows this.

Bar chart of model time per request before and after capping CPU threads: images and video drop sharply, text and speech barely change

Our example backend derives the thread count from the cgroup quota. Without a quota, you set it by hand. We verified this on one container only, so time CPU preprocessing and GPU encoding separately, per input type, on yours.

Freed GPU memory stays reserved

The complete model occupied 1.4 GB of GPU memory, or 0.5 GB with only the text part loaded. Single-input requests peaked at 2.6 GB. Batching is where it grows: eight documents of roughly 6,000 tokens each (a token is about a word or part of one) in one request peaked near 10 GB.

The surprise came afterwards. PyTorch caches released GPU memory for reuse, so the process held 11.6 GB until a restart. Calling torch.cuda.empty_cache() after very large requests fixed it: the same batch now leaves the process at 1.6 GB. For bulk indexing, size the GPU for your biggest batch.

Shorter vectors: same speed, some quality loss

Google trained the model with Matryoshka Representation Learning, so the first 512, 256 or 128 numbers of each vector work as a smaller embedding. Truncating was not faster for us (34 to 35 ms per query at every size). Only storage and payload shrink.

The quality cost depends on the input type, per the model card:

Benchmark 768 256 128
MTEB multilingual v2 (text) 61.36 60.41 57.89
MMEB v2 (images and video) 59.01 56.24 45.65
MSEB retrieval (speech) 69.54 66.76 56.71

The card calls 256 close to lossless. At 128, text holds, but image and video search drops by almost a quarter, so the card keeps 128 for text. Re-normalize each vector after cutting it (scale it back to length 1), and embed queries and documents at the same size. Our EmbeddingGemma 2 measurements of VRAM and latency per input type show what each option costs on a 16 GB card.

What the search test does and does not show

Our test index had twelve files and zero text: photos, short videos and spoken recordings. Six English queries each ranked the correct file first. "How to knead bread dough" matched a Spanish recording of a bread recipe. We generated all twelve files, so this confirms the cross-modal pipeline runs end to end, not how the model compares with others. The images came from Sperid Labs' Iris-3B, which we ran in an earlier test of the Iris-3B image model, and the voices from Pocket TTS.

On serving: vLLM, the open-source inference server, did not support this model in version 0.31.0 when we checked, not even for text, so we wrote a small Python backend on sentence-transformers. Google's launch post names vLLM, SGLang, Ollama and llama.cpp as supported tooling, so check current versions first. The complete EmbeddingGemma 2 deployment write-up, with latency and VRAM tables for every input type, covers the request format, start-up times and all five problems we ran into.

Credits

Google DeepMind built EmbeddingGemma 2: the model, its image, video and audio components, the benchmark results and the float16 warning. Sahil Dua and Henrique Schechter Vera wrote the launch post. The truncatable vectors rely on Matryoshka Representation Learning by Kusupati and colleagues. Hugging Face transformers supports the model from version 5.19.0, and sentence-transformers runs it. PyAV decoded our audio and video, and ffmpeg cut the clips. We wrote the backend, took the measurements and fixed the thread count and the memory cache.

The backend code, the SDK example and every number above are in the EmbeddingGemma 2 deployment article on inferencekey.com.

Top comments (0)