Graph Pass: scaled_dot_product_attention_sliced_q #2418
+328
−1
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
junpeiz
requested review from
TobyRoseman,
aseemw,
junpeiz,
DawerG,
YifanShenSZ,
jakesabathia2 and
cymbalrush
|
Great! Could you also add some concrete numbers in the PR description? For example, the memory and execution time of a model before/after using the slice Q algorithm? |
junpeiz
reviewed
Dec 17, 2024
TobyRoseman
reviewed
Dec 18, 2024
coremltools/converters/mil/mil/passes/defs/scaled_dot_product_attention_sliced_q.py
Show resolved
Hide resolved
jirmasek
force-pushed
the
jim/graph-pass-sdpa-sliced-q
branch
from
62daa78 to
05c5fc3
Compare
junpeiz
approved these changes
Jan 2, 2025
YifanShenSZ
reviewed
Jan 3, 2025
| @@ -44,5 +44,6 @@ | |||
| optimize_state, | |||
| optimize_tensor_operation, | |||
| preprocess, | |||
| scaled_dot_product_attention_sliced_q, | |||
There was a problem hiding this comment.
Shall we name the source file as transformer.py?
Other names here are "category", so a larger categorical name might sound more suitable than a concrete graph pass name. Wdyt?
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.
For longer Q sequence lengths (typically >1024), it's beneficial to calculate the attention by an algorithm (inspired by Lazy Softmax) that is processing Q in chunks. The overall memory usage and execution time (given it's executed concurrently, e.g. on ANE) should be better, and in certain cases when models encounter OOMs for longer sequence lengths, models using this algorithm still work.
This PR implements a new graph pass that can optionally transform the MIL operation
ios18.scaled_dot_product_attentioninto a set of operations calculating the attention by chunks of Q.Parameters of the new graph pass:
min_seq_length(default: 1280) - the original MIL operation will only be transformed if the sequence length of Q is greater than or equal to this value.seq_length_divider(default: 16) - defines the size of chunks (based on:chunk_size = sequence_length / seq_length_divider)Example of performance of Depth-Anything model running on ANE:
execution time: 131.55 ms
memory usage: 169.67 MB
execution time: 86.84 ms
memory usage: 93.34 MB
CI pipeline run: https://gitlab.com/coremltools1/coremltools/-/pipelines/1600785656