Repository navigation
Keep the output Reshape when fusing MultiHeadAttention into a non-3D output - #3157
Open
Sankur Kundu (Shakhtar-Sankur) wants to merge 1 commit into
Open
Sankur Kundu (Shakhtar-Sankur) wants to merge 1 commit into
Sankur Kundu (Shakhtar-Sankur) wants to merge 1 commit into
Conversation
…output The MHA pattern ends with a Reshape that is assumed to restore (B, S, D), the shape MultiHeadAttention returns. GPT-2 exports reshape the attention output straight to (B*S, D) for the Conv1D Gemm instead, so the fusion replaced a 2D value with a 3D one and optimize_for_ort produced a model ONNX Runtime rejects (Gemm: input 0 expected to have rank 2 but has rank 3). SmolLM hits the same pattern after Reshape folding, but its 2D value feeds a MatMul, which happens to accept the 3D tensor. When the matched output isn't 3D, reshape the fused output back to the original target shape, so consumers see the same shape as before. Only do so for constant target shapes without zeros, since a 0 would copy a dimension of the Reshape's input, which changes from 4D to 3D.
Sankur Kundu (Shakhtar-Sankur)
changed the title
optimize_for_ort turns GPT-2 into a model ONNX Runtime can't load:
Author
|
@microsoft-github-policy-service agree |
This branch has not been deployed
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.
optimize_for_ortturns GPT-2 into a model ONNX Runtime can't load:The
MultiHeadAttentionrule (ort_fusions/mha.py) matches... -> SDPA -> Transpose -> Reshapeand assumes the final Reshape restores(B, S, D), the shapeMultiHeadAttentionreturns (the code has a# TODO: verify Reshapes). GPT-2's export reshapes the attention output straight to(B*S, D)for the Conv1DGemm, so the fusion replaces a 2D value with a 3D one. Applyingoptimize_for_ort's steps one at a time onhf-internal-testing/tiny-random-GPT2LMHeadModelconfirms that every step beforefuse_mha2leaves the model valid and matching, andfuse_mha2breaks it.SmolLM (
test_smollm) hits the same pattern after two Reshapes are folded into one(B*S, D)Reshape, but its value feeds aMatMul, which happens to accept the 3D tensor, so it works by chance.This PR keeps the fusion and, when the matched output isn't 3D, reshapes the
MultiHeadAttentionoutput back to the original target shape, so consumers see exactly the shape they saw before. It only does this when the target shape is a constant without zeros: a0copies a dimension of the Reshape's input, and that input changes from 4D to 3D. Otherwise the rule doesn't match. The past/present variant reshapes only the attention output.Results (CPU, onnxruntime 1.31.1, main at 9112df0)
optimize_for_ortpytest onnxscript/rewriter/ort_fusions: 143 passed on main, 144 passed with this PR (0 failed).test_output_reshaped_to_2d_keeps_reshapeinmha_unit_test.py: fails on main, passes with this PR.ruff checkandruff formatpass on the changed files.