Gemma 2 Inference with BF16 fails
Author: murthyrudra
Pull request: no

### System Info

```
- `transformers` version: 4.42.4
- Platform: Linux-4.18.0-513.11.1.el8_9.x86_64-x86_64-with-glibc2.28
- Python version: 3.11.9
- Huggingface_hub version: 0.23.4
- Safetensors version: 0.4.2
- Accelerate version: 0.28.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.3.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 distributed or parallel set-up in script?: No
- Using GPU in script?: Yes
- GPU type: NVIDIA A100-SXM4-80GB
```

### Who can help?

@ArthurZucker 

### Information

- [X] The official example scripts
- [ ] My own modified scripts

### Tasks

- [X] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

If I load Gemma-2 model in `fp16` and perform inference on A100 80GB GPU, the output seems normal
```
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-9b")
model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-2-9b",
    device_map="cuda:0",
    load_in_8bit=False,
    torch_dtype=torch.float16
)

input_text = ["Write me a poem about Machine Learning.", "I want to eat a", "I go to the", "I want", "I"]
input_ids = tokenizer.batch_encode_plus(input_text, return_tensors="pt", padding=True)

for t in input_ids:
    if torch.is_tensor(input_ids[t]):
        input_ids[t] = input_ids[t].to("cuda:0")

outputs = model.generate(**input_ids)
print(tokenizer.batch_decode(outputs))
```

Output:
```
'<bos>Write me a poem about Machine Learning.\n\nI’m not sure if I’m the', 
"<pad><pad><pad><bos>I want to eat a lot of food, but I don't want to", 
"<pad><pad><pad><pad><bos>I go to the gym 3-4 times a week and I'", 
'<pad><pad><pad><pad><pad><pad><bos>I want to start by saying that I am not a fan of', 
'<pad><pad><pad><pad><pad><pad><pad><bos>I have a 2000 323']
```

However, the same model when loaded in `bf16` goes haywire
Output:
```
['<bos>Write me a poem about Machine Learning.\n\nI’m not sure if I’m the', 
'<pad><pad><pad><bos>I want to eat a<pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>', 
'<pad><pad><pad><pad><bos>I go to the<pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>', 
'<pad><pad><pad><pad><pad><pad><bos>I want<pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>', 
'<pad><pad><pad><pad><pad><pad><pad><bos>I<pad><pad><pad><pad><pad><pad><pad><pad><pad><pad><pad>']
```

### Expected behavior

The expected output should be the same output using `fp16` inference

Comments:
Hey! can you make sure that you are indeed using either `eager` or `flash_attn` for the `attn_implementation` (should be eager by default). 
Was this broken with the patches or has this always been the case? 
The outputs are not expected to be the same, however bfloat16 should work better than float16 for sure! Which is why this is a bit weird!

Thanks, I will try and see if this works.

With `attn_implementation='eager'`
```
from transformers import AutoTokenizer, AutoModelForCausalLM
import torch

tokenizer = AutoTokenizer.from_pretrained("google/gemma-2-9b")
model = AutoModelForCausalLM.from_pretrained(
    "google/gemma-2-9b",
    device_map="cuda:0",
    load_in_8bit=False,
    torch_dtype=torch.bfloat16,
    attn_implementation="eager"
)

input_text = ["Write me a poem about Machine Learning.", "I want to eat a", "I go to the", "I want", "I"]
input_ids = tokenizer.batch_encode_plus(input_text, return_tensors="pt", padding=True)

for t in input_ids:
    if torch.is_tensor(input_ids[t]):
        input_ids[t] = input_ids[t].to("cuda:0")

outputs = model.generate(**input_ids)
print(tokenizer.batch_decode(outputs))
```

Output:
```
['<bos>Write me a poem about Machine Learning.\n\nI’m not sure if I’m the',
 "<pad><pad><pad><bos>I want to eat a lot of food, but I don't want to", 
"<pad><pad><pad><pad><bos>I go to the gym 3-4 times a week and I'", 
'<pad><pad><pad><pad><pad><pad><bos>I want to start by saying that I am not a fan of', 
'<pad><pad><pad><pad><pad><pad><pad><bos>I have a 2000 323']
```



Thanks @ArthurZucker setting `attn_implementation="eager"` is giving reasonable generations

glad I could be of help! 🤗 


I can confirm I suffered from the same issue, Autoloading the model with torch.bfloat16 and a batch size of 5 produced nonsensical results when padding was introduced only working on the 5th item in the batch. The issue was fixed after changed my attn_implementation to "eager"