onnxruntime
46126133 - Add pre-layer normalization support to attention fusion (#27418)

Commit
144 days ago
Add pre-layer normalization support to attention fusion (#27418) ### Description Add support for pre-layer normalization (pre-LN) transformer architectures in the Python attention fusion optimizer. ### Motivation and Context Fixes #11684 Pre-LN models (used in GPT-3, ViT variants, and many modern architectures) apply LayerNormalization **before** attention rather than after. The first block of a pre-LN model has no `Add` node before its first `LayerNormalization` — its input comes directly from a graph input. This caused `FusionAttention.fuse()` to bail out early because it assumed every `LayerNormalization` anchor has an `Add` parent (the residual connection from the previous block). This PR makes four surgical changes to `fusion_attention.py` so that pre-LN first-block models fuse correctly, while preserving all existing post-LN behavior: 1. **Allow LN with graph-input parent** — instead of returning early when no `Add` parent is found, check whether the input is a graph input and continue 2. **Include graph inputs in residual collection** — the `other_inputs` loop previously skipped anything not in `output_name_to_node`; graph inputs are now recognized 3. **Extend child-LN resolution to SkipLN anchors** — after `fuse_skip_layer_norm()` runs, the anchor becomes `SkipLayerNormalization`; the redirect from `root_input` (graph input) to the first LN's output now fires for SkipLN anchors too 4. **Guard `output_name_to_node` lookup** — graph inputs are not node outputs, so the dictionary access is now guarded ### Changes - `onnxruntime/python/tools/transformers/fusion_attention.py` — 4 targeted edits to `FusionAttention.fuse()` - `onnxruntime/test/python/transformers/bert_model_generator.py` — new `create_bert_attention_pre_ln()` test graph generator - `onnxruntime/test/python/transformers/test_attention_fusion.py` — new `test_attention_fusion_pre_ln()` test ### Test Plan - [x] New unit test `test_attention_fusion_pre_ln` passes — verifies `Attention` fused op appears in the optimized graph - [x] Lintrunner passes on all changed files (no lint issues) - [x] Changes are minimal and scoped to the pre-LN first-block gap
Author
Parents
Loading