DenseOn produces a non-finite FP16 embedding for a long-tail query from the Amazon ESCI dataset
#2
by jshree - opened
DenseOn returns non-finite FP16 embeddings for a valid long-tail query
Running DenseOn over 25,000 unique English e-commerce queries (Amazon ESCI, US
locale), exactly one query produces a non-finite CLS output in FP16. The values
are already non-finite in last_hidden_state, before any pooling or L2
normalization, so it is not a normalization artifact. BF16 is stable on the same
input, and switching between eager and SDPA attention makes no difference.
Environment
| GPU | NVIDIA L4 (compute capability 8.9, 23,034 MiB) |
| Driver | 580.82.07 (CUDA 13.0 runtime available) |
| PyTorch | 2.11.0+cu128 (CUDA toolkit 12.8) |
| transformers | 5.5.3 |
| sentence-transformers | 5.4.1 |
| Platform | Google Colab |
| Batch size | 1, no truncation (39 tokens) |
Results
| dtype | attention | finite | max abs |
|---|---|---|---|
| float16 | eager | no | – |
| float16 | sdpa | no | – |
| bfloat16 | sdpa | yes | 4.40625 |
Reproduction
import torch
from transformers import AutoModel, AutoTokenizer
MODEL_ID = "lightonai/DenseOn"
QUERY = 'query: Not sure. Maybe a a standard size (15 or 17??) screen middle of the range laptop that is either black or purple. It might be a Dell or Acer'
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
features = tokenizer(QUERY, return_tensors="pt", return_token_type_ids=False).to("cuda")
for dtype, attention in [(torch.float16, "eager"), (torch.float16, "sdpa"), (torch.bfloat16, "sdpa")]:
model = AutoModel.from_pretrained(MODEL_ID, dtype=dtype, attn_implementation=attention).to("cuda").eval()
with torch.inference_mode():
cls = model(**features).last_hidden_state[:, 0]
print(dtype, attention, bool(torch.isfinite(cls).all()))
del model
torch.cuda.empty_cache()
Hey @jshree , i recommend you loading the model with SentenceTransformers directly, as indicated on the model card.