Repository navigation
Consideration of Flash Attention in Generative Components #7944
Description
Regarding the comments here: #7715 (comment)
We have removed the flash attention from the generative components merged into the core. However, based on experiments conducted by @dongyang0122, there appears to be a significant difference between using and not using flash attention. We should consider adding this option back. @dongyang0122 will share more detailed comparison results from the experiments.
Metadata
Metadata
Assignees
Labels
No labels
Activity
KumoLiu commentedon Jul 23, 2024
cc @ericspod @virginiafdez
dongyang0122 commentedon Jul 25, 2024
We use the following Python script for comparison. When flash attention is enabled, we are able to train the diffusion model with batch size 1. When flash attention is disabled, the training process will cause out of memory. For the experiments, we use a A100 GPU with 80GB of memory. If flash attention is enabled, 30GB+ of memory is utilized.
verify_training.py.txt
mingxin-zheng commentedon Jul 26, 2024
Thanks @dongyang0122 for the verification.
It seems to me that xformers/flash_attention brought benefits and we should weight including them in the dependency list.
On the other hand, installation of the package could be challenging, because the range of
torch,cuda, and os platform seems narrow: https://github.com/facebookresearch/xformers/issuesIf
monaileaves the installation ofxformersto the user but we keep theoptional_importin the code base, is this acceptable?ericspod commentedon Jul 26, 2024
We could look at using the flash attention in Pytorch as well: https://pytorch.org/blog/pytorch2-2/
KumoLiu commentedon Jul 29, 2024
Thanks to @dongyang0122 for the script. I tried the script and tested it on both the original generative implementation and the PyTorch implementation. The results are shown below. We can see that PyTorch almost achieves the same results as the xformer implementation. Based on this, I believe we can use PyTorch instead.
KumoLiu commentedon Jul 29, 2024
cc @guopengf for vis. Here is the code I use: https://github.com/KumoLiu/MONAI/tree/flash-atten
Flash attention (#7977)