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 xformersTip
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 it with [~ModelMixin.set_attention_backend] as shown in the Attention backends guide.
Warning
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.