[SDPA vs. xformers] Discussions on benchmarking SDPA and xformers and implications

[sayakpaul]

Recently, we have been receiving issues from users complaining that SDPA leads to OOMs whereas xformers doesn't not.

So, to help identify the root cause of this, I started a simple benchmark to compare the timings of the different efficient implementations of attention provided by SDPA and xformers. I have restricted this to forward-only for now.

My Colab Notebook is available here: https://colab.research.google.com/drive/1Ne0YPY16G2gmr9H1eCa9iB1wG5jLwf-h?usp=sharing

The setup is simple. Let's start by defining some variables:

batch_size = 32
max_sequence_len = 1024
num_heads = 32
embed_dimension = 32

dtype = torch.float16
device = "cuda"

query = torch.rand(
    batch_size, num_heads, max_sequence_len, embed_dimension, device=device, dtype=dtype
)
key = torch.rand(
    batch_size, num_heads, max_sequence_len, embed_dimension, device=device, dtype=dtype
)
value = torch.rand(
    batch_size, num_heads, max_sequence_len, embed_dimension, device=device, dtype=dtype
)

And then I followed this tutorial to perform the benchmark with F.scaled_dot_product_attention(). And I extended that to use xformers.ops.memory_efficient_attention(). Refer to the Colab Notebook for more details.

Here's the gist of the results:

SDPA

[---------------------  ---------------------]
                     |  flash  |  mem  |  math
6 threads: -----------------------------------
      torch.float16  |   2.7   |  3.0  |  7.0 

Times are in milliseconds (ms).

xformers

[-------------------  ------------------]
                     |  cutlass  |  flash
6 threads: ------------------------------
      torch.float16  |    2.6    |   1.7 

Times are in milliseconds (ms).

We see that the flash variant obtains better performance here. This small difference can compound quickly when we're dealing with dee models. Additionally, I don't think we satisfy the required conditions in the Unet2DConditionModel to be able to take advantage of flash attention. More details: #3594 (comment).

The above issue comment confirms that when using SDPA, the dispatcher uses the SDPBackend.EFFICIENT_ATTENTION variant. This is 3.0 as far as timing is concerned.

For cutlass, we have 2.6 in xformers which is still less than the above.

Triton variant of the memory-efficient implementation errors out. See: facebookresearch/xformers#769.

Tagging a few people here to hear their thoughts:

@williamberman @takuma104 @patrickvonplaten @pcuenca

Let me know if anything is unclear.

[patrickvonplaten]

What versions of PyTorch and xformers are you using @sayakpaul ? the difference in xformers flash and SDPA flash seems pretty significant

cc @Chillee is this expected here?

[williamberman]

yes some of the lower resolution sequence lengths have larger inner dims and the flash kernel iirc only supports inner dims of up to 128 so that makes sense, the other kernel is still the same algo + where flash attention isn't supported the sequence lengths are small so I don't think its a big deal we don't get the best kernel. For a batch size of 32, the lowest stable diffusion attention resolution time is on the order of hundreds of micro seconds for the other kernel so like idk that's probable fine lol.

A difference of 1 ms for that large of a batch size might be fine depending on the variance of the measurements. There are 32 attention blocks (however at different resolutions) in stable diffusion for a single forward pass (assuming all have a 1ms difference), that's a difference of 32 ms. For our default number of inference steps of 50, that's a total difference of 1.6 seconds. iirc for uncompiled stable diffusion, 30s for 50 steps is reasonable so that's ~5% of total inference time? Is 5% ok? for inference and finetuning that sounds ok to me to not look at that much?

What's the variance of inference time for both the whole pipeline and the kernel itself? It could very well be that a 1ms difference is w/in the standard amount the inference time varies

Was the original issue that people were seeing OOMs, speed issues or both? If the issues are OOMs, we should probably just be looking at peak memory during kernel execution

[sayakpaul]

What versions of PyTorch and xformers are you using @sayakpaul ?

PyTorch: 2.0.1+cu118, xformers: 0.0.20 (info available from the Colab I linked in the OP)

[sayakpaul]

Currently, obtaining memory usage with and without xformers (when not using xformers defaulting SDPA in PT 2.0) with the DreamBooth LoRA script. Will report back.

Will report the memory usage for the following:

• xformers in PT 1.13.1
• xformers in PT 2.0
• SDPA in PT 2.0

[sayakpaul]

I took the LoRA DreamBooth training script and added logging for memory usage. See here:

diffusers/examples/dreambooth/train_dreambooth_lora.py

Line 1309 in a7d1171

# Printing the GPU memory usage details such as allocated memory, peak memory, and total memory usage

With this change, I set up my benchmarking runs with the following:

export MODEL_NAME="runwayml/stable-diffusion-v1-5"
export INSTANCE_DIR="dog"
export OUTPUT_DIR="lora-trained"

accelerate launch train_dreambooth_lora.py \
  --pretrained_model_name_or_path=$MODEL_NAME  \
  --instance_data_dir=$INSTANCE_DIR \
  --output_dir=$OUTPUT_DIR \
  --mixed_precision="fp16" \
  --instance_prompt="a photo of sks dog" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
  --checkpointing_steps=100 \
  --learning_rate=1e-4 \
  --report_to="wandb" \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=500 \
  --validation_prompt="A photo of sks dog in a bucket" \
  --validation_epochs=50 \
  --seed="0" \
  --push_to_hub

For training with xformers I just enabled the --enable_xformers_memory_efficient_attention flag. Additionally, dog was downloaded like so:

from huggingface_hub import snapshot_download

local_dir = "./dog"
snapshot_download(
    "diffusers/dog-example",
    local_dir=local_dir, repo_type="dataset",
    ignore_patterns=".gitattributes",
)

I tested this setup in the following configurations on two devices: A100 and V100:

• xformers in PT 1.13.1 (xformers == 0.0.16)
• xformers in PT 2.0 (xformers == 0.0.20)
• SDPA in PT 2.0

Here are the memory logs on wandb:

A100 (fp16):

• xformers in PT 2.0 (0.0.20): https://wandb.ai/sayakpaul/dreambooth-lora/runs/sqzas5aa/logs
• xformers in PT 1.13.1 (0.0.16): https://wandb.ai/sayakpaul/dreambooth-lora/runs/kv083ygw/logs
• SDPA: https://wandb.ai/sayakpaul/dreambooth-lora/runs/wcobi2p4/logs

V100 (fp16):

• xformers in PT 2.0 (0.0.20): https://wandb.ai/sayakpaul/dreambooth-lora/runs/o8ggqbe7/logs
• xformers in PT 1.13.1 (0.0.16): https://wandb.ai/sayakpaul/dreambooth-lora/runs/qfbjxyrg/logs
• SDPA: https://wandb.ai/sayakpaul/dreambooth-lora/runs/jfmwnyp8/logs

If you check the last GPU memory occupancy from the V100 logs for each of the settings, you'd notice xformers in PT 2.0 and SDPA are similar. FWIW, xformers in PT 1.13.1 seems to have more memory occupancy.

Let me know if I am looking at it in the right way. @pcuenca @williamberman

[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.