Streamlining of Scaled Dot-Product Attention #12
Merged
+194
−52
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
Flips the order of AbsorbSignBiasIntoMultiThreshold and MoveScalarLinearPastInvariants streamlining transforms to prefer absorbing adds into multi-thresholds instead of propagating them downwards. This should prevent accumulation of scalar adds in front of two-input matmuls in scaled dot-product attention operators (they cannot be moved past the matmul operation in that case).
The MoveScalarMulPastMatMul transformation can now handle matmul operations with both inputs preceded by a scalar multiplication. This change is required for streamlining scaled dot-product attention operations, which are essentially two-input matmuls.
Assertions are to restrictive, causing the program to terminate in cases the streamlining simply encounters nodes to which the transforms are not applicable: Just skip those nodes. Only the two transforms currently affecting the streamlining of scaled dot-product attention have been changed.
This is pretty much copy and paste of the existing test case, just replacing the MatMul initializer by a second top-input followed by a scalar Mul.
Folding quantized initializers into add-like nodes did not repsect the order of inputs to the add node correctly. This is fixed by testing for one of the two possible orders and selecting the following indices accordingly. Shape inference following the transformation is fixed by deleting the annotations instead of propagating them incorrectly. Deleting the shape annotations should not hurt, as these are redone by running shape inference after each transformation anyways.
Add is commutative and thus the export does not always generate the initializer as the second input. However, this was always assumed by this transformation, failing via assertion if the inputs were simply ordered differently. The transformation now handles both of the two possible input orderings.
This is required for streamlining packed input projections of multi-head scaled dot-product attention. Adds support for Squeeze and Unsqueeze as well. Skip moving of fork-node producers as this is not handled correctly. However, the same effect can be attained by applying the MoveLinearPastFork transformation first.
Explicitly rejects absorbing into fork-nodes. Previously, this probably would have failed, silently resulting in a wrong model. Not sure whether this happened in any practically relevant models?
This probably is still rather sketchy, but at least it tries to check the data layout annotation. For now seems to be enough for getting the thresholds of multi-head attention right, IF qonnx properly annotates the 3D layouts.
15 tasks
|
@iksnagreb Could you fix the conflicts (especially the one in qonnx_activation_handlers.py) by merging dev into this? |
|
Conflicts should be resolved now |
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.
See Xilinx#901
Merge following PRs first - maybe into this feature branch first to avoid conflicts