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

[ehsanrs2]

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,
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)

# 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")

# 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

# 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.

[hiSandog]

Tracing the first divergence to layer 0 makes a token-ID-based fixture especially valuable: it can avoid tokenizer/font/language variables while preserving the activation pattern. Comparing the pre-attention norm, sliding-window mask, Q/K scaling, and attention output between PyTorch and ONNX for that fixed sequence should isolate whether export or runtime evaluation introduces the error. Keeping one FP32 fixture in CI would also ensure quantization is not blamed for a graph-level mismatch.