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

Open
#2,277 1 comment 1 reaction 2 assignees Claimed by @kunal-vaishnavi View on GitHub
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.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.