Streamlining of Scaled Dot-Product Attention #901
Conversation
|
The " |
|
In addition to the unit tests included here, I have now rebuilt all finn-examples up to the |
|
Reverted this to a draft as we have found more issues while streamlining more realistic dummy models. I will add the issues to the list in #878 and start working on fixes soon. |
|
Might depend on this PR at QONNX fastmachinelearning/qonnx#85, introducing a new cleanup transformation ensuring input order assumptions. |
|
The |
|
We need this one as well: fastmachinelearning/qonnx#89 |
|
We might want to have this one included as well to support python-mode execution of the graph at any step for verification purposes: fastmachinelearning/qonnx#92 |
|
This one is relevant as well to properly streamline attention head splitting/merging: fastmachinelearning/qonnx#107 |
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.
iksnagreb
force-pushed
the
feature/attention-streamline
branch
from
2b638d0 to
a4fc498
Compare
Trying ideas and bug fixes for streamlining the scaled dot-product attention operator. Related to issue/discussion #878
MoveScalarMulPastMatMulfor two-input join-node matmulsAbsorb1BitMulIntoMatMulandAbsorb1BitMulIntoConvtest for the presence of weight initializersInferShapesfails afterFoldTransposeIntoQuantInitMoveScalarAddPastMatMulby preferringAbsorbSignBiasIntoMultiThresholdFoldQuantWeightstransformation currently propagating shapes backwards and maybe generating the inverse of the scale factorAbsorbAddIntoMultiThresholdtransformation assuming input and initializer order which might not always hold trueFix (and include?) the(Seems to be fixed by fixing one of the other issues, was probably caused by faulty rewiring of the graph inMoveLinearPastEltwiseAddtransformation which does not correctly propagate the shapesFoldQuantWeights, transformation seems not to be required anymore, maybe reopen)Suggest Brevitas to change all the quantizers to signed quantizers to be finn compatibleRemoveIdentityOpshandling fork-node producer