Extremely degenerated performance of StarCoder (GPTBigCodeForCausalLM) when using flash attention 2 
Author: yzhang1918
Pull request: no

### System Info

- `transformers` version: 4.37.2
- Platform: Linux-5.15.0-26-generic-x86_64-with-glibc2.31
- Python version: 3.10.12
- Huggingface_hub version: 0.20.3
- Safetensors version: 0.4.2
- PyTorch version (GPU?): 2.2.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?: (True)
- Using distributed or parallel set-up in script?: (False)

----
Extra Info:
- `flash-attn` version: 2.5.2
- GPU: NVIDIA A100-SXM4-80GB x2
- Output of `nvcc --version`:
```
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2023 NVIDIA Corporation
Built on Tue_Aug_15_22:02:13_PDT_2023
Cuda compilation tools, release 12.2, V12.2.140
Build cuda_12.2.r12.2/compiler.33191640_0
```


### Who can help?

@susnato @pacman100 @fxmarty 

### Information

- [X] 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

MWE:

```python
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer


model_path = "bigcode/starcoderbase"

tokenizer = AutoTokenizer.from_pretrained(model_path)

dtype = torch.bfloat16  # also works for torch.float16
device_map = 'auto'  # also works for single GPU
model_flash = AutoModelForCausalLM.from_pretrained(model_path, device_map=device_map, attn_implementation="flash_attention_2", torch_dtype=dtype)
model_sdpa = AutoModelForCausalLM.from_pretrained(model_path, device_map=device_map, attn_implementation="sdpa", torch_dtype=dtype)
# sdpa and eadge give the same results, skip here
# model_eager = AutoModelForCausalLM.from_pretrained(model_path, device_map='cuda:0', attn_implementation="eager", torch_dtype=dtype)

def gen(prompt, model):
    model_inputs = tokenizer([prompt], return_tensors="pt").to(model.device)
    generated_ids = model.generate(**model_inputs, max_new_tokens=30, do_sample=False)
    text = tokenizer.batch_decode(generated_ids)[0]
    return text


print('=' * 50)
prompt = '<filename>spec.yaml'
print('[flash attention 2]')
print(gen(prompt, model_flash))
print('-' * 50)
print('[sdpa]')
print(gen(prompt, model_sdpa))

print('=' * 50)
prompt = "<fim_prefix>def fib(n):<fim_suffix>    else:\n        return fib(n - 2) + fib(n - 1)<fim_middle>"
print('[flash attention 2]')
print(gen(prompt, model_flash))
print('-' * 50)
print('[sdpa]')
print(gen(prompt, model_sdpa))
```

Outputs:

```
==================================================
[flash attention 2]
<filename>spec.yaml
#
#
#
#
#
#
#
#
#
#
#
#
#
#
#
--------------------------------------------------
[sdpa]
<filename>spec.yaml
name: "k8s-cluster"
version: "0.1.0"
summary: "Kubernetes cluster"
description:

==================================================
[flash attention 2]
<fim_prefix>def fib(n):<fim_suffix>    else:
        return fib(n - 2) + fib(n - 1)<fim_middle>

def fib(n):
    return n * 2

def fib(n):
    return n * 2

def fib
--------------------------------------------------
[sdpa]
Setting `pad_token_id` to `eos_token_id`:0 for open-end generation.
<fim_prefix>def fib(n):<fim_suffix>    else:
        return fib(n - 2) + fib(n - 1)<fim_middle>
    if n < 2:
        return n


```

### Expected behavior

The following `model` x `attn_implementation` combinations work perfectly on my env:
- {mistral family models} x {`eager`, `sdpa`, `flash-attention-2`}
- {starcoder} x {`eager`, `sdpa`}

However, the results of starcoder with flash-attention-2 are really wired as shown above.

I'm not sure if this problem is on `transformers`, `StarCoder`, `flash-attn`, or my side.

Thanks!


Comments:
cc @younesbelkada and unsubscribing myself

Hello, the above PR should resolve this issue.

Wow! Thanks for your quick response and fix!