Transformer in FINN: Scaled Dot-Product Attention #13
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
Currently this is not a HLSCustomOp, but a QONNX CustomOp. Implemented are first operator attributes, ONNX graph/model construction and a rather improvised python mode node execution for debugging.
This causes the C++ simulation to fail as multithreshold activations are not implemented on the HLS side yet.
Note: The threshold parameters are generated and included but not connected to the attention operator yet. The attention operator uses uninitialized thresholds of the same type and shape.
Note: Currently there is no method for optimizing the accumulator width of both, the HLSCustomOp and the python simulation. Thus, to make the tests pass, both must be specified manually to the maximum possible accumulator bitwidth. Doing the MinimizeAccumulatorWidth transform would cause the HLS and python operator behavior to diverge.
These are the new Infer* transformations for detecting the scaled dot-product attention and multi-head attention patterns to convert these to the hardware custom operations, as well as the InferReplicateStream transformation, corresponding to the new ReplicateStream operation. This also adds a few reordering transformations for optimizing the attention pattern and eventually unrolling the attention heads within the ONNX graph to achieve a fully-parallel hardware implementation of multi-head attention, which might look like streamline transformations, but must be applied when/after converting the attention pattern to hardware operators.
Note: This is currently not controlling the memory used by the internal threshold operations and also not controlling the resoruce type used for implementing the floating-point operations within the softmax. These are all still handled by the tools' automatic strategy.
This is a temporary solution to get at least node-by-node RTL simulation of models working by simply skipping the attention operator.
The inferred shape is not taken from the model graph but from the node attributes specifying the shape.
Instead of manually squeezing all shapes, explicit Squeeze and Unsqueeze operations are inserted into the graph before deleting and redoing all shape annotations from scratch. This should be more robust and keeps the interface (data layout) the model exposes to the outside. Wraps Im2Col operations in Unsqueeze-Squeeze operators to shield it from squeezing as Im2Col always operates on 4-dimensional layouts.
iksnagreb
requested review from
DeepCowProductions,
fpjentzsch and
bwintermann
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.
Adds support for multi-head scaled dot-product attention, i.e., the core operation of a Transformer, to FINN. This includes compiler integration of hardware operators for the attention mechanism and multi-head splitting/and merging as well as related graph transformations. Heavily depends on the related streamlining of scaled dot product attention: #12
fetch-repos.sh, see https://github.com/iksnagreb/attention-hlslibTo support a complete Transformer, the following PRs must be merged:
WIP: Merge branch for testing the integration of all the Transformer related PRs until they are fully merged into dev: https://github.com/eki-project/finn-plus/tree/transformer