Different loss when using different "attn_implemetation"， when using "flash_attention_2" is notebly larger than "eager" and "spda".
Author: ttw1018
Pull request: no
Labels: name: trainer

name: bug

name: Flash Attention

### System Info

<img width="1033" alt="image" src="https://github.com/user-attachments/assets/3134af0c-3895-482f-9e88-6f750316fe6c">


### Who can help?

@ArthurZucker

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

The loss of `spda`

```
{'loss': 0.5651, 'grad_norm': 0.4917651414871216, 'learning_rate': 9.989949748743719e-05, 'epoch': 0.01}
{'loss': 0.1834, 'grad_norm': 0.5474741458892822, 'learning_rate': 9.979899497487438e-05, 'epoch': 0.02}
{'loss': 0.0285, 'grad_norm': 0.09983359277248383, 'learning_rate': 9.969849246231156e-05, 'epoch': 0.03}
{'loss': 0.0142, 'grad_norm': 0.2753373086452484, 'learning_rate': 9.959798994974875e-05, 'epoch': 0.04}
{'loss': 0.0116, 'grad_norm': 0.006525201722979546, 'learning_rate': 9.949748743718594e-05, 'epoch': 0.05}
{'loss': 0.0175, 'grad_norm': 0.11793188750743866, 'learning_rate': 9.939698492462311e-05, 'epoch': 0.06}
{'loss': 0.0563, 'grad_norm': 0.13649235665798187, 'learning_rate': 9.929648241206031e-05, 'epoch': 0.07}
{'loss': 0.0094, 'grad_norm': 0.127386212348938, 'learning_rate': 9.91959798994975e-05, 'epoch': 0.08}
{'loss': 0.0322, 'grad_norm': 0.6825785040855408, 'learning_rate': 9.909547738693468e-05, 'epoch': 0.09}
{'loss': 0.0069, 'grad_norm': 0.012013601139187813, 'learning_rate': 9.899497487437186e-05, 'epoch': 0.1}
```

The loss of `eager`

```
{'loss': 13.072, 'grad_norm': 297.276123046875, 'learning_rate': 9.987951807228915e-05, 'epoch': 0.01}
{'loss': 10.0997, 'grad_norm': 8193.296875, 'learning_rate': 9.975903614457832e-05, 'epoch': 0.02}
{'loss': 8.2837, 'grad_norm': 3707.72900390625, 'learning_rate': 9.963855421686748e-05, 'epoch': 0.04}
{'loss': 9.1116, 'grad_norm': 686.0602416992188, 'learning_rate': 9.951807228915663e-05, 'epoch': 0.05}
{'loss': 8.439, 'grad_norm': 44.402984619140625, 'learning_rate': 9.939759036144579e-05, 'epoch': 0.06}
{'loss': 7.4174, 'grad_norm': 1022.7124633789062, 'learning_rate': 9.927710843373495e-05, 'epoch': 0.07}
{'loss': 6.5367, 'grad_norm': 30.798368453979492, 'learning_rate': 9.91566265060241e-05, 'epoch': 0.08}
{'loss': 4.2007, 'grad_norm': 704.5174560546875, 'learning_rate': 9.903614457831326e-05, 'epoch': 0.1}
{'loss': 3.5107, 'grad_norm': 4172.1484375, 'learning_rate': 9.891566265060241e-05, 'epoch': 0.11}
{'loss': 4.6052, 'grad_norm': 980.4346313476562, 'learning_rate': 9.879518072289157e-05, 'epoch': 0.12}
{'loss': 5.478, 'grad_norm': 1532.0015869140625, 'learning_rate': 9.867469879518073e-05, 'epoch': 0.13}
```


The loss of `flash_attention_2`

```
{'loss': 15.1953, 'grad_norm': 1.0505084991455078, 'learning_rate': 9.989339019189766e-05, 'epoch': 0.01}
{'loss': 14.849, 'grad_norm': 0.15318499505519867, 'learning_rate': 9.978678038379531e-05, 'epoch': 0.02}
{'loss': 14.5066, 'grad_norm': 0.7192463278770447, 'learning_rate': 9.968017057569297e-05, 'epoch': 0.03}
{'loss': 15.1727, 'grad_norm': 0.13244017958641052, 'learning_rate': 9.957356076759062e-05, 'epoch': 0.04}
{'loss': 14.855, 'grad_norm': 0.08410739153623581, 'learning_rate': 9.946695095948828e-05, 'epoch': 0.05}
{'loss': 15.0897, 'grad_norm': 0.25853678584098816, 'learning_rate': 9.936034115138593e-05, 'epoch': 0.06}
{'loss': 14.8282, 'grad_norm': 0.1422465592622757, 'learning_rate': 9.925373134328359e-05, 'epoch': 0.07}
{'loss': 14.9486, 'grad_norm': 0.00950670801103115, 'learning_rate': 9.914712153518124e-05, 'epoch': 0.09}
{'loss': 14.6899, 'grad_norm': 0.4061712324619293, 'learning_rate': 9.90405117270789e-05, 'epoch': 0.1}
```

### Expected behavior

The loss of three attn_implemetation is close.

Comments:
Hey @ttw1018, do you have a reproducible snippet for what you get? Thanks

I repeat the previous experiments on my code, the only difference between them is the parameter of `attn_implementation`. I tried three kinds of `attn_implementation`, 'green' is `sdpa`, 'orange' is `eager` and 'black' is `flash_attention_2`. The result is shown below, the loss of `flash_attention_2` is notably lager than the other two and has no tendency to converge. 

<img width="1071" alt="image" src="https://github.com/user-attachments/assets/e5d245ba-ced7-40e9-a298-23d7e75178c9">

I want to reproduce it, and implement a simplified demo below. 

```python
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
from peft import get_peft_model, LoraConfig, TaskType
from matplotlib import pyplot as plt

seed = 42
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False

model_name = "./PTM/llama3.1"

tokenizer = AutoTokenizer.from_pretrained(model_name)

tokenizer.padding_side = "left"

if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

prompt = [
    "Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n### Instruction:\nTrack the state of the slot <hotel-parking> in the input dialogue.\n\n### Input:\n[user] I need information on a hotel that includes free parking please.[domain] hotel [slot] parking.This slot is categorical and you can choose from the following available values: yes,no,free,not care. In the previous dialogue state, the value of <hotel-parking> is none. If the slot is not mentioned in the dialogue, just return none. So the value of <hotel-parking> is\n\n### Response:\n",
    "Below is an instruction that describes a task, paired with an input that provides further context. Write a response that appropriately completes the request.\n\n### Instruction:\nTrack the state of the slot <taxi-destination> in the input dialogue.\n\n### Input:\n[system] I have booked your taxi look for a red tesla and the contact number is 07685178468. [user] Sure thank you so much, that is all.[domain] taxi [slot] destination. In the previous dialogue state, the value of <taxi-destination> is alexander bed and breakfast. If the slot is not mentioned in the dialogue, just return none. So the value of <taxi-destination> is\n\n### Response:\n",
]
answer = ["yes", "alexander bed and breakfast"]

input = tokenizer(
    [i[0] + i[1] for i in zip(prompt, answer)], return_tensors="pt", padding=True
)
labels = input["input_ids"].clone()
print(tokenizer.batch_decode(labels))
for i in range(len(prompt)):
    c = input["attention_mask"][i].tolist().count(0)
    labels[i, c : c + len(tokenizer.tokenize(prompt[i]))] = -100
input["labels"] = labels
input = input.to("cuda")

print(input)


def train(attn_implementation):
    base_model = AutoModelForCausalLM.from_pretrained(
        model_name, torch_dtype=torch.bfloat16, attn_implementation=attn_implementation
    )

    lora_config = LoraConfig(
        task_type=TaskType.CAUSAL_LM,
        r=8,
        target_modules=["q_proj", "k_proj", "v_proj", "o_proj"],
        lora_dropout=0.05,
        bias="none",
    )

    model = get_peft_model(base_model, lora_config).to("cuda")

    optimizer = torch.optim.Adam(model.parameters(), lr=1e-4)

    loss_list = []

    for i in range(20):
        optimizer.zero_grad()
        loss = model(**input).loss
        loss.backward()
        optimizer.step()
        loss_list.append(loss.item())
        print(loss.item())
    return loss_list


plt.figure(figsize=(10, 8))

loss_list = []
for attn_implementation in ["flash_attention_2", "eager", "sdpa"]:
    loss = train(attn_implementation)
    loss_list.append(loss)

for i, label in zip(loss_list, ["flash_attention_2", "eager", "sdpa"]):
    plt.plot(range(len(i)), i, label=label)

plt.legend()
plt.savefig("loss.png")
``` 
<img width="656" alt="image" src="https://github.com/user-attachments/assets/af446779-86d1-4edb-bff9-c2fc3149f225">


This is very interesting! TBH I think this might be more an issue for the FA2 repo, as I believe the way we call `FA2` is fairly simple. But, if this comes form us, happy to fix!

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.