xFormers
We recommend xFormers for both inference and training. In our tests, the optimizations performed in the attention blocks allow for both faster speed and reduced memory consumption.
Install xFormers from pip:
pip install xformers
The xFormers pip package requires the latest version of PyTorch. If you need to use a previous version of PyTorch, then we recommend installing xFormers from the source.
After xFormers is installed, you can use enable_xformers_memory_efficient_attention() for faster inference and reduced memory consumption as shown in this section.
According to this issue, xFormers v0.0.16 cannot be used for training (fine-tune or DreamBooth) in some GPUs. If you observe this problem, please install a development version as indicated in the issue comments.