Repository: transformers
Issue: Difference in LlamaAttention & LlamaFlashAttention2 attn_output
Labels: WIP

### System Info

- `transformers` version: 4.34.1
- Platform: Linux-5.15.0-86-generic-x86_64-with-glibc2.31
- Python version: 3.11.5
- Huggingface_hub version: 0.17.3
- Safetensors version: 0.4.0
- Accelerate version: 0.23.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.1.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
- Using distributed or parallel set-up in script?: no

### Who can help?

@ArthurZucker and @younesbelkada

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

We notice `LlamaFlashAttention2._flash_attention_forward` returns a different `attn_output` than `LlamaAttention` computes.

`flash_attn_non_determinism.py`:
```python
import argparse

import torch
import torch.backends.cudnn
import transformers
from transformers.models import llama


def main() -> None:
    torch.backends.cudnn.deterministic = True

    parser = argparse.ArgumentParser()
    parser.add_argument("--use-flash-attention-2", action="store_true")
    args = parser.parse_args()
    use_flash_attention_2 = args.use_flash_attention_2

    tokenizer = transformers.AutoTokenizer.from_pretrained(
        "/models/huggingface/meta-llama/llama-2-7b-chat-hf", local_files_only=True, use_safetensors=True, device_map=torch.device("cuda")
    )
    tokenizer.pad_token = tokenizer.eos_token
    tokenizer.padding_side = "left"

    text = "Hello world!"
    tokenized_text = tokenizer(text)
    tokenized_text = {key: torch.tensor(value).unsqueeze(dim=0).to(torch.device("cuda")) for key, value in tokenized_text.items()}
    tokenized_text["labels"] = tokenized_text["input_ids"].clone()

    torch.manual_seed(0)
    model = llama.LlamaForCausalLM.from_pretrained(
        "/models/huggingface/meta-llama/llama-2-7b-chat-hf",
        local_files_only=True,
        use_safetensors=True,
        device_map=torch.device("cuda"),
        use_flash_attention_2=use_flash_attention_2,
        torch_dtype=torch.bfloat16,
    )
    assert isinstance(model, llama.LlamaForCausalLM)
    model.eval()
    for param in model.parameters():
        param.requires_grad = False

    model.model.layers[0].train()
    for param in model.model.layers[0].parameters():
        param.requires_grad = True

    optim = torch.optim.AdamW(model.parameters())

    torch.manual_seed(0)

    for i in range(10):
        output = model(**tokenized_text)
        loss = output["loss"]
        if i in (0, 9):
            print(loss)
        loss.backward()
        optim.step()
        optim.zero_grad()


if __name__ == "__main__":
    main()
```

```console
$ python flash_attn_non_determinism.py --use-flash-attention-2
tensor(5.6612, device='cuda:0', grad_fn=<NllLossBackward0>)
tensor(0.3542, device='cuda:0', grad_fn=<NllLossBackward0>)
$ python flash_attn_non_determinism.py
tensor(5.6589, device='cuda:0', grad_fn=<NllLossBackward0>)
tensor(0.2275, device='cuda:0', grad_fn=<NllLossBackward0>)
```

### Expected behavior

I am not expecting the magnitude of the difference between the 2 implementations. A difference of `0.1267` compared to `0.3542` seems very large.