Model.generate use_cache=True generates different results than use_cache=False
Author: edenlum
Pull request: no
Labels: name: bug

### System Info

- `transformers` version: 4.49.0
- Platform: Linux-6.5.0-35-generic-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.29.1
- Safetensors version: 0.5.3
- Accelerate version: 1.4.0
- Accelerate config:    not found
- DeepSpeed version: not installed
- PyTorch version (GPU?): 2.6.0+cu124 (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?: reproduced the issue with and without GPU
- GPU type: NVIDIA A40 

### Who can help?

@gante 

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

I see different outputs when generating text with model.generate with and without use_cache argument. Is this intended and how can I combat this?
The scores when I use cache (from the second token generated and onwards) are different. AFAIK use_cache is an optimization that shouldn’t effect the outputs. I also got this error on gpu (here in the code I use 'cpu). Code to reproduce:


```
from transformers import AutoTokenizer, AutoModelForCausalLM, DataCollatorWithPadding
import torch

MODEL = "meta-llama/Meta-Llama-3.1-8B-Instruct"


tokenizer = AutoTokenizer.from_pretrained(MODEL, padding_side='left')
model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.bfloat16) # , cache_dir="/workspace/fmrai/")

tokenizer.pad_token_id = tokenizer.eos_token_id
tokenizer.pad_token = tokenizer.eos_token

model.eval()

device = 'cpu' # 'cuda' if torch.cuda.is_available() else 'cpu'
model.to(device)
# Example text to generate from
prompt = "Tell me something that is very exciting"

# Format the prompt using the chat template
formatted_prompt = tokenizer.apply_chat_template([{"role": "user", "content": prompt}], tokenize=False)

# Tokenize the input
inputs = tokenizer(formatted_prompt, return_tensors="pt", padding=True)

def generate(model, tokenizer, inputs, use_cache, output_attentions=False, device='cpu'):
    with torch.no_grad():
        outputs = model.generate(
            inputs["input_ids"].to(device),
            attention_mask=inputs["attention_mask"].to(device),
            max_new_tokens=2,
            pad_token_id=tokenizer.eos_token_id,
            do_sample=False,
            return_dict_in_generate=True,
            use_cache=use_cache,
            output_scores=True,
            # output_hidden_states=True,
            output_attentions=output_attentions,
        )
    return outputs

outputs_cache = generate(model, tokenizer, inputs, use_cache=True, device=device)
outputs_no_cache = generate(model, tokenizer, inputs, use_cache=False, device=device)
outputs_cache_attentions = generate(model, tokenizer, inputs, use_cache=True, output_attentions=True, device=device)
outputs_no_cache_attentions = generate(model, tokenizer, inputs, use_cache=False, output_attentions=True, device=device)

for i in range(2):
    print(f"Cache {i}:        {outputs_cache.scores[i]}")
    print(f"No Cache {i}:     {outputs_no_cache.scores[i]}")
    print(f"Cache Att {i}:    {outputs_cache_attentions.scores[i]}")
    print(f"No Cache Att {i}: {outputs_no_cache_attentions.scores[i]}")
```


### Expected behavior

I think all output scores should be the same, no matter if use_cache, output_attentions is True/False.

Comments:
Hi Eden, I tried to reproduce this bug on my device, and it turns out this is very likely a problem with __model precision__. Using `torch_dtype=torch.float32` instead of `torch_dtype=torch.bfloat16` produces consistent outputs no matter `use_cache` is set to `True` or `False`.

__float32 results__
```bash
Cache 0:        tensor([[-12.6331, -11.2944, -13.4185,  ...,   8.6258,   8.6257,   8.6258]])
No Cache 0:     tensor([[-12.6331, -11.2944, -13.4185,  ...,   8.6258,   8.6257,   8.6258]])
Cache Att 0:    tensor([[-12.6331, -11.2944, -13.4185,  ...,   8.6258,   8.6257,   8.6258]])
No Cache Att 0: tensor([[-12.6331, -11.2944, -13.4185,  ...,   8.6258,   8.6257,   8.6258]])
Cache 1:        tensor([[-0.0462,  1.2840,  1.5762,  ...,  3.0992,  3.0989,  3.0987]])
No Cache 1:     tensor([[-0.0462,  1.2840,  1.5762,  ...,  3.0992,  3.0989,  3.0987]])
Cache Att 1:    tensor([[-0.0462,  1.2840,  1.5762,  ...,  3.0992,  3.0989,  3.0987]])
No Cache Att 1: tensor([[-0.0462,  1.2840,  1.5762,  ...,  3.0992,  3.0989,  3.0987]])
```
__bfloat16 results__
```bash
Cache 0:        tensor([[-12.8125, -11.3125, -13.4375,  ...,   8.6875,   8.6875,   8.6875]])
No Cache 0:     tensor([[-12.8125, -11.3125, -13.4375,  ...,   8.6875,   8.6875,   8.6875]])
Cache Att 0:    tensor([[-12.5625, -11.1875, -13.2500,  ...,   8.5625,   8.5625,   8.5625]])
No Cache Att 0: tensor([[-12.5625, -11.1875, -13.2500,  ...,   8.5625,   8.5625,   8.5625]])
Cache 1:        tensor([[-0.1299,  1.3359,  1.5938,  ...,  3.0625,  3.0625,  3.0625]])
No Cache 1:     tensor([[-0.0889,  1.2969,  1.5703,  ...,  3.0469,  3.0469,  3.0469]])
Cache Att 1:    tensor([[-0.0281,  1.3984,  1.6953,  ...,  3.0312,  3.0312,  3.0312]])
No Cache Att 1: tensor([[0.0310, 1.3438, 1.7266,  ..., 3.0781, 3.0781, 3.0781]])
```

If you are interested, I checked `LlamaModel.forward()` in [modeling_llama.py](https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py) and it turns out the hidden states after each decode layer are slightly different w/ and w/o cache (bfloat16 precision), and after 32 layers, this difference becomes somewhat significant. However, I believe this difference brings minimal harm to the final generation.

> Hi Eden, I tried to reproduce this bug on my device, and it turns out this is very likely a problem with **model precision**. Using `torch_dtype=torch.float32` instead of `torch_dtype=torch.bfloat16` produces consistent outputs no matter `use_cache` is set to `True` or `False`.
> 
> **float32 results**
> 
> Cache 0:        tensor([[-12.6331, -11.2944, -13.4185,  ...,   8.6258,   8.6257,   8.6258]])
> No Cache 0:     tensor([[-12.6331, -11.2944, -13.4185,  ...,   8.6258,   8.6257,   8.6258]])
> Cache Att 0:    tensor([[-12.6331, -11.2944, -13.4185,  ...,   8.6258,   8.6257,   8.6258]])
> No Cache Att 0: tensor([[-12.6331, -11.2944, -13.4185,  ...,   8.6258,   8.6257,   8.6258]])
> Cache 1:        tensor([[-0.0462,  1.2840,  1.5762,  ...,  3.0992,  3.0989,  3.0987]])
> No Cache 1:     tensor([[-0.0462,  1.2840,  1.5762,  ...,  3.0992,  3.0989,  3.0987]])
> Cache Att 1:    tensor([[-0.0462,  1.2840,  1.5762,  ...,  3.0992,  3.0989,  3.0987]])
> No Cache Att 1: tensor([[-0.0462,  1.2840,  1.5762,  ...,  3.0992,  3.0989,  3.0987]])
> **bfloat16 results**
> 
> Cache 0:        tensor([[-12.8125, -11.3125, -13.4375,  ...,   8.6875,   8.6875,   8.6875]])
> No Cache 0:     tensor([[-12.8125, -11.3125, -13.4375,  ...,   8.6875,   8.6875,   8.6875]])
> Cache Att 0:    tensor([[-12.5625, -11.1875, -13.2500,  ...,   8.5625,   8.5625,   8.5625]])
> No Cache Att 0: tensor([[-12.5625, -11.1875, -13.2500,  ...,   8.5625,   8.5625,   8.5625]])
> Cache 1:        tensor([[-0.1299,  1.3359,  1.5938,  ...,  3.0625,  3.0625,  3.0625]])
> No Cache 1:     tensor([[-0.0889,  1.2969,  1.5703,  ...,  3.0469,  3.0469,  3.0469]])
> Cache Att 1:    tensor([[-0.0281,  1.3984,  1.6953,  ...,  3.0312,  3.0312,  3.0312]])
> No Cache Att 1: tensor([[0.0310, 1.3438, 1.7266,  ..., 3.0781, 3.0781, 3.0781]])
> If you are interested, I checked `LlamaModel.forward()` in [modeling_llama.py](https://github.com/huggingface/transformers/blob/main/src/transformers/models/llama/modeling_llama.py) and it turns out the hidden states after each decode layer are slightly different w/ and w/o cache (bfloat16 precision), and after 32 layers, this difference becomes somewhat significant. However, I believe this difference brings minimal harm to the final generation.

Just a follow-up, a relevant previous issue can be found [here](https://github.com/huggingface/transformers/issues/25420#issuecomment-1775317535)

Thank you very much :)