Mistral flash attention 2 is not work, training speed is equal to the original way which not use flash attn
Author: FangxuLiu
Pull request: no

### System Info

transformers==4.36.2
torch==2.0

model = transformers.AutoModelForCausalLM.from_pretrained(script_args.model_path, trust_remote_code=True, use_cache=False, attn_implementation="flash_attention_2", torch_dtype="auto")

I am pretraining Mistral model with deepspeed zero2, when I used flash attention 2, the training speed not improved.
And some log are there:
You are attempting to use Flash Attention 2.0 with a model not initialized on GPU. Make sure to move the model to GPU after initializing it on CPU with `model.to('cuda')`.

so I want to know what should I do? @ArthurZucker @younesbelkada 

### Who can help?

_No response_

### Information

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

### Tasks

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

### Reproduction

model = transformers.AutoModelForCausalLM.from_pretrained(script_args.model_path, trust_remote_code=True, use_cache=False, attn_implementation="flash_attention_2", torch_dtype="auto")



torchrun \
    --nnode 1 \
    --master_port 10000 \
    --nproc_per_node 4 \
    training/train_instruction.py \
    --model_path /mnt/bn/ecom-nas-lfx/mrgt/models/Mistral-7B-v0.1 \
    --train_data /mnt/bn/ecom-nas-lfx/mrgt/data/v12_1/v2code_train.jsonl \
    --output_dir /mnt/bn/ecom-nas-lfx/mrgt/models/mistral-v12-base-4gpu-flash-test \
    --max_length 2048 \
    --evaluation_strategy no \
    --per_device_train_batch_size 1 \
    --gradient_accumulation_steps 1 \
    --learning_rate 1e-5 \
    --weight_decay 0.1 \
    --optim adamw_torch \
    --num_train_epochs 2 \
    --max_steps -1 \
    --lr_scheduler_type cosine \
    --warmup_steps 100 \
    --logging_strategy steps \
    --logging_steps 1 \
    --save_strategy steps \
    --save_steps 2000 \
    --save_total_limit 1 \
    --seed 42 \
    --bf16 True \
    --report_to none \
    --deepspeed config/zero2.json

### Expected behavior

You are attempting to use Flash Attention 2.0 with a model not initialized on GPU. Make sure to move the model to GPU after initializing it on CPU with `model.to('cuda')`.

Comments:
Hi @FangxuLiu 
Hmm this is strange, FA-2 should lead to some speedup https://huggingface.co/docs/transformers/model_doc/mistral#expected-speedups
Can you try on a larger bs / seq_len ? also what is the hardware you are using ? Are you sure you don't train with padding tokens? 
If you use TRL library you can make sure to pack your input sentences using `packing=True` https://huggingface.co/docs/trl/sft_trainer#packing-dataset--constantlengthdataset-

With the same error. @FangxuLiu , have you solved this problem? 

Hi @FangxuLiu could you share your code


with the same problem

sam problem, any feedback,transformer 3.6,  seq lenght 800, 4 a100, deepspeed zero 2, had no any speedup
install command, pip install -U flash-attn

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.

@aaronlyt and all, did you explicitly set `attn_implementation="flash_attention_2"` ? 

@ArthurZucker @FangxuLiu 
Same problem while inferencing.

Tried the mistral 7b model inference with flash attention and without flash attention as below
```
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer
device = "cuda"
model_with_flash = AutoModelForCausalLM.from_pretrained(
     "mistralai/Mistral-7B-v0.1", 
    torch_dtype=torch.float16, 
    attn_implementation="flash_attention_2")
model_with_flash.to(device)

model_without_flash = AutoModelForCausalLM.from_pretrained(
    "mistralai/Mistral-7B-v0.1", 
    torch_dtype=torch.float16, 
)
model_without_flash.to(device)
tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")
prompt = "My favourite condiment is"
model_inputs = tokenizer([prompt], return_tensors="pt").to(device)

import time
start = time.time()
generated_ids = model_without_flash.generate(**model_inputs, max_new_tokens=100, do_sample=True)
end = time.time()
print(tokenizer.batch_decode(generated_ids)[0])
print("time taken ",end-start)

start = time.time()
generated_ids = model_with_flash.generate(**model_inputs, max_new_tokens=100, do_sample=True)
end = time.time()
print(tokenizer.batch_decode(generated_ids)[0])
print("time taken ",end-start)
```
and here is the output
```

 My favourite condiment is without doubt, mayonnaise. A bit of an addiction. My parents were a little concerned.

I’m a classic mayo kinda girl. A dash of Tabasco and there’s no stopping me.

There was that phase when I used to put mayo inside my sandwiches (a habit that was broken back in high school when the food tech teacher banned it). I then discovered how wonderful it was when mixed with my pasta (still do) and then
time taken 3.065387487411499


 My favourite condiment is without doubt, mayonnaise. A bit of an addiction. My parents were a little concerned.

I’m a classic mayo kinda girl. A dash of Tabasco and there’s no stopping me.

There was that phase when I used to put mayo inside my sandwiches (a habit that was broken back in high school when the food tech teacher banned it). I then discovered how wonderful it was when mixed with my pasta (still do) and then
time taken 3.07542353211499
```

Flash attention seems not to help improve the inference speed nor decrease the inference latency.
Is there any change in the method I followed?

versions
```
Python 3.8.16
cuda 12.2
A100 GPU
flash-attn==2.5.6
```



Running @pradeepdev-1995 's code on a RTX 4090 (by sequentially evaluating each model), and inference is much slower using flash attention 2:

```

 My favourite condiment is Worcestershire sauce. We call it simply ‘Sos’ in my country, but it’s not just your regular sauce. It’s so much more. It can enhance a dish as the secret ingredient, or even stand alone as a cocktail.

Worcestershire sauce was first introduced in 1837 at Edward Blake’s chemist shop. It started with vinegar and then other ingredients were added including spices, lemon peels and anchov
time taken  2.5174171924591064



 My favourite condiment is ketchup, and I'll happily slather it on the most bananas of dishes. However, what the French call a vinaigrette has become an indispensable element of my culinary lexicon as well, at least in its simplest form.

The name itself, literally meaning a "vinegar mixture" is an accurate encapsulation of most vinaigrettes. In simplest form, a vinaigrette is a
time taken  3.0108351707458496

```

I have also run additional tests using Optimum Benchmark on longer sequences, and flash attention 2 is consistently slower for Mistral.

versions:

```
- `transformers` version: 4.39.3
- Platform: Linux-6.5.0-14-generic-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.22.2
- Safetensors version: 0.4.2
- Accelerate version: 0.27.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.2.1+cu121 (True)
```


cc @fxmarty or @younesbelkada would be nice if we can provide the benchmark we run 😉 
cc @ydshieh fa2 generation is important in the benchmark. 

For both @pradeepdev-1995 and @andxalex , we don't know the settings you are using, so we need a small reproducer. Also note that we need to know what was the prompt, the number of generated tokens, the batch size etc.

@ArthurZucker I used the script given by @pradeepdev-1995 in their last comment, slightly modified, see below. Generated 100 tokens on unit batch size with prompt: "My favourite condiment is".

```[python]
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

device = "cuda"

model_without_flash = AutoModelForCausalLM.from_pretrained(
    "mistralai/Mistral-7B-v0.1",
    device_map="auto",
    torch_dtype=torch.float16,
)

tokenizer = AutoTokenizer.from_pretrained("mistralai/Mistral-7B-v0.1")
prompt = "My favourite condiment is"
model_inputs = tokenizer([prompt], return_tensors="pt").to(device)

import time

start = time.time()
generated_ids = model_without_flash.generate(
    **model_inputs, max_new_tokens=100, do_sample=True
)
end = time.time()

del model_without_flash
torch.cuda.empty_cache()

print(tokenizer.batch_decode(generated_ids)[0])
print("time taken ", end - start)

model_with_flash = AutoModelForCausalLM.from_pretrained(
    "mistralai/Mistral-7B-v0.1",
    device_map="auto",
    torch_dtype=torch.float16,
    attn_implementation="flash_attention_2",
)

start = time.time()
generated_ids = model_with_flash.generate(
    **model_inputs, max_new_tokens=100, do_sample=True
)
end = time.time()
print(tokenizer.batch_decode(generated_ids)[0])
print("time taken ", end - start)
```

I've also run for different sequence / batch size combinations up to 8k tokens, inference using FA is consistently slower. Also, probably related: https://github.com/huggingface/transformers/issues/30041



can you share the outputs of `transformers-cli env` ? 🤗 

And I believe the linked issue has the answers! A TLDR here would be better