Best way to work around xformers backwards K limit on sm80/86/89 devices

[hafriedlander]

XFormers currently can't run backwards on many Sm80/86/89 devices. It runs fine at any size on A100s, but on other devices in the class (including common ones like 3080, 3090, 4090) it only works backwards when dim_heads <= 64.

This is a pretty narrow use case (most people don't run backwards as standard inference doesn't need it, and SD2 and greater have dim_heads <= 64 for all layers), and kind of working around a bug in xformers, so I'm not sure if there should be a specific API to support working around this issue.

If this is appropriate to include in DIffusers, any preference on an API? (My naive though is to add a max_dim_head argument to set_use_memory_efficient_attention_xformers). I'm happy to do a PR if an API is agreed.

[hafriedlander]

Here's some logic I wrote that detects the issue (for Diffusers 0.10.x) - https://github.com/cloneofsimo/lora/blob/bc497114969a629c0cf909c0a5837d1dacb6909f/lora_diffusion/xformers_utils.py

It's similar to the test in https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/cross_attention.py#L120 but deliberately runs backwards, and checks per dimhead.

Any opinion on whether this should be in Diffusers @patil-suraj or @patrickvonplaten ?

[patil-suraj]

As you said, it's a narrow use case, so I'm not sure if we should handle this in diffusers. I think it's okay to let it fail explicitly during backward if it not working. Maybe open an issue in xformers to see if they have any plans for this.

[hafriedlander]

There is already an issue facebookresearch/xformers#563 - TLDR; they aren't planning on fixing it any time soon.

The downside is that Dreambooth and other training will fail on SD1.5 based models with xformers on, in a way that isn't well documented or communicated.

It's relatively easy to detect the error, and you already have a "does xformers work reliably" test in unet. My personal vote would be that the training scripts at least should have a test & warning if user tries to enable xformers when it won't work.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[qslug]

This is (at least partly) fixed now as part of facebookresearch/xformers@615175f.

[github-actions]

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the contributing guidelines are likely to be ignored.

[litaotju]

It's a known issue of xformers library on 30x GPUs, try install the dev version of the xformers.
This works on my 3060.

pip install xformers==0.0.17.dev466

See facebookresearch/xformers#631