Output from model.Generate & model.forward not same when output attention/hidden_state is True
Author: AamodThakur
Pull request: no
Labels: name: bug

name: Generation

### System Info

- `transformers` version: 4.42.4
- Platform: Linux-5.15.0-107-generic-x86_64-with-glibc2.31
- Python version: 3.10.13
- Huggingface_hub version: 0.24.0
- Safetensors version: 0.4.3
- Accelerate version: 0.31.0
- Accelerate config: 	not found
- PyTorch version (GPU?): 2.2.2+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?: No
- GPU type: Tesla V100-SXM2-32GB

### Who can help?

@gante @Arth

### Information

- [X] The official example scripts
- [ ] 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

**Common Step:**
```
tokenizer = GemmaTokenizer.from_pretrained(model_path, device_map = device)
model = GemmaForCausalLM.from_pretrained(model_path, device_map = device)

prompt = ["
What is USA?
"]
input_tokens = tokenizer(prompt, return_tensors = 'pt', padding=True).to(device)

## Lets call this common_output
generate_ids = model.generate(**input_tokens, return_dict_in_generate=True,
                            temperature=None,
                            top_p=None, max_new_tokens=5, use_cache=False)
```

**Case 1:**
```
test_input_tokens = input_tokens

generate_ids_1 = model.forward(**test_input_tokens, output_attentions=True, output_hidden_states=True)
next_token_logits = generate_ids_1.logits[:, -1, :]
test_input_tokens["input_ids"] = torch.concat((test_input_tokens["input_ids"], 
                                   next_token_logits.argmax().unsqueeze(dim=0).unsqueeze(dim=0)), dim=-1)
test_input_tokens["attention_mask"] = torch.concat((test_input_tokens["attention_mask"], torch.ones((1,1), dtype=torch.long).to("cuda")), dim=-1)
```

**Case 2:**

```
test_input_tokens = input_tokens

generate_ids_1 = model.forward(**test_input_tokens)
next_token_logits = generate_ids_1.logits[:, -1, :]
test_input_tokens["input_ids"] = torch.concat((test_input_tokens["input_ids"], 
                                   next_token_logits.argmax().unsqueeze(dim=0).unsqueeze(dim=0)), dim=-1)
test_input_tokens["attention_mask"] = torch.concat((test_input_tokens["attention_mask"], torch.ones((1,1), dtype=torch.long).to("cuda")), dim=-1)

```

We are getting output of generate function & output of Case 2 same. But for output Case 1 is different because we have set "output_attentions=True, output_hidden_states=True" which we believe should be used for debugging purpose.

### Expected behavior

Output of generate should be same as Case 1 & Case 2's output.

Comments:
Hi @AamodThakur 👋 

The differences are expected :) You are hitting an edge case of the optimized SDPA attention (the implementation used by default) when `output_attentions=True`, where it falls back to the eager attention implementation. As with most optimized FP operations, there are tiny fluctuations to be expected when compared to other implementations (see [here](https://github.com/huggingface/transformers/issues/25420#issuecomment-1775317535) for an explanation of why it happens).

You should also be getting a warning like "GemmaModel is using GemmaSdpaAttention, but `torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True` (...)" 🤗 

For reference, here are the relevant code lines: https://github.com/huggingface/transformers/blob/0a7af19f4dc868bafc82f35eb7e8d13bac87a594/src/transformers/models/gemma/modeling_gemma.py#L489-L503

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.