Gemma3 CPU export: replace fused SkipSimplifiedLayerNormalization with explicit Add + RMSNorm path - #2296
Draft
kunal-vaishnavi with Copilot wants to merge 10 commits into
Draft
Gemma3 CPU export: replace fused SkipSimplifiedLayerNormalization with explicit Add + RMSNorm path#2296kunal-vaishnavi with Copilot wants to merge 10 commits into
kunal-vaishnavi with Copilot wants to merge 10 commits into
Conversation
|
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
…ndow layer assignment
…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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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)
Gemma2Model.make_layernorm_op, CPU handling forSkipSimplifiedLayerNormalizationis now emitted as:Add(residual combine,output_3)SimplifiedLayerNormalization(normalized path,output_0)com.microsoft::SkipSimplifiedLayerNormalizationop unchanged.Scope and compatibility
src/python/py/models/builders/gemma.py).Regression coverage
test/python/builder/test_gemma.py) to lock in:Add + SimplifiedLayerNormalizationSkipSimplifiedLayerNormalization