CodeLLama generates very different outputs with flash_attention_2
Author: ikergarcia1996
Pull request: no

### System Info

- `transformers` version: 4.35.0.dev0
- Platform: Linux-4.18.0-477.10.1.el8_8.x86_64-x86_64-with-glibc2.28
- Python version: 3.9.7
- Huggingface_hub version: 0.16.4
- Safetensors version: 0.3.1
- Accelerate version: 0.23.0
- Accelerate config: 	
        - compute_environment: LOCAL_MACHINE
	- distributed_type: NO
	- mixed_precision: bf16
	- use_cpu: False
	- debug: False
	- num_processes: 1
	- machine_rank: 0
	- num_machines: 1
	- rdzv_backend: static
	- same_network: False
	- main_training_function: main
	- downcast_bf16: False
	- tpu_use_cluster: False
	- tpu_use_sudo: False
- 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 1xNvidia A100
- Using distributed or parallel set-up in script?: No



### Who can help?

@muellerzr @pacman100 @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

First, we load codellama 7B with and without Flash Attention. I use `torch_dtype=torch.bfloat16`

```python
from transformers import AutoModelForCausalLM, AutoTokenizer, AutoConfig
import torch

model_weights_name_or_path = "codellama/CodeLlama-7b-hf"

tokenizer = AutoTokenizer.from_pretrained(model_weights_name_or_path)
tokenizer.pad_token_id = tokenizer.unk_token_id
            
            
config = AutoConfig.from_pretrained(
    model_weights_name_or_path,
    trust_remote_code=False,
    # pretraining_tp=1, Setting pretraining_tp=1 doesn't solve the issue
)

model = AutoModelForCausalLM.from_pretrained(
        pretrained_model_name_or_path=model_weights_name_or_path,
        device_map=None,
        max_memory=None,
        quantization_config=None,
        torch_dtype=torch.bfloat16,
        config=config,
        trust_remote_code=True,
        use_flash_attention_2=False
).to("cuda")

model_flash = AutoModelForCausalLM.from_pretrained(
        pretrained_model_name_or_path=model_weights_name_or_path,
        device_map=None,
        max_memory=None,
        quantization_config=None,
        torch_dtype=torch.bfloat16,
        config=config,
        trust_remote_code=True,
        use_flash_attention_2=True
).to("cuda")

model_flash.eval()
model.eval()

```

Prepare the input. I am using a NER example from: https://github.com/hitz-zentroa/GoLLIE/
```python

prompt = '''
# The following lines describe the task definition
@dataclass
class PrivateSpaceCompany(Entity):
    """Refers to private companies primarily focused on space exploration, transportation,
    satellite launch, or space-based services. These are non-governmental entities that have
    a commercial interest in space activities."""

    span: str  # Such as: "Blue origin", "Boeing", "Northrop Grumman", "Arianespace"


@dataclass
class PublicSpaceCompany(Entity):
    """Refers to governmental entities or agencies that are primarily focused on space
    exploration, research, transportation, satellite launch, or other space-based services.
    These entities are state-owned and operated and are generally funded through public funds.
    """

    span: str  # Such as "ESA", "ISRO", "CNSA"


@dataclass
class Planet(Entity):
    """Refers to celestial bodies that orbit a star. Planets are large enough
    to have cleared their orbits of other debris and have a nearly round shape
    due to their self-gravity."""

    span: str  # Such as: "Earth", "Jupiter", "Venus", "Mercury", "Saturn"


@dataclass
class Launcher(Entity):
    """Refers to a vehicle designed primarily to transport payloads from the Earth's
    surface to space. Launchers can carry various payloads, including satellites,
    crewed spacecraft, and cargo, into various orbits or even beyond Earth's orbit.
    They are usually multi-stage vehicles that use rocket engines for propulsion."""

    span: str  # Such as: "Sturn V", "Atlas V", "Soyuz", "Ariane 5"


# This is the text to analyze
text = "SpaceX is colaborating with NASA in the mission to bring humans to Mars using their new Starship rocket."

# The annotation instances that take place in the text above are listed here
result ='''.strip()
    
model_input = tokenizer(prompt, add_special_tokens=True, return_tensors="pt")
```


Finally, we can run both models with the same input

```python
model_flash_ouput = model_flash.generate(
    **model_input.to(model_flash.device),
    max_new_tokens=128,
    do_sample=False,
    min_new_tokens=0,
    num_beams=1,
    num_return_sequences=1,
)

model_ouput = model.generate(
    **model_input.to(model.device),
    max_new_tokens=128,
    do_sample=False,
    min_new_tokens=0,
    num_beams=1,
    num_return_sequences=1,
)

print(tokenizer.batch_decode(model_flash_ouput)[0])
print(tokenizer.batch_decode(model_ouput)[0])
```

The result (I skip the input tokens for brevity) is:

```python

# FLASH ATTENTION 2 MODEL
result = [ <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE> <PRE>

# BASELINE MODEL
result = [
    PrivateSpaceCompany(span="SpaceX"),
    PublicSpaceCompany(span="NASA"),
    Planet(span="Mars"),
    Launcher(span="Starship"),
]
```

### Expected behavior

Both models should produce the same result.

Comments:
cc @younesbelkada pretty sure we have logits tests and generation for Llama but not CodeLlama

@ArthurZucker @younesbelkada It seems that the issue is resolved after updating Flash Attention from version `2.0.8` to `2.3.2`. 

Very nice @ikergarcia1996 ! Indeed, we have a similar issue here: https://github.com/huggingface/transformers/issues/26697