microsoft / microsoft/onnxruntime
EmbeddingGemma-300M fp16 CPU inference regression: NaN output on Linux aarch64 since 1.24.1
- Dominant language
- C++
- Stars
- 21.9k
- Forks
- 4.2k
- Avg merge
- 4d 11h
- Merged PRs (30d)
- 184
Description
## Describe the bug
Starting in **1.24.1** (the release immediately after 1.23.2), running EmbeddingGemma-300M's
public fp16 ONNX export (`onnx-community/embeddinggemma-300m-ONNX`, `onnx/model_fp16.onnx`)
through the CPU execution provider on **Linux aarch64** produces all-`NaN` output for batches
that ran fine on 1.23.2. Confirmed via a full release bisection from 1.19.0 through 1.28.0:
**1.23.2 is the last good version, 1.24.1 is the first bad one**, and it is still present as
of the latest release, **1.30.0**. It does **not** reproduce on macOS aarch64 (Apple Silicon)
with the same onnxruntime version and inputs.
## Root cause (via graph surgery)
I exposed every intermediate tensor in the graph as a model output and ran inference on the
real model with real inputs directly on the affected Linux aarch64 node. The model's residual
stream grows roughly exponentially across its 24 transformer layers — from a magnitude of
~900 at layer 0 to ~50,000–65,000 by layer 22. fp16's max representable value is 65504. At
`/model/layers.22/Add_2` (the residual add after the feed-forward block), two already-huge
fp16 tensors sum to just over that ceiling and overflow to `+Inf`, which propagates to `NaN`
by layer 23.
This overflow is inherent to the fp16 export in general — for sequences padded near the
tokenizer's 512-token max, even 1.23.2 can overflow. What changed between 1.23.2 and 1.24.1
is that some CPU-EP kernel now produces a sub-1% numeric difference that's enough to flip
already-marginal batches (comfortably under 65504 on 1.23.2) to just over it.
## To Reproduce
Self-contained repro using synthetic filler text (no real user data — chosen purely to match
the character-length distribution that triggers this in our production traffic):
```python
import json
from huggingface_hub import hf_hub_download
from tokenizers import Tokenizer
import numpy as np
import onnxruntime as ort
REPO = "onnx-community/embeddinggemma-300m-ONNX"
model_path = hf_hub_download(REPO, "onnx/model_fp16.onnx", subfolder=None)
hf_hub_download(REPO, "onnx/model_fp16.onnx_data", subfolder=None) # external data, same dir
tokenizer_path = hf_hub_download(REPO, "tokenizer.json")
LOREM = ("Lorem ipsum dolor sit amet, consectetur adipiscing elit, sed do eiusmod tempor "
"incididunt ut labore et dolore magna aliqua. Ut enim ad minim veniam, quis nostrud "
"exercitation ullamco laboris nisi ut aliquip ex ea commodo consequat. Duis aute irure "
"dolor in reprehenderit in voluptate velit esse cillum dolore eu fugiat nulla pariatur. "
"Excepteur sint occaecat cupidatat non proident, sunt in culpa qui officia deserunt "
"mollit anim id est laborum. ")
# Character lengths chosen to match a real failing production batch (mix of short and long,
# one long enough to require truncation to the 512-token model max).
LENGTHS = [279, 264, 261, 186, 87, 167, 297, 644, 926, 77, 625, 1164, 4621, 198, 358, 873,
1670, 123, 260, 124]
texts = [(LOREM * (n // len(LOREM) + 2))[:n] for n in LENGTHS]
tokenizer = Tokenizer.from_file(tokenizer_path)
tokenizer.enable_padding(pad_id=0, pad_token="")
tokenizer.enable_truncation(max_length=512)
encs = tokenizer.encode_batch(texts)
input_ids = np.array([e.ids for e in encs], dtype=np.int64)
attention_mask = np.array([e.attention_mask for e in encs], dtype=np.int64)
session = ort.InferenceSession(model_path, providers=["CPUExecutionProvider"])
out = session.run(["sentence_embedding"], {"input_ids": input_ids, "attention_mask": attention_mask})[0]
nan_rows = sum(1 for row in out if np.isnan(row).any())
print(f"onnxruntime {ort.__version__}: {nan_rows}/{len(out)} rows contain NaN")
```
## Expected behavior
`nan_rows` should be `0/20` regardless of onnxruntime version (as it is on 1.23.2).
## Actual behavior
- `onnxruntime==1.23.2` → `0/20` (clean)
- `onnxruntime==1.24.1` through `1.30.0` → `20/20` (all rows NaN)
Verified with the pip `onnxruntime` package on Linux aarch64 (AWS Graviton, Amazon Linux
2023, Python 3.11) and independently via the Java Maven artifact
(`com.microsoft.onnxruntime:onnxruntime`) on the same hardware — same result in both bindings.
## Additional context
- Setting `GraphOptimizationLevel` (all 5 levels, `ORT_DISABLE_ALL` through `ORT_ENABLE_ALL`)
and `ExecutionMode` (`SEQUENTIAL`/`PARALLEL`) in every combination does **not** change the
outcome — the raw, unoptimized graph (`ORT_DISABLE_ALL`) overflows identically to the fully
optimized one, which rules out a graph-fusion pass as the cause and points to the base
kernel arithmetic itself (most likely in an op used heavily in this architecture — RMSNorm,
RoPE/rotary embedding, or the residual `Add`/`MatMul` chain).
- Not reproducible on macOS aarch64 (Apple Silicon) with the identical onnxruntime version,
model file (verified by MD5), and inputs — suggesting the regression is specific to the
Linux aarch64 build's kernel implementation, not a universal numeric change.
- Happy to share the full per-layer intermediate-tensor trace (tensor names, shapes, min/max
values at every layer) if useful for narrowing down the specific op.
## System information
- OS: Amazon Linux 2023, Linux aarch64 (AWS Graviton)
- ONNX Runtime version: 1.24.1 through 1.30.0 (regression); 1.23.2 and earlier (working)
- Python version: 3.11 (pip package); also reproduced via Java bindings on the same hardware
- Model: `onnx-community/embeddinggemma-300m-ONNX`, `onnx/model_fp16.onnx`, revision
`5090578d9565bb06545b4552f76e6bc2c93e4a66`
- Execution provider: CPUExecutionProvider
Contributor guide
Assessment
This issue has not been assessed yet.