Support (Bias)SkipLayerNormalization fusion in GPT2 #13988
Merged
Conversation
This file contains 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
…crosoft/onnxruntime into hari/skip_layer_normalization
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
hariharans29
commented
Dec 20, 2022
hariharans29
changed the title
hariharans29
commented
Dec 20, 2022
tianleiwu
requested changes
Dec 20, 2022
| @@ -537,7 +537,14 @@ def fuse_gpt2(self, layernorm, add_before_layernorm, input_name_to_nodes, output | |||
| if two_gather is None: | |||
| return False | |||
|
|
|||
| # If the add_before_layernorm node is an Add node, then the add_output output is the first index | |||
There was a problem hiding this comment.
nit: It is better to also update the comment at the beginning of this function. That comment only contains first case but not the second.
There was a problem hiding this comment.
I donm't see any other comment at the beginnign of this function. Are you referring to the args description of is_embedding_sum_needed() before this by any chance ?
There was a problem hiding this comment.
Streamlined some logic around optional_embedding_sum_output and add_output which both need the comment. It should be easier to understand the code now.
wangyems
reviewed
Dec 20, 2022
tianleiwu
approved these changes
Dec 21, 2022
tianleiwu
added a commit
that referenced
this pull request
Dec 22, 2022
### Description Add support of ONNX conversion of GPT-2 for two stages: * Stage 1 is the initial stage that has empty past state. * Stage 2 has non-empty past state and sequence_length is 1. Add a parameter --stage to specify such stage. For stage 1, we will enable mask_index for Attention so that we can use fused attention in CUDA. Other changes: (1) use int32 inputs as default (otherwise, there is error in inference) (2) update gpt2_parity to include SkipLayerNormalization (see #13988) and EmbedLayerNormalization (3) get all environment variables that might impact GPT-2 latency in benchmark_gpt2 ### Motivation and Context To test fused attention for GPT-2 model for #13953.
henrywu2019
pushed a commit
to henrywu2019/onnxruntime
that referenced
this pull request
Dec 26, 2022
henrywu2019
pushed a commit
to henrywu2019/onnxruntime
that referenced
this pull request
Dec 26, 2022
### Description Add support of ONNX conversion of GPT-2 for two stages: * Stage 1 is the initial stage that has empty past state. * Stage 2 has non-empty past state and sequence_length is 1. Add a parameter --stage to specify such stage. For stage 1, we will enable mask_index for Attention so that we can use fused attention in CUDA. Other changes: (1) use int32 inputs as default (otherwise, there is error in inference) (2) update gpt2_parity to include SkipLayerNormalization (see microsoft#13988) and EmbedLayerNormalization (3) get all environment variables that might impact GPT-2 latency in benchmark_gpt2 ### Motivation and Context To test fused attention for GPT-2 model for microsoft#13953.
wangyems
pushed a commit
that referenced
this pull request
Jan 3, 2023
### Description 1. SkipLayerNormalization has a new output (#13988) and the symbolic shape inference script needs corresponding updates 2. The greedy sampling op (#13426) shouldn't re-use the logits buffer as its corresponding kernel doesn't seem to support it yet. ### Motivation and Context Fix some transformer issues
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.
Description
The GPT2 model has a slightly different SkipLayerNormalization pattern from BERT in that the residual to the next layer is not the output of LayerNorm like in BERT but the added residual and the input fed to the LayerNorm node thus making the Add and LayerNorm to not be fused as there is a consumer of the intermediate output of Add. This also means that any Add after MatMuls feeding into SkipLayerNortmalization won't be fused as well (as the pre-requisite to that fusion is that SkipLayerNormalization be fused first). This change adds support for this variant of SkipLayerNormalization (SLN) by adjusting the schema of SLN to have an optional output which will be the output of the addition of the residual and the input.
TODO: Add kernel test
Motivation and Context
Improve fusion coverage for GPT2 and help improve its perf and any language model using it