Migrate transformers fusions to onnx-ir + onnxscript rewriter (QuickGelu, BiasAdd, GroupNorm)#29662
Open
justinchuby wants to merge 2 commits into
Open
Migrate transformers fusions to onnx-ir + onnxscript rewriter (QuickGelu, BiasAdd, GroupNorm)#29662justinchuby wants to merge 2 commits into
justinchuby wants to merge 2 commits into
Conversation
Start migrating the transformers-optimizer graph fusions off the proto-based
OnnxModel helper API onto onnx-ir + onnxscript.rewriter. Each fusion becomes a
declarative pattern()/check()/rewrite() rule instead of imperative ModelProto
surgery, reusing the shared rewriter's subgraph-safety and traversal handling.
This first, self-contained batch (new onnx_ir_fusions package) covers fusions
not already implemented by other onnx-ir rewriters:
* QuickGelu (replaces fusion_quickgelu.py) -> com.microsoft.QuickGelu
* BiasAdd (replaces fusion_bias_add.py) -> com.microsoft.BiasAdd
* GroupNorm (replaces fusion_group_norm.py) -> com.microsoft.GroupNorm
(channels-last transposes + optional swish activation)
The existing proto fusions are left in place; this package is additive until the
IR path is wired into the optimizer driver in follow-up PRs. Includes unit tests
(positive + negative cases) and a README describing the migration approach, and
adds onnxscript / onnx-ir to the transformers requirements.
All 8 unit tests pass; ruff check + format clean.
Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
| import unittest | ||
|
|
||
| import numpy as np | ||
| import onnx |
| import unittest | ||
|
|
||
| import numpy as np | ||
| import onnx |
| import unittest | ||
|
|
||
| import numpy as np | ||
| import onnx |
…rm, SkipSimplifiedLayerNorm) Copy the already-implemented onnx-ir / onnxscript.rewriter fusion rules into the onnx_ir_fusions package rather than re-deriving them, extending the batch that migrates the transformers optimizer off the proto OnnxModel API: * bias_gelu_rules() Add + Gelu(tanh) -> com.microsoft.BiasGelu * gelu_fusion_rules() Erf/tanh Gelu subgraphs -> onnx Gelu (exact + tanh) * layer_norm_fusion_rules() mean/var normalization -> LayerNormalization * skip_layer_norm_rules() Add + LayerNorm -> com.microsoft.SkipLayerNormalization * skip_norm_rules() Add + RMSNorm -> com.microsoft.SkipSimplifiedLayerNormalization Each ported rule gets standalone unit tests that build the matched subgraph, apply the rule and assert the fused op (plus negative cases where relevant). Exported from __init__ and added to all_rules(); README updated to distinguish newly authored vs ported fusions. All 17 package unit tests pass; ruff check + format clean. Signed-off-by: Justin Chu <justinchuby@users.noreply.github.com>
| import unittest | ||
|
|
||
| import numpy as np | ||
| import onnx |
| import math | ||
| import unittest | ||
|
|
||
| import onnx |
| import unittest | ||
|
|
||
| import numpy as np | ||
| import onnx |
| import unittest | ||
|
|
||
| import numpy as np | ||
| import onnx |
| import unittest | ||
|
|
||
| import numpy as np | ||
| import onnx |
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.
Motivation
The transformers optimizer's graph fusions are built on the proto-based
OnnxModelhelper API (onnx_model.py) — imperativeModelProtosurgery viamatch_parent_path/is_safe_to_fuse_nodes/nodes_to_add/nodes_to_remove.This PR begins migrating those fusions onto
onnx-ironnxscript.rewriter, where eachfusion is a small declarative
pattern()/check()/rewrite()rule thatreuses the shared, well-tested rewriter (subgraph-safety, topological handling,
subgraph traversal) instead of re-deriving it per fusion.
Scope
This is deliberately an incremental first step, not a wholesale rewrite —
OnnxModelis used across ~80 files / 16k+ LOC and migrating it all at oncewould be unreviewable and regression-prone. This PR adds a new, additive
onnx_ir_fusionspackage with a first batch of self-contained fusions thatare not already covered by other onnx-ir rewriters. The existing proto
fusions are left untouched; the IR path will be wired into the optimizer driver
in follow-ups.
quick_gelu_rules()fusion_quickgelu.pycom.microsoft.QuickGelubias_add_rules()fusion_bias_add.pycom.microsoft.BiasAddgroup_norm_rules()fusion_group_norm.pycom.microsoft.GroupNorm(channels-last transposes + optional swish activation)What's included
onnx_ir_fusions/package:_quick_gelu.py,_bias_add.py,_group_norm.py,shared
_common.pyhelpers,__init__.py(per-fusion factories +all_rules()),and a
README.mddescribing the migration approach and how to add a fusion._*_test.py) with positive and negative cases(e.g. QuickGelu rejects a non-1.702 multiplier; BiasAdd rejects a non-1-D bias;
GroupNorm rejects a non-identity InstanceNormalization).
onnxscript/onnx-irtotransformers/requirements.txt.Usage
Testing
All 8 unit tests pass (
python -m unittest onnx_ir_fusions._quick_gelu_test onnx_ir_fusions._bias_add_test onnx_ir_fusions._group_norm_test).ruff checkandruff formatare clean.Follow-ups
Subsequent PRs will migrate additional non-duplicated fusions (e.g. BiasSplitGelu,
SkipGroupNorm, GemmFastGelu) using the same pattern, and then wire the IR rule
sets into the optimizer driver.