microsoft / microsoft/onnxruntime-genai
Gemma3ForCausalLM ONNX export produces garbage output for certain prompts (Persian, and any prompt triggering Gemma3's known "massive activation" pattern) - divergence traced to layer 0 (sliding-attention) output
- Dominant language
- C++
- Stars
- 1.1k
- Forks
- 354
- Avg merge
- 2d 16h
- Merged PRs (30d)
- 85
Description
## Environment
- `onnxruntime-genai` version: 0.14.1 (pip)
- `onnxruntime` version: (installed alongside, CPU execution provider)
- `transformers` version: 4.57.6
- Platform: Linux, CPU execution provider
- Model: `google/gemma-3-4b-it`, text-only submodule extracted from the multimodal
`Gemma3ForConditionalGeneration` checkpoint into a standalone `Gemma3ForCausalLM`
checkpoint (config `architectures: ["Gemma3ForCausalLM"]`, `model_type: "gemma3_text"`).
Extraction verified correct: `model.language_model` state_dict loads into the new
`Gemma3ForCausalLM.model` with zero missing/unexpected keys, and PyTorch inference
on this extracted checkpoint (CPU, fp32) produces fully correct output.
- Builder invocation:
```
python -m onnxruntime_genai.models.builder \
-i ./merged-model-text-only \
-o ./onnx-model-fp32-diagnostic \
-p fp32 \
-e cpu
```
(Also reproduced with `-p int4 -e cpu --extra_options int4_accuracy_level=4 int4_algo_config=rtn int4_block_size=32`
— same failure signature, ruling out quantization as the cause.)
## Summary
For certain prompts, the ONNX-exported model produces incoherent/repetitive garbage
output (e.g. endless "Hello" repetition, or a single replacement character) from the
very first generated token, while the identical checkpoint run directly through
`transformers` (PyTorch, CPU, fp32) produces fully correct output for the same prompt.
This reproduces on **both**:
- a Persian LoRA fine-tune merged checkpoint (`mshojaei77/gemma-3-4b-persian-v0`), and
- the plain, untouched `google/gemma-3-4b-it` base model,
put through the identical extraction+export pipeline, which rules out the checkpoint
itself as the cause.
## Failure pattern / prompt matrix
Tested with greedy decoding (`do_sample=False`) on both an fp32 and an int4/RTN/CPU
export, crossing 5 prompts designed to isolate token-ID magnitude vs. script/content:
| prompt | fp32 (base google) | int4 (Persian merged) |
|---|---|---|
| "What is the capital of France?" (mostly low token ids) | correct | correct |
| "The quixotic archaeologist discovered peculiar hieroglyphic anomalies." (high ids, no Persian) | partially garbled | partially garbled |
| "你好! What is the capital of France?" (starts with a high-id non-Latin token) | correct | correct |
| "من است یک دو سه نان" (Persian, but only low token ids, max_id=43865) | degenerate (near-empty output) | **broken** ("Hello" looped) |
| "پایتخت ایران کجاست؟" (normal Persian, high ids) | broken (single garbage char) | **broken** ("Hello" looped) |
This rules out a simple index-overflow theory (e.g. Gather/Cast truncating indices
above 65535) — the low-token-id Persian prompt still fails, while the high-token-id
non-Persian prompt succeeds. The failure correlates with Persian-script content, not
with numeric token ID magnitude.
Progressively truncating the Persian prompt down to a **single Persian character**
already produces garbled output — the bug is not sequence-length-dependent (well
under the model's `sliding_window: 1024`).
## Root cause isolation (hidden-state diff)
Compared intermediate tensors between:
- PyTorch (`transformers`, CPU, fp32, `output_hidden_states=True`)
- the raw ONNX graph (`onnxruntime.InferenceSession` directly on `model.onnx`,
bypassing `onnxruntime_genai`, with two extra graph outputs added via `onnx.helper`:
`/model/embed_tokens/Mul/output_0` and `/model/layers.0/pre_feedforward_layernorm/output_3`)
for the single-character prompt `"پ"` (10 tokens after chat templating):
| checkpoint stage | max abs diff | notes |
|---|---|---|
| raw embedding lookup (`Gather`, pre-scale) | `0.000000` | identical |
| embedding after `* sqrt(hidden_size)` scale (input to layer 0) | `0.002056` | negligible fp rounding |
| **output of layer 0** (residual stream feeding layer 1) | **`12582.95`** | PyTorch max activation ≈ `13118.36`, ONNX max ≈ `671.16` at the same position |
| final logits | `50.80` | consequence of the above |
Layer 0 is a `sliding_attention` layer (local RoPE, `rope_local_base_freq=10000`,
`sliding_window_pattern=6`). The huge PyTorch activation value (~13k) is consistent
with Gemma3's well-documented "massive activations" phenomenon (see
[Daniel Han / Unsloth's writeup on infinite activations in Gemma 3](https://x.com/danielhanchen/status/1902396261875249346),
where post-layernorm activations can reach ~800,000). PyTorch handles this fine in
fp32; the ONNX graph computes a *different*, much smaller value at the same tensor,
not NaN/Inf — suggesting a genuine numerical/algorithmic divergence in how the
exported graph's sliding-attention block (`GroupQueryAttention` /
`SimplifiedLayerNormalization` / `SkipSimplifiedLayerNormalization` fused ops)
handles this specific activation regime, rather than a simple overflow.
Note the whole graph's `io_dtype` is `FLOAT` (fp32) in both the `-p fp32` and
`-p int4 -e cpu` builds (confirmed via `set_io_dtype`'s `int4_cpu` branch), and the
builder's `layernorm_attrs["cast"]["use_fp32"]` option is a no-op when `io_dtype`
is already `FLOAT` — so this is not the known fp16-overflow issue; it reproduces
in a fully-fp32 exported graph.
## Repro script (base google/gemma-3-4b-it, no fine-tune needed)
```python
# 1. Extract Gemma3ForCausalLM from Gemma3ForConditionalGeneration (text-only submodule)
import torch
from transformers import AutoModelForImageTextToText, AutoTokenizer, Gemma3ForCausalLM
SRC = "google/gemma-3-4b-it" # or a local snapshot
full_model = AutoModelForImageTextToText.from_pretrained(SRC, dtype=torch.bfloat16)
text_model = Gemma3ForCausalLM(full_model.config.text_config).to(torch.bfloat16)
text_model.model.load_state_dict(full_model.model.language_model.state_dict(), strict=True)
text_model.lm_head.weight.data.copy_(full_model.lm_head.weight.data)
text_model.tie_weights()
text_model.save_pretrained("./gemma3-4b-text-only", safe_serialization=True)
AutoTokenizer.from_pretrained(SRC).save_pretrained("./gemma3-4b-text-only")
```
```bash
# 2. Export fp32 CPU ONNX
python -m onnxruntime_genai.models.builder \
-i ./gemma3-4b-text-only -o ./gemma3-4b-onnx-fp32 -p fp32 -e cpu
```
```python
# 3. Run a Persian prompt through onnxruntime_genai and observe garbage output
import onnxruntime_genai as og
from transformers import AutoTokenizer
model = og.Model("./gemma3-4b-onnx-fp32")
tok = og.Tokenizer(model)
hf_tok = AutoTokenizer.from_pretrained("./gemma3-4b-text-only")
prompt = "پایتخت ایران کجاست؟" # "Where is the capital of Iran?"
templated = hf_tok.apply_chat_template(
[{"role": "user", "content": prompt}], add_generation_prompt=True, tokenize=False
)
ids = tok.encode(templated)
params = og.GeneratorParams(model)
params.set_search_options(max_length=len(ids) + 60, do_sample=False)
gen = og.Generator(model, params)
gen.append_tokens(ids)
out = []
while not gen.is_done():
gen.generate_next_token()
out.append(int(gen.get_next_tokens()[0]))
print(tok.decode(out))
# Expected: a coherent Persian answer ("The capital of Iran is Tehran...")
# Actual: a single garbage/replacement character, or endless repetition
```
## Question for maintainers
Is this a known limitation of the ONNX-exported sliding-attention block for Gemma3
when input activations hit the "massive activation" regime? Is there a recommended
export option (analogous to Unsloth's fp32-for-layernorm-only fix) to work around
this on CPU, or is this an open bug in the fused `GroupQueryAttention` /
`SimplifiedLayerNormalization` CPU kernels for this shape/precision combination?
Happy to provide the exact modified `.onnx` graph (with the two extra debug outputs)
and the raw hidden-state dumps used for the diff above, if useful.
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.