Optimize LayerNormalization for better cache efficiency + SIMD #465
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
robertknight
force-pushed
the
simd-layer-norm
branch
4 times, most recently
from
b9041a6 to
6ed9b72
Compare
Instead of performing each step of normalization on the whole input before moving onto the next, perform the full normalization over each input slice before moving on to the next. This is more cache efficient. Also fuse and vectorize the steps that scale the input to normalize the variance and apply elementwise scales. With these changes the operator is ~2.5-3x faster on x64 assuming the input is already contiguous. The `LayerNormalization` operator specification allows for the `bias` and `scale` values to have any shape that can be broadcast to the input shape. However actual models seen so far always set these shapes to match the normalized axes of the input. Hence this change drops support for other bias/scale input shapes for the time being.
robertknight
force-pushed
the
simd-layer-norm
branch
from
6ed9b72 to
7685b8b
Compare
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 the LayerNormalization implementation by:
biasandscaleinto one vectorized passTested on docvqa this speeds up
LayerNormalizationin the encoder by 2.5-3x.