[Gemma3] NaN embeddings on GPU when batching sequences of mixed length (sliding window attention + all-padding windows)

@ArthurZucker

System Info

System Info

transformers: 4.45.1
sentence-transformers: 5.1.2
tokenizers: 0.20.0
safetensors: 0.4.5
PyTorch: ≥ 2.6.0
Serving runtime: pytorch/torchserve-kfs:0.12.0 (Python 3.9, Linux x86_64)
GPU: NVIDIA (CUDA, via KServe / TorchServe on Kubernetes)
CPU inference: not affected

Description

When encoding a batch of sentences of mixed lengths with google/EmbeddingGemma (300M) via SentenceTransformer.encode_query(), short sentences receive an all-NaN embedding vector on GPU. The same inputs encoded individually (batch_size=1) return correct, fully-normalized embeddings.

Trigger condition: the shortest sequence in the batch must be short enough that, when padded to the length of the longest sequence, one or more sliding-window attention windows fall entirely on padding tokens.

In the reported case:

Chunk 0: 142 tokens → padded to 728 tokens → NaN=768/768 in the batch
Chunks 1–3: 452 / 555 / 728 tokens → valid embeddings

The model's sliding window size is 512. With 728 − 142 = 586 padding tokens appended, the first attention windows of Chunk 0 are 100% padding, causing softmax([-inf, -inf, …]) = 0/0 = NaN in the GPU kernel — a silent numerical error (no exception raised, no warning).

Actual behavior: Chunk 0 (shortest, ~142 tokens) returns a vector of 768 NaN values when batched with longer sequences. No exception or warning is raised.

Workaround: detect NaN after encoding and re-encode the affected inputs with batch_size=1.

Who can help?

@ArthurZucker

@Cyrilvallez

Information

The official example scripts
My own modified scripts

Tasks

An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
My own task or dataset (give details below)

Reproduction

Steps to reproduce

from sentence_transformers import SentenceTransformer
import numpy as np

model = SentenceTransformer("google/EmbeddingGemma-300M")  # GPU required

chunks = [
    "ipso literam...",  # ~142 tokens
    "ipso literam...",                      # ~452 tokens
    "ipso literam...",           # ~555 tokens
    "ipso literam...",                # ~728 tokens
]

# Batch — Chunk 0 produces all-NaN on GPU
embs = model.encode_query(chunks, batch_size=64)
for i, emb in enumerate(embs):
    print(f"Chunk {i}: NaN={np.any(np.isnan(emb))}")
# Output:
# Chunk 0: NaN=True   ← 768/768 NaN values
# Chunk 1: NaN=False
# Chunk 2: NaN=False
# Chunk 3: NaN=False

# Individual — all correct
for chunk in chunks:
    emb = model.encode_query([chunk], batch_size=1)
    print(f"NaN={np.any(np.isnan(emb))}")  # False for all

Expected behavior

Expected behavior: all embeddings in the batch are valid (no NaN).

Fix NaN in Gemma3/EmbeddingGemma when batching mixed-length sequences…

I believe the root cause is Sliding window attention + dynamic padding in batched processing → all-padding windows → numerical instability in softmax.

I have added pr fixing the same

@tomaarsen

Honestly, this is not a transformers issue itself but a padding side issue. Regardless of the size of the chunk, having only pad input is invalid in itself (i.e. what values would you expect on all padding input --> NaNs, 0s, ...). So this doesn't need a fix by us but you need to force left padding side instead - not sure why right is the default padding side here tbh

cc @tomaarsen for ST

Many encoders use right padding, and it seems that https://huggingface.co/google/embeddinggemma-300m is not an exception. Left padding is a bit of a newer development from decoders. I'll try and reproduce this myself locally, but if right-padding with sliding window attention breaks if there's more than window_size padding tokens, then that might reasonably be a bug. I'll try and reproduce this in a bit.

Tom Aarsen

@RiccardoTOTI

@RiccardoTOTI embeddinggemma-300m its architecture (gemma3_text) was added in https://github.com/huggingface/transformers/releases/tag/v4.56.0-Embedding-Gemma-preview, but you report being on 4.45.1. I get this error when loading with your exact environment:

...
KeyError: 'gemma3_text'

During handling of the above exception, another exception occurred:
...
ValueError: The checkpoint you are trying to load has model type `gemma3_text` but Transformers does not recognize this architecture. This could be because of an issue with the checkpoint, or because your version of Transformers is out of date.

Could you share a bit more details? Also, on the latest Transformers/Sentence Transformers, I'm unable to get any NaNs even with inputs of the same shape:

from sentence_transformers import SentenceTransformer, __version__ as st_version
from transformers import __version__ as transformers_version
import numpy as np

print(f"ST version: {st_version}")
print(f"Transformers version: {transformers_version}")

model = SentenceTransformer("google/embeddinggemma-300m")  # GPU required

chunks = [
    "and " * 139,  # ~142 tokens
    "and " * 449,  # ~452 tokens
    "and " * 552,  # ~555 tokens
    "and " * 725,  # ~728 tokens
]

inputs = model.preprocess(chunks)
print(inputs["attention_mask"].sum(1))

# Batch — Chunk 0 produces all-NaN on GPU
embs = model.encode_query(chunks, batch_size=64)
for i, emb in enumerate(embs):
    print(f"Chunk {i}: NaN={np.any(np.isnan(emb))}")  # False for all

# Individual — all correct
for chunk in chunks:
    emb = model.encode_query([chunk], batch_size=1)
    print(f"NaN={np.any(np.isnan(emb))}")  # False for all

Tom Aarsen

@RiccardoTOTI embeddinggemma-300m its architecture (gemma3_text) was added in https://github.com/huggingface/transformers/releases/tag/v4.56.0-Embedding-Gemma-preview, but you report being on 4.45.1. I get this error when loading with your exact environment:
...
KeyError: 'gemma3_text'

During handling of the above exception, another exception occurred:
...
ValueError: The checkpoint you are trying to load has model type `gemma3_text` but Transformers does not recognize this architecture. This could be because of an issue with the checkpoint, or because your version of Transformers is out of date.
Could you share a bit more details? Also, on the latest Transformers/Sentence Transformers, I'm unable to get any NaNs even with inputs of the same shape:

from sentence_transformers import SentenceTransformer, version as st_version
from transformers import version as transformers_version
import numpy as np

print(f"ST version: {st_version}")
print(f"Transformers version: {transformers_version}")

model = SentenceTransformer("google/embeddinggemma-300m") # GPU required

chunks = [
"and " * 139, # ~142 tokens
"and " * 449, # ~452 tokens
"and " * 552, # ~555 tokens
"and " * 725, # ~728 tokens
]

inputs = model.preprocess(chunks)
print(inputs["attention_mask"].sum(1))

Batch — Chunk 0 produces all-NaN on GPU

embs = model.encode_query(chunks, batch_size=64)
for i, emb in enumerate(embs):
print(f"Chunk {i}: NaN={np.any(np.isnan(emb))}") # False for all

Individual — all correct

for chunk in chunks:
emb = model.encode_query([chunk], batch_size=1)
print(f"NaN={np.any(np.isnan(emb))}") # False for all
* Tom Aarsen

The pip list | grep transformers output showing 4.45.1 is from the base image before model dependencies are installed. At model load time, requirements.txt upgrades transformers to 4.50.0 into an isolated directory, so the actual runtime version is 4.50.0. The model loads fine on 4.50.0 — gemma3_text is already present there. The NaN occurs in its GPU attention kernel when running on that version

embeddinggemma-300m requires the aforementioned version also for the inclusion of the use_bidirectional_attention parameter: https://huggingface.co/google/embeddinggemma-300m/blob/main/config.json#L57

More details in huggingface/sentence-transformers#3725

Tom Aarsen

embeddinggemma-300m requires the aforementioned version also for the inclusion of the use_bidirectional_attention parameter: https://huggingface.co/google/embeddinggemma-300m/blob/main/config.json#L57

More details in huggingface/sentence-transformers#3725
* Tom Aarsen

Thank you — this is a key clarification. Our config.json does include use_bidirectional_attention: true. If that parameter is ignored on 4.50.0, the model was running with causal attention, which would explain the abnormal behaviour. We will upgrade to transformers ≥ 4.56 and re-test. We may need to revisit this issue entirely after the upgrade — the NaN under bidirectional attention on a correct runtime may or may not reproduce.

Yes, please give it another try with a newer version and let us know! Good luck.

Tom Aarsen

Navigation Menu

Repository navigation

[Gemma3] NaN embeddings on GPU when batching sequences of mixed length (sliding window attention + all-padding windows) #45491

System Info

Who can help?

Information

Tasks

Reproduction

Expected behavior

Batch — Chunk 0 produces all-NaN on GPU

Individual — all correct

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

[Gemma3] NaN embeddings on GPU when batching sequences of mixed length (sliding window attention + all-padding windows) #45491

Description

System Info

Who can help?

Information

Tasks

Reproduction

Expected behavior

Activity

Kabir08 commented on Apr 19, 2026

vasqu commented on Apr 20, 2026

tomaarsen commented on Apr 20, 2026

tomaarsen commented on Apr 20, 2026

RiccardoTOTI commented on Apr 20, 2026

Batch — Chunk 0 produces all-NaN on GPU

Individual — all correct

tomaarsen commented on Apr 20, 2026

RiccardoTOTI commented on Apr 20, 2026

tomaarsen commented on Apr 20, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions