Skip to content

Gemma3 CPU export: replace fused SkipSimplifiedLayerNormalization with explicit Add + RMSNorm path - #2296

Draft
kunal-vaishnavi with Copilot wants to merge 10 commits into
mainfrom
copilot/fix-onnx-export-garbage-output
Draft

Gemma3 CPU export: replace fused SkipSimplifiedLayerNormalization with explicit Add + RMSNorm path#2296
kunal-vaishnavi with Copilot wants to merge 10 commits into
mainfrom
copilot/fix-onnx-export-garbage-output

Conversation

Copilot AI commented Jul 13, 2026

Copy link
Copy Markdown
Contributor

Gemma3 (Gemma3ForCausalLM) ONNX exports on CPU can diverge from PyTorch for prompts that trigger large early-layer activations (notably in sliding-attention layers), producing degenerate first-token behavior. The divergence was traced to layer-0 residual stream behavior around fused normalization.

  • Root-cause-targeted export change (Gemma2/Gemma3 on CPU)

    • In Gemma2Model.make_layernorm_op, CPU handling for SkipSimplifiedLayerNormalization is now emitted as:
      • Add (residual combine, output_3)
      • SimplifiedLayerNormalization (normalized path, output_0)
    • Non-CPU paths continue to use the fused com.microsoft::SkipSimplifiedLayerNormalization op unchanged.
  • Scope and compatibility

    • Change is isolated to Gemma builder logic (src/python/py/models/builders/gemma.py).
    • No behavior change for non-Gemma architectures or non-CPU execution providers.
  • Regression coverage

    • Added focused builder unit tests (test/python/builder/test_gemma.py) to lock in:
      • CPU emits expanded Add + SimplifiedLayerNormalization
      • non-CPU still emits fused SkipSimplifiedLayerNormalization
# CPU Gemma export path for skip RMSNorm now expands to:
Add(root_input, skip_input) -> output_3
SimplifiedLayerNormalization(output_3, weight) -> output_0

@azure-pipelines

Copy link
Copy Markdown
Azure Pipelines:
There may be pipelines that require an authorized user to comment /azp run to run.

Copilot AI changed the title [WIP] Fix ONNX export for Gemma3ForCausalLM to prevent garbage output Gemma3 CPU export: replace fused SkipSimplifiedLayerNormalization with explicit Add + RMSNorm path Jul 13, 2026
Copilot AI requested a review from kunal-vaishnavi July 13, 2026 23:48
…e output)

Restores the `not q_norm` / `not k_norm` guards for `use_packed_matmul` in
`base.py make_attention_init` that were dropped during the v0.9.0→branch
refactor.  Without these guards, Gemma-3 (q_norm=True, k_norm=True) had
packed matmul enabled, causing `make_qk_norm` to apply layer-norm to the
full packed-QKV tensor (instead of Q alone) and to use an empty string as
the k_path input — producing a malformed ONNX graph and garbage output.

Added `TestPackedMatmulQKNorm` regression tests to prevent recurrence.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

2 participants