Fix: Qwen2 eager attention bug
Author: HandH1998
Pull request: yes

## Bug
There are three supported attention modes: `eager`, `sdpa`, and `flash_attention_2` for Qwen2 in transformers.  I have evaluated wikitext2 PPL on the Qwen2-7b model under different attention modes. The results are as follows. 

I noticed that overflow might occur during the q * k calculation in `eager` attention mode when the dtype is `float16` or `bfloat16`. The bold data represents abnormal results which can verify my findings.

| precision |  eager | sdpa  | flash_attention 2 | fixed eager |
|-----------|---------------|-------|-------------------| --------|
| FP32      | 7.1378     | 7.1378| not support fp32  | 7.1378 | 
| FP16      | **nan**   | 7.1381| 7.1377            | 7.1379 |
| BF16      | **7.9967**    | 7.1330| 7.1338            | 7.1334 |

## Fix
I converted query and key to `float32` before performing the query * key operation to fix the bug in `eager` attention mode, when dtype is `float16` or `bfloat16`. The corrected results are displayed in the last column of the aforementioned table.


## Reproduce
You can utilize the following script to reproduce the table results with the latest transformers.
```
from tqdm import tqdm
import argparse
import random
import numpy as np
import torch
import torch.nn as nn
from transformers import AutoTokenizer, AutoModelForCausalLM
from datasets import load_dataset


DTYPE_MAP = {
    "float16": torch.float16,
    "bfloat16": torch.bfloat16,
    "float32": torch.float32,
}


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model_path",
        type=str,
        default="",
        help="path contains model weight",
    )
    parser.add_argument(
        "--tokenizer_path",
        type=str,
        default="",
        help="path contains tokenizer",
    )
    parser.add_argument(
        "--attn_implementation",
        type=str,
        default="eager",
        choices=["eager", "sdpa", "flash_attention_2"],
    )
    parser.add_argument(
        "--dtype",
        type=str,
        default="float16",
        choices=["float16", "bfloat16", "float32"],
    )
    parser.add_argument("--batch_size", type=int, default=8)
    parser.add_argument("--max_length", type=int, default=2048)
    parser.add_argument("--seed", type=int, default=0)
    return parser.parse_args()


def setup_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    np.random.seed(seed)
    random.seed(seed)
    torch.backends.cudnn.deterministic = True


def get_wikitext2(nsamples, seed, seqlen, tokenizer):
    print("get_wikitext2")
    traindata = load_dataset("wikitext", "wikitext-2-raw-v1", split="train")
    testdata = load_dataset("wikitext", "wikitext-2-raw-v1", split="test")

    trainenc = tokenizer("\n\n".join(traindata["text"]), return_tensors="pt")
    testenc = tokenizer("\n\n".join(testdata["text"]), return_tensors="pt")

    random.seed(seed)
    trainloader = []
    for _ in range(nsamples):
        i = random.randint(0, trainenc.input_ids.shape[1] - seqlen - 1)
        j = i + seqlen
        inp = trainenc.input_ids[:, i:j]
        tar = inp.clone()
        tar[:, :-1] = -100
        trainloader.append((inp, tar))
    return trainloader, testenc


@torch.no_grad()
def eval_model(model, tokenizer, args):
    max_length = args.max_length
    _, testloader = get_wikitext2(128, seed=0, seqlen=max_length, tokenizer=tokenizer)
    testenc = testloader.input_ids

    nsamples = testenc.numel() // max_length

    nlls = []
    for i in tqdm(range(nsamples)):
        batched_inps = testenc[:, (i * max_length) : ((i + 1) * max_length)].to(
            model.device
        )
        outputs = model.model(batched_inps)
        hidden_states = outputs[0]
        logits = model.lm_head(hidden_states)
        shift_logits = logits[:, :-1, :]
        shift_labels = testenc[:, (i * max_length) : ((i + 1) * max_length)][:, 1:].to(
            model.lm_head.weight.device
        )
        loss_fct = nn.CrossEntropyLoss()
        loss = loss_fct(
            shift_logits.view(-1, shift_logits.size(-1)),
            shift_labels.view(-1),
        )
        neg_log_likelihood = loss.float() * max_length
        nlls.append(neg_log_likelihood)

    ppl = torch.exp(torch.stack(nlls).sum() / (nsamples * max_length))
    print("Wikitext2 PPL: ", ppl)


if __name__ == "__main__":
    args = parse_args()
    setup_seed(args.seed)
    model = AutoModelForCausalLM.from_pretrained(
        args.model_path,
        device_map="sequential",
        attn_implementation=args.attn_implementation,
        torch_dtype=DTYPE_MAP[args.dtype],
    )
    tokenizer = AutoTokenizer.from_pretrained(
        args.tokenizer_path,
        trust_remote_code=True,
    )
    eval_model(model, tokenizer, args)

```


Comments:
The official Qwen documentation mentioned this modification method.  
https://qwen.readthedocs.io/en/latest/inference/chat.html 
--> https://github.com/jklj077/transformers/commit/53dba62c9c35c4a632c81f0406d844ea3679d4dd
It might have some reference value.

Thanks, I see it. But why not merge the official qwen2 patch? 

Thanks for your PR @HandH1998 ! cc @ArthurZucker 

Hey! Since we merged #33312, I am guessing we can close this one no? 

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_33317). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@ArthurZucker @GeLee-Q It is great that you guys have fixed this issue in a more elegant way. Before closing this PR, I'm curious if replacing NaN with zero impacts performance, even though the outputs seem normal in https://github.com/huggingface/transformers/pull/33312. I would appreciate it if you could provide evaluation results on public datasets, such as Wikitext2.

I won't have time to provide this 😢 
However, this was used in `Deberta` models for example. 
It brings your model closer to actual outputs AFAIK since it properly negates paddings instead of having a very very small value associated to it! (the usual torch.finfo.min)

@ArthurZucker OK, thanks.

> @ArthurZucker OK, thanks.

Hi, have you done some evaluations after this PR was closed? I find that when using bf16, flash_attention_2 performs better than eager, is it normal? Thanks for your time and I look forward to your reply!

@xinlyyy I didn't evaluated it. I think it is normal if the difference is not large between `flash_attn_2` and `eager`.

Yep, in general eager will still attend to masked token, even if a tad little bit, while fa2 won't