fix: handling of default attrs in SimplifiedLayerNormalization + LayerNormalization🐛 #2396

KarelZe · 2025-06-17T03:37:03Z

SkipLayerNormFusion does currently not fuse ops, if stash_type is at default (=1) or epsilon is at default (=1e-5) for LayerNormalization and SimplifiedLayerNormalization

This pr:

fixes handling default attrs in LayerNormalization, SimplifiedLayerNormalization
adds BART encoder as new test model. I added this model as some of the stash types are at default. The model is versatile and can also be used to test other fusions e.g., EmbedLayerNormalization.
allows for commuted inputs.

Closes #2378.

@shubhambhokare1 @justinchuby Could you please review? Any feedback is greatly appreciated.

…rNormalization🐛

codecov · 2025-06-17T04:45:44Z

❌ 5 Tests Failed:

Tests completed	Failed	Passed	Skipped
1383	5	1378	688

View the top 3 failed test(s) by shortest run time

::onnxscript.tools.training_helper

Stack Traces | 0s run time

ImportError while importing test module '.../onnxscript/tools/training_helper.py'.
Hint: make sure your test modules/packages have valid Python names.
Traceback:
.../hostedtoolcache/Python/3.11.13.../x64/lib/python3.11/importlib/__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
onnxscript/tools/training_helper.py:6: in <module>
    from torch.onnx import _OrtBackend, _OrtBackendOptions
E   ImportError: cannot import name '_OrtBackend' from 'torch.onnx' (.../onnxscript/onnxscript/.nox.../test_torch_nightly/lib/python3.11.../torch/onnx/__init__.py)

::onnxscript.tools.transformers_models.llama_test

Stack Traces | 0s run time

ImportError while importing test module '.../tools/transformers_models/llama_test.py'.
Hint: make sure your test modules/packages have valid Python names.
Traceback:
.../hostedtoolcache/Python/3.11.13.../x64/lib/python3.11/importlib/__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
.../tools/transformers_models/llama_test.py:12: in <module>
    import onnxscript.tools.training_helper
onnxscript/tools/training_helper.py:6: in <module>
    from torch.onnx import _OrtBackend, _OrtBackendOptions
E   ImportError: cannot import name '_OrtBackend' from 'torch.onnx' (.../onnxscript/onnxscript/.nox.../test_torch_nightly/lib/python3.11.../torch/onnx/__init__.py)

::onnxscript.tools.transformers_models.mistral_test

Stack Traces | 0s run time

ImportError while importing test module '.../tools/transformers_models/mistral_test.py'.
Hint: make sure your test modules/packages have valid Python names.
Traceback:
.../hostedtoolcache/Python/3.11.13.../x64/lib/python3.11/importlib/__init__.py:126: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
.../tools/transformers_models/mistral_test.py:14: in <module>
    import onnxscript.tools.training_helper
onnxscript/tools/training_helper.py:6: in <module>
    from torch.onnx import _OrtBackend, _OrtBackendOptions
E   ImportError: cannot import name '_OrtBackend' from 'torch.onnx' (.../onnxscript/onnxscript/.nox.../test_torch_nightly/lib/python3.11.../torch/onnx/__init__.py)

To view more test analytics, go to the Test Analytics Dashboard
_{📋 Got 3 mins? Take this short survey to help us improve Test Analytics.}

onnxscript/rewriter/ort_fusions/skip_normalization.py

justinchuby · 2025-06-17T15:57:52Z

@gramalingam

onnxscript/rewriter/ort_fusions/models/_bart_encoder.py

Copilot

Pull Request Overview

This PR fixes how default attributes (epsilon, stash_type) are handled in both LayerNormalization and SimplifiedLayerNormalization fusions, adds a BART encoder model to the fusion tests, and introduces commuted-input support for SkipLayerNormalization rules.

Extract default epsilon from the matched node instead of requiring it in the pattern signature
Add test_bart_encoder to validate fusion with default-attribute cases
Enable commuted-input variants by applying .commute() to fusion rules

Reviewed Changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated 1 comment.

File	Description
skip_normalization_test.py	Added `test_bart_encoder` to cover default-attribute fusions
skip_normalization.py	Refactored patterns to drop default attrs, extract `epsilon` in rewrite, and apply rule commutation

Comments suppressed due to low confidence (2)

onnxscript/rewriter/ort_fusions/skip_normalization_test.py:73

The test uses fuse_skip_layer_normalization(model) but there is no import for that symbol in this file. Please add from onnxscript.rewriter.ort_fusions.skip_normalization import fuse_skip_layer_normalization (or adjust the import path) to ensure the function is available.

        fuse_skip_layer_normalization(model)

onnxscript/rewriter/ort_fusions/skip_normalization.py:231

The new .commute() calls are applied only to the full SkipLayerNormalization rules. To allow commuted inputs for SkipSimplifiedLayerNormalization as well, you should apply .commute() to the simplified-layer ruleset (if defined) or include those here before applying apply_fusion_rules.

skip_layer_normalization_ruleset = pattern.RewriteRuleSet(

onnxscript/rewriter/ort_fusions/skip_normalization.py

justinchuby · 2025-07-03T18:58:49Z

Looks like there are some comments unaddressed

KarelZe · 2025-07-06T15:19:04Z

@justinchuby Sorry, was busy at work. Resolved one issue, for the other I'd need @gramalingam s input. 👍

onnxscript/rewriter/ort_fusions/skip_normalization.py

justinchuby · 2025-07-14T15:54:38Z

@KarelZe feel free to take a look at the final round of comments. After that I think we are ready to merge. Thanks!

KarelZe · 2025-07-15T20:00:59Z

@justinchuby Should be complete now🙌

Just one more question, would you be open to /interested in adding EmbedLayerNormalization fusion to onnxscript? I could maybe contribute some rewrite rules in a separate pr.

justinchuby · 2025-07-15T20:23:14Z

@justinchuby Should be complete now🙌

Just one more question, would you be open to /interested in adding EmbedLayerNormalization fusion to onnxscript? I could maybe contribute some rewrite rules in a separate pr.

Thanks! Would let @gramalingam comment

gramalingam · 2025-07-18T22:57:43Z

@justinchuby Should be complete now🙌

Just one more question, would you be open to /interested in adding EmbedLayerNormalization fusion to onnxscript? I could maybe contribute some rewrite rules in a separate pr.

Sure, that would be great, thanks!

fix: handling of default attrs in SimplifiedLayerNormalization + Laye…

d550d38

…rNormalization🐛

github-project-automation bot added this to ONNX Script Review Board Jun 17, 2025

github-project-automation bot moved this to Todo in ONNX Script Review Board Jun 17, 2025

KarelZe added 4 commits June 17, 2025 09:54

tests: add new BART encoder test model⚡️

b5ffeac

test: fix typos in BART encoder test model⚡️

87122e8

style: rename variable names of bart encoder to lowercase💅

e9c0235

chore: final clean up of skip layer norm fusion

ac155dc

KarelZe commented Jun 17, 2025

View reviewed changes

onnxscript/rewriter/ort_fusions/skip_normalization.py Show resolved Hide resolved

KarelZe marked this pull request as ready for review June 17, 2025 12:13

KarelZe marked this pull request as draft June 17, 2025 12:47

justinchuby requested review from gramalingam and Copilot and removed request for gramalingam June 17, 2025 15:57

github-advanced-security bot found potential problems Jun 17, 2025

View reviewed changes

onnxscript/rewriter/ort_fusions/models/_bart_encoder.py Fixed Show fixed Hide fixed

onnxscript/rewriter/ort_fusions/models/_bart_encoder.py Fixed Show fixed Hide fixed

Copilot AI reviewed Jun 17, 2025

View reviewed changes

onnxscript/rewriter/ort_fusions/skip_normalization.py Outdated Show resolved Hide resolved

KarelZe added 2 commits June 17, 2025 19:28

fix: set allow_other_attributes for layer norm fusion

edad0ca

fix: address copilot review comments

ba4a971

KarelZe commented Jun 17, 2025

View reviewed changes

onnxscript/rewriter/ort_fusions/skip_normalization.py Show resolved Hide resolved

KarelZe marked this pull request as ready for review June 18, 2025 04:14

justinchuby approved these changes Jun 22, 2025

View reviewed changes

justinchuby enabled auto-merge (squash) July 3, 2025 18:57

justinchuby approved these changes Jul 3, 2025

View reviewed changes

justinchuby disabled auto-merge July 3, 2025 18:58

revert: use non-communative version of skip layer norm

3267861

justinchuby assigned gramalingam Jul 9, 2025

titaiwangms self-requested a review July 10, 2025 17:37

gramalingam reviewed Jul 10, 2025

View reviewed changes

onnxscript/rewriter/ort_fusions/skip_normalization.py Outdated Show resolved Hide resolved

gramalingam reviewed Jul 10, 2025

View reviewed changes

onnxscript/rewriter/ort_fusions/skip_normalization.py Show resolved Hide resolved

justinchuby added the module: rewriter label Jul 14, 2025

fix: apply review comments for skip layer normalization✨

b3467ea

KarelZe requested a review from gramalingam July 15, 2025 20:01

Merge branch 'main' into fix-skip-layer-norm

c5a0c11

justinchuby enabled auto-merge (squash) July 16, 2025 15:19

justinchuby approved these changes Jul 16, 2025

View reviewed changes

gramalingam approved these changes Jul 18, 2025

View reviewed changes

justinchuby merged commit 061f62b into microsoft:main Jul 18, 2025
26 of 32 checks passed

github-project-automation bot moved this from Todo to Done in ONNX Script Review Board Jul 18, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

fix: handling of default attrs in SimplifiedLayerNormalization + LayerNormalization🐛 #2396

fix: handling of default attrs in SimplifiedLayerNormalization + LayerNormalization🐛 #2396

Uh oh!

KarelZe commented Jun 17, 2025 •

edited

Loading

Uh oh!

codecov bot commented Jun 17, 2025 •

edited

Loading

Uh oh!

Uh oh!

justinchuby commented Jun 17, 2025

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Uh oh!

Uh oh!

Uh oh!

justinchuby commented Jul 3, 2025

Uh oh!

KarelZe commented Jul 6, 2025 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

justinchuby commented Jul 14, 2025

Uh oh!

KarelZe commented Jul 15, 2025 •

edited

Loading

Uh oh!

justinchuby commented Jul 15, 2025

Uh oh!

gramalingam commented Jul 18, 2025

Uh oh!

Uh oh!

Uh oh!

fix: handling of default attrs in SimplifiedLayerNormalization + LayerNormalization🐛 #2396

fix: handling of default attrs in SimplifiedLayerNormalization + LayerNormalization🐛 #2396

Uh oh!

Conversation

KarelZe commented Jun 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov bot commented Jun 17, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

❌ 5 Tests Failed:

Uh oh!

Uh oh!

justinchuby commented Jun 17, 2025

Uh oh!

Uh oh!

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull Request Overview

Reviewed Changes

Uh oh!

Uh oh!

Uh oh!

justinchuby commented Jul 3, 2025

Uh oh!

KarelZe commented Jul 6, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

justinchuby commented Jul 14, 2025

Uh oh!

KarelZe commented Jul 15, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

justinchuby commented Jul 15, 2025

Uh oh!

gramalingam commented Jul 18, 2025

Uh oh!

Uh oh!

Uh oh!

KarelZe commented Jun 17, 2025 •

edited

Loading

codecov bot commented Jun 17, 2025 •

edited

Loading

KarelZe commented Jul 6, 2025 •

edited

Loading

KarelZe commented Jul 15, 2025 •

edited

Loading