Mistral7B unexpected performance with FlashAttention2
Author: franchukpetro
Pull request: no

### System Info

- `transformers` version: 4.39.3
- Platform: Linux-5.4.0-166-generic-x86_64-with-glibc2.35
- Python version: 3.10.13
- Huggingface_hub version: 0.22.2
- Safetensors version: 0.4.2
- Accelerate version: 0.28.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.2.0+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: Yes, NVIDIA RTX 4090
- Using distributed or parallel set-up in script?: Not sure whether optimum-benchmark runs benchmarks in parallel

### Who can help?

@ArthurZucker, @younesbelkada

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

Hi,

I was trying to define the impact of different attention implementations on Mistral7B model by benchmarking attention, FA1 (SDPA) and FA2 using [optimum-benchmark](https://github.com/huggingface/optimum-benchmark/tree/main).

I used script below, and by changing `attn_implementation` param to "eager"/"sdpa"/"flash_attention_2" I obtained 3 benchmarks on multiple input seqlens.

```python
from optimum_benchmark.backends.pytorch.config import PyTorchConfig
from optimum_benchmark.benchmarks.inference.config import InferenceConfig
from optimum_benchmark.experiment import ExperimentConfig, launch
from optimum_benchmark.launchers.process.config import ProcessConfig
from optimum_benchmark.logging_utils import setup_logging
import torch
import os


def benchmark_mistral7b(input_seqlen=1024):

    setup_logging(level="INFO")
    launcher_config = ProcessConfig(device_isolation=False)

    # batch_size = int(16384 / input_seqlen)
    batch_size = 1 

    print("Batch size: ", batch_size)

    benchmark_config = InferenceConfig(
        memory=True,
        latency=True,
        input_shapes={"batch_size": batch_size, "sequence_length": input_seqlen},
        generate_kwargs={"max_new_tokens": 512, "min_new_tokens": 512},
    )
    backend_config = PyTorchConfig(
        model="mistralai/Mistral-7B-v0.1",
        device="cuda",
        device_ids="0",
        torch_dtype = "float16",
        attn_implementation = "eager",
        no_weights=True,
    )
    experiment_config = ExperimentConfig(
        experiment_name="mistral-bfp16-base",
        benchmark=benchmark_config,
        launcher=launcher_config,
        backend=backend_config,
    )
    benchmark_report = launch(experiment_config)
    benchmark_report.to_json(f"/root/flash_attention/orig_attention/mistral7b_f16_orig_attention_seqlen_{input_seqlen}_batch_size_{batch_size}_benchmark_report.json", flat=False)
    experiment_config.to_json(f"/root/flash_attention/orig_attention/mistral7b_f16_orig_attention_seqlen_{input_seqlen}_batch_size_{batch_size}_experiment_config.json", flat=False)



if __name__ == "__main__":
    os.environ["PROCESS_SPECIFIC_VRAM"] = "0"

    benchmark_mistral7b(input_seqlen=8192)
```

### Expected behavior

However, after I gathered results together and visualized them, I noticed that FA2 is not even as fast as FA1, but even slower than original attention implementation (eager). I though that problem could be in the dtype, so I tried to load model in both float16 and bfloat16, but that had no affect on the results.

![image](https://github.com/huggingface/transformers/assets/35275869/459eeab9-aefa-495f-8a53-c3afab36124f)

I wasn't expecting that, since in this [Transformers doc](https://huggingface.co/docs/transformers/model_doc/mistral#expected-speedups) it is stated that FA2 should have ~x2 speed up compared to original attention.

I'm not sure, whether it is problem in Transformers, optimum-benchmark or original FA2 implementation, but since optimum-benchmark uses this framework under the hood I came to you firstly. If you think the problem is on the other side, just let me know.

Thanks in advance for the help!


Comments:
It is known that FA2 is slow for inference due to lack of parallelism: https://pytorch.org/blog/flash-decoding/#benchmarks-on-codellama-34b . It is best for long-sequence training and prefill.

@learning-chip  thanks for the answer!

This blog confused me now even more. I wasn't aware about FlashDecoding as something standalone, it seems to be cool. But they are reffering to it's implementation from FA2 repo, at the same time comparing it with FA2. Is there a way to use different FlashAttentions 2 from the same implementation? Or is that blog just about the comparing older version of Flash Attention 2 with newer one, where Flash Decoding was introduced?

And just to clarify - is FA2 useful only for training, or we can get significant optimization at inference as well, but only with very long input seqlen (> 10-20k)? Then why in [this Transformers doc](https://huggingface.co/docs/transformers/model_doc/mistral#expected-speedups), where **inference time** was benchmarked, FA2 gets boost compared to default attention? I was benchmarking 4k seqlen as well, but haven't seen that x2 improvement, could the reason be in the GPU (4090 vs A100)?

No you should be getting improved performance when generating as well. 
Latest benchmark still show 1.7x speedup when generating, but we did not go to 8000 batchsize. 
cc @fxmarty and @younesbelkada 

@ArthurZucker sorry, my bad, that is typo in visualization - x axis is input sequence length - 256, 512, 1024, 2048, 4096 and 8192 respectively. Batch size for this experiment was 1.

![image](https://github.com/huggingface/transformers/assets/48595927/57694db8-6435-469f-8471-cf8d881ef0d3)
Here is a benchmark we ran with `llama` fairly recently. That is why I am suprised

I've done additional experiments with Gemma7B, and got more or less expected throughput (again, x axis should be input seqlen, not batch size):

![image](https://github.com/huggingface/transformers/assets/35275869/6978c9b6-7ec9-4bb3-bf21-867b0be648f0)

However, for Gemma 7B decode memory consumption isn't affected by FA or FA2, which is also suspicious for me. On contrast, Mistral7B get's some memory consumption decrease from FA or FA2 incorporation:

![image](https://github.com/huggingface/transformers/assets/35275869/863940ff-c72b-46a5-b151-52f7a4d1fd36)

![image](https://github.com/huggingface/transformers/assets/35275869/9c8cc1c4-75fc-418a-8f25-e2a89b0c2ef1)


@IlyasMoutawwakil could that be an issue with optimum-benchmark? Or probably that can be caused by the fact that my benchmarking is done from the docker image with global VRAM consumption monitoring?

@franchukpetro you will get better forward pass perf with FA2 but generally not decoding (becz of the small seq len) unless the batch size is huge (making up for the small seq len). that's why it's usually recomended for training. in FA2.2 [flash attention with kv cache was introduced](https://github.com/Dao-AILab/flash-attention?tab=readme-ov-file#22-optimize-for-inference), which might be interesting to integrate as an attn implementation as well.
mistral/gemma use a lot of diff attn optimizations so it's hard to tell without digging and checking if GQA, MQA, Sliding Window Attention, were used.
For memory just look at allocated device memory instead of pynvml's vram, as vram is only useful to debug OOMs and with frameworks like onnxruntime which don't expose a memory allocator api.
Make sure you consider stdevs as well in your benchmarks because the gemma plot for example seems to be a lucky shot, in which case augmenting duration might make sense.

@IlyasMoutawwakil thanks for the answer!
I've rerun some of my experiments to address some of your suggestions.

1. In the original [transformers documentation page for Mistral](https://huggingface.co/docs/transformers/model_doc/mistral#expected-speedups) "inference time" is mentioned, so I was considering that it is meant real inference, not the training phase. Is that just a typo in documentation? Cause it is reported x2 speed-up for Mistral with batch size=1 and input seqlen=4096. I did the same experiments but get totally opposite results, and the only difference which I see is the GPU used (A100 vs RTX 4090).

2.  I also thought that probably eager attention is not "native" for Mistral, so I rerun experiment with optimum-benchmark without explicitly stating which attention to use with corresponding parameter `attn_implementation `. Please see the plots for memory & throughput below (for memory plot red green and blue lines have ~same values, manually checked that).

![image](https://github.com/huggingface/transformers/assets/35275869/af3a22af-c049-4bb1-912e-f610f1059f81)
![image](https://github.com/huggingface/transformers/assets/35275869/0ad62f1a-5a13-4524-935c-b4f716fe55b8)

3. I didn't really get your point about "considering stdevs", so I would be grateful if you can elaborate it in more details, so I could make my experiments reproducible. In order to check how "lucky" that shot was I re-run Gemma7B benchmark entirely once again, with native attention (not assigning explicit attention), eager, sdpa and fa2. Now, the throughput results are slightly different in terms eager vs sdpa, but still fa2 is best. While for vram not changes - all 4 experiments yield same memory consumption.

![image](https://github.com/huggingface/transformers/assets/35275869/c9df52cf-cb93-45cc-a28e-1c011fe89733)
![image](https://github.com/huggingface/transformers/assets/35275869/56e92c8a-5d6d-4fe5-b6f5-ffbf82cd36df)


4.  I wasn't explicitly mentioning that in 2nd and 3rd point but,  I'm attaching just decode VRAM, since decode allocated has exactly same pattern, but just slightly lower absolute values.

UPD: I re-run Mistral7B benchmark, since it was strange for me that native attention yields different results from the 3 other options, and got next results:

![image](https://github.com/huggingface/transformers/assets/35275869/7dfc0a6b-1274-4302-a009-9a41eddf7ab7)

So i guess plot in previous comment for Mistral throughput was wrong cause of some reason.

For now, I have next conclusions:
- For Mistral7B and Gemma 7B native attention (attention which is used by default, when we don't specify specific implementation via corresponding params) yields same performance as Flash (SDPA) Attention. I assume that those are just the same, and SDPA is used by default (at least for those two models).
- For Mistral7B FA2 works slower than eager attention, while FA (SDPA) works better. For Gemma7B results are opposite - FA2 is better than eager, while SDPA is slower.
- In terms of memory consumption, for Mistral7B FA and FA2 yield better results, while for Gemma7B there is no difference.


Hi @franchukpetro 
thanks for all the benchmarks ! indeed it's strange that the FA2 impl leands to slower results than SDPA, to force-use the eager attention you should use `attn_implementation="eager"`. This might indicate something is off with Mistral FA2 attention class. Will have a look

@franchukpetro thanks for the insightful benchmarks
1. the diff is that you're comparing decoding throughput while the benchmark in the documentation was comparing forward latency (which is used in prefill and during training).
2. I agree that the flash attention v2 implementation is not optimized for inference with language models, since we're mostly interested in decoding with LLMs. You can take a look at the [LLM-Perf](https://huggingface.co/spaces/optimum/llm-perf-leaderboard), on SDPA and FAv2, you'll see how with a batch size of 1, SDPA had better speedups than FAv2 in the decoding plots, FAv2 even had negative speedups on most of the popular architectures. My conclusion was that current FAv2 impl in transformers is just not optimized for decoding 😅.

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](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.