Add Qwen3 model type support to Python transformer optimizer - #27556
Merged
tianleiwu merged 1 commit intoMar 13, 2026
Merged
Conversation
Register qwen3 in MODEL_TYPES (using Gpt2OnnxModel for RoPE/RMSNorm support) with appropriate fusion option defaults: disable EmbedLayerNorm and use NoMask for this decoder-only architecture. Fix SkipLayerNormalization fusion fallback when symbolic shape inference fails. Previously fusion was skipped entirely; now it proceeds with the safe default skip_index=1 since both Add inputs are already verified as non-initializer dynamic tensors. Add Qwen3 test graph generator and unit test verifying SimplifiedLayerNormalization and SkipSimplifiedLayerNormalization fusion on a synthetic Qwen3 decoder layer.
tianleiwu
approved these changes
Mar 6, 2026
tianleiwu
enabled auto-merge (squash)
March 6, 2026 07:39
Contributor
|
/azp run Linux QNN CI Pipeline, Win_TRT_Minimal_CUDA_Test_CI, Windows ARM64 QNN CI Pipeline, Windows GPU Doc Gen CI Pipeline |
|
Azure Pipelines successfully started running 4 pipeline(s). |
tianleiwu
added a commit
that referenced
this pull request
Mar 13, 2026
…sion.py (#27642) # Description This PR addresses a build error and subsequent test failures related to recent changes in GridSample and the transformer optimizer. Related PRs: #27201, #27556. ## Changes ### 1. Fix GridSample Build Error - Removed an unused local variable `mode_str` in `onnxruntime/core/providers/cuda/tensor/grid_sample.cc` that was causing a warning (treated as error) about shadowing a member variable. - Ref: [`grid_sample.cc`](https://github.com/microsoft/onnxruntime/blob/c979a2407f/onnxruntime/core/providers/cuda/tensor/grid_sample.cc#L54) ### 2. Update GridSample Tests - Updated `onnxruntime/test/providers/cpu/tensor/grid_sample_test_custom.inc` to use default execution providers in `RunTests` instead of a hardcoded opset version, ensuring compatibility across different environments. ### 3. Revert Transformer Fusion Fallback - Reverted a recent change in `onnxruntime/python/tools/transformers/fusion_skiplayernorm.py` that enabled a fallback for `SkipLayerNormalization` fusion when symbolic shape inference fails. - This revert was necessary to avoid regressions in GPT-2 tests where model definitions contain typos that intentionally (or coincidentally) break shape inference. - Ref: [`fusion_skiplayernorm.py`](https://github.com/microsoft/onnxruntime/blob/c979a2407f/onnxruntime/python/tools/transformers/fusion_skiplayernorm.py#L113) ### 4. Restore Transformer Test Parity - Updated `onnxruntime/test/python/transformers/test_attention_fusion.py` specifically `test_qwen3_normalization_fusion` to match the expected node counts after reverting the fusion fallback. - Ref: [`test_attention_fusion.py`](https://github.com/microsoft/onnxruntime/blob/c979a2407f/onnxruntime/test/python/transformers/test_attention_fusion.py#L398) ## Verification - `build_cuda.sh` completed successfully. - `onnxruntime/test/python/transformers/test_attention_fusion.py` passes with "OK". - `lintrunner -a` reports no issues.
tianleiwu
pushed a commit
that referenced
this pull request
Mar 16, 2026
### Description Extend `FusionRotaryEmbeddings` to handle Qwen3's on-the-fly rotary position embedding computation, where cos/sin values are computed from `inv_freq` at runtime instead of being looked up from a pre-computed cache. This is a follow-up to #27556 (Qwen3 basic model type support). Depends on #27556. Part of #25083. ### Motivation and Context Qwen3 models (ranked 4th on MTEB) compute RoPE differently from existing supported models (Phi, LLaMA, etc.). Instead of pre-computing cos/sin caches and looking them up via `Gather(cache, position_ids)`, Qwen3 computes them on-the-fly: ```python freqs = inv_freq_expanded @ position_ids_expanded # MatMul emb = torch.cat((freqs, freqs), dim=-1) # Concat cos = emb.cos() * attention_scaling # Cos, Mul sin = emb.sin() * attention_scaling # Sin, Mul ``` Additionally, TorchScript exports of Qwen3 insert `Cast` nodes in the `rotate_half` pattern (from `torch.floor_divide` tracing), which the existing path patterns don't account for. ### Changes **`fusion_rotary_attention.py`:** - Add Cast-tolerant `rotate_half` path patterns (`rotate_half_x2_path_2_3`, `_2_4`, `rotate_half_x1_path_2_3`, `_2_4`) that allow 1-2 Cast nodes between Unsqueeze and Div in the dynamic Slice index computation - Add `sin_path_5` / `cos_path_5` patterns matching the on-the-fly computation: `MatMul → Transpose → Concat → Cos/Sin → Mul(scaling) → Unsqueeze → Mul`, with optional Cast variant (the optimizer's earlier Cast fusion pass may remove the Cast) - Add `create_cos_sin_cache_from_on_the_fly_rope()` helper that extracts `inv_freq` weights, computes cos/sin caches as model initializers, and traces `position_ids` from the graph - Handle per-layer vs shared node removal correctly (only remove per-layer Unsqueeze/outer Mul nodes; shared MatMul/Cos/Sin nodes are pruned automatically by the optimizer) **`qwen3_model_generator.py`:** - Add `include_rope=True` parameter to `create_qwen3_decoder_layer()` - Generate full on-the-fly RoPE computation graph: `inv_freq` initializer, `position_ids` input, MatMul/Transpose/Concat/Cos/Sin/Mul nodes, and `rotate_half` pattern with dynamic Slice indices (including Cast nodes from floor division) - Apply RoPE to both Q and K paths **`test_attention_fusion.py`:** - Add `test_qwen3_rotary_embedding_fusion` verifying 2 RotaryEmbedding nodes are fused along with 3 SimplifiedLayerNormalization and 1 SkipSimplifiedLayerNormalization ### Verification - **Unit tests**: All 15 `test_attention_fusion.py` tests pass (14 existing + 1 new) - **Real model**: Verified on Qwen3-Embedding-0.6B (28 layers): 56 RotaryEmbedding nodes fused (28 layers × 2 per layer for Q and K), reducing total node count from 7416 → 4661 (37% reduction) - **No regressions**: All changes are additive alternative path patterns — existing models that use dynamic Slice indices or cache-based RoPE never hit the new paths - **Lint**: `lintrunner -a` clean on all modified files
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
2 participants
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.
Description
Add
qwen3to the Python transformer optimizer's model type registry, enabling graph optimization for Qwen3 models (e.g., Qwen3-Embedding-0.6B, ranked 4th on MTEB).Motivation
Fixes #25083
Running
optimum-cli export onnx --optimize O3on Qwen3 models fails with:This PR resolves that by registering the model type and fixing a fusion gap that blocked normalization fusions.
Changes
Model type registration (
optimizer.py):"qwen3": (Gpt2OnnxModel, "pytorch", 0)toMODEL_TYPESGpt2OnnxModel(notBertOnnxModel) because itsfuse_attention()callsFusionRotaryAttention, which searches onSkipSimplifiedLayerNormalizationanchors — needed for RMSNorm-based modelsFusion option defaults (
fusion_options.py):EmbedLayerNormalization(decoder-only, no BERT-style embedding)AttentionMaskFormat.NoMask(causal masking is implicit)SkipLayerNormalization fusion fallback (
fusion_skiplayernorm.py):SkipLayerNormalization/SkipSimplifiedLayerNormalizationfusionsskip_index=1(second Add input is skip), since both inputs are already verified as non-initializer dynamic tensors (lines 88-90)SkipSimplifiedLayerNormalizationfusion on Qwen3 models where shape inference failsTest (
test_attention_fusion.py,qwen3_model_generator.py):SimplifiedLayerNormalization(pre-attn, Q-norm, K-norm) + 1×SkipSimplifiedLayerNormalization(residual + post-attn RMSNorm)Verified on real model: Running the optimizer on an exported Qwen3-Embedding-0.6B (2-layer) reduces nodes from 208 → 150 (28% reduction). All 9 RMSNorm patterns fuse correctly: 5×
SimplifiedLayerNormalization+ 4×SkipSimplifiedLayerNormalization.Scope note: Full RotaryEmbedding + MultiHeadAttention fusion for Qwen3's dynamo-exported graphs requires additional pattern matching work (static Slice indices, on-the-fly sin/cos computation, QK-Norm in Q/K paths, GQA expansion). That will be addressed in a follow-up PR.
Test Plan
test_attention_fusion.py::TestFusion::test_qwen3_normalization_fusionpassestest_attention_fusion.pypass (no regressions)test_optimizer_huggingface_bert.pypass (bert, distillbert, roberta, xlm_roberta — no regressions from the SkipLayerNorm fallback change)lintrunner -aclean