Should I be getting more speedup/memory reduction from FlashAttention2 with Mistral?
Author: cassianlewis
Pull request: no

### System Info

transformers: 4.35.0
python: 3.9.13

### Who can help?

@SunMarc 
@younesbelkada 
@gant

### 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

## Setup model 
```
model_id = "mistralai/Mistral-7B-Instruct-v0.1"

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_use_double_quant=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_compute_dtype=torch.bfloat16
)

# load base LLM model and tokenizer
model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float16, 
                                              quantization_config=bnb_config, use_flash_attention_2 = True)
```

## Run code for different batch sizes
```
results = []
for n in range(1, 25):
    
    print(f'Processing {n} examples')
    tokenized_prompt = tokenizer([context]*n, return_tensors="pt")
    
    
    length = len(tokenized_prompt['input_ids'][0])+1
    print(length)
    t0 = time()
    with torch.no_grad():
        output = model.generate(
            inputs = tokenized_prompt['input_ids'],
            max_new_tokens = 400,
            repetition_penalty = 1.2
        )
    t1 = time()

    time_taken = t1 - t0
    mem_usage = memory()  
    new_token_length = len(output[0]) - length
    tokens_per_second = new_token_length * n / time_taken
    time_per_batch = time_taken/n

    print('Time taken = ', time_taken)
    print(f'Tokens/s = {tokens_per_second}')

    gc.collect()
    torch.cuda.empty_cache()

    results.append({'batch_size': n, 'time_taken': time_taken, 
                    'tokens_per_second': tokens_per_second, 'memory_usage': mem_usage, 'time_per_batch':time_per_batch})
```

### Expected behavior

## Results
Very little speedup/memory improvement:
![flash](https://github.com/huggingface/transformers/assets/131266258/2b722a0b-67d4-4a58-be21-a8eab9cc2f09)

### Profiling
With FA2:
<img width="1255" alt="Screenshot 2023-11-06 at 18 22 46" src="https://github.com/huggingface/transformers/assets/131266258/ab75b997-9225-495f-9e9d-f86162039edc">

Without FA2
<img width="1248" alt="Screenshot 2023-11-06 at 18 16 58" src="https://github.com/huggingface/transformers/assets/131266258/96a84c7a-2f37-48f7-8682-bf256ab2490a">

Would expect better performance given these


Comments:
Hi @cassianlewis 
Thanks a lot for the extensive benchmark here! 
Many things in your setup can add an overhead here, I think we should test it without `use_double_quant=True` and make sure to use the same dtype for `compute_dtype` and `torch_dtype`. For better efficiency, I would try with `torch_dtype=torch.float16` and `bnb_4bit_compute_dtype=torch.float16`. Moreover, note that the input context length has a big impact in the final latency. Please take a look at the experiments I made here: https://github.com/huggingface/transformers/pull/26464#issuecomment-1743273513 and let me know if you have more questions.

Hi @younesbelkada, thanks for the reply.
I did actually do testing with fp16 and the results were largely the same:
```
bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_compute_dtype=torch.float16
)

# load base LLM model and tokenizer
model = AutoModelForCausalLM.from_pretrained(model_id, cache_dir = 'm', torch_dtype=torch.float16, 
                                               quantization_config=bnb_config, use_flash_attention_2 = True)
```

### Results
Slight improvements in time/memory. FYI the input is ~2k tokens (but no padding) so this may benefit from longer input sequences.
![flash_fp16](https://github.com/huggingface/transformers/assets/131266258/30da13b7-242c-4212-a4f8-09b8dae80d6e)

I then tested it with `max_new_tokens = 1` (ie just looking at prefill) and saw a much larger discrepancy:

![flash_s](https://github.com/huggingface/transformers/assets/131266258/84ae9be9-faa6-416f-90e0-e323a55634e6)

This is more in line with what you posted in https://github.com/huggingface/transformers/pull/26464#issuecomment-1743273513

So it looks like, for decoding at least, the speedups are fairly minimal for the kind of input sequence lengths/batch sizes I'm using.



Thanks a lot for this benchmark ! Yes I think this result is pretty much inline with my findings. 
Flash Attention is great for prefill indeed (which is proven by one of your plots), although for generation it is not 'as fast' as prefill you can note that you can fit larger batch size with the same memory budget (~10 for native vs ~16 for FA2). I think FA-2 + HF transformers really shines in the context of training / fine-tuning because you can fit much larger sequence length / batch size, hence improve efficiency of your training setup.
If we want to increase generation throughput one needs to use static KV cache + flash decoding