Bug in contrastive search with GPT2-decoders crossattention using batches
Author: josh-oo
Pull request: no

### System Info

- `transformers` version: 4.24.0
- Platform: Linux-5.10.133+-x86_64-with-Ubuntu-18.04-bionic
- Python version: 3.7.15
- Huggingface_hub version: 0.11.0
- PyTorch version (GPU?): 1.12.1+cu113 (True)
- Tensorflow version (GPU?): 2.9.2 (True)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: yes
- Using distributed or parallel set-up in script?: no

### Who can help?

@patrickvonplaten @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

Load the models and prepare example inputs:
```python
import torch
from transformers import EncoderDecoderModel, AutoTokenizer

model = EncoderDecoderModel.from_encoder_decoder_pretrained(
    "bert-base-uncased", "gpt2"
)

in_tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
out_tokenizer = AutoTokenizer.from_pretrained("gpt2")

model.config.decoder_start_token_id = out_tokenizer.bos_token_id

inputs = in_tokenizer(["This is a simple test","This is a simple test"], return_tensors="pt")
input_ids = inputs['input_ids']
```

Calling the generate method:

```python
model.generate(input_ids, top_k=4, penalty_alpha=0.6)
```

Leads to a Runtime Error in GPTs crossattention layer:

```
[...]  
File "[...]/transformers/generation_utils.py", line 1511, in generate  
    return self.contrastive_search(  
[...]  
File "[..]/transformers/models/gpt2/modeling_gpt2.py", line 895, in forward  
    outputs = block(  
File "[..]/torch/nn/modules/module.py", line 1190, in _call_impl  
    return forward_call(*input, **kwargs)  
File "[..]/transformers/models/gpt2/modeling_gpt2.py", line 417, in forward  
    cross_attn_outputs = self.crossattention(  
File "[...]/transformers/models/gpt2/modeling_gpt2.py", line 336, in forward  
    attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)  
File "[...]/transformers/models/gpt2/modeling_gpt2.py", line 182, in _attn  
    attn_weights = torch.matmul(query, key.transpose(-1, -2))  
RuntimeError: The size of tensor a (8) must match the size of tensor b (2) at non-singleton dimension 0
```
# Notes:  
I have tried some combinations. The error appears only with plain gpt as decoder. And only when using batches greater than 1 . 
Beam search works although it also sends batches of size batch_size*num_beams to the gpt model. Therefore the error is probably in the contrastive search.

To produce the same Error without calling generate:
```python
encoder_output = (torch.rand(2, 7, 768), None) # shape: (batch_size, seq_len, hidden_dim)
encoder_attention_mask = torch.ones(8, 7) # shape: (batch_size*top_k, seq_len)
decoder_input_ids = torch.tensor([[0],[0]]) # shape: (batch_size, 1)

model(decoder_input_ids=decoder_input_ids, encoder_outputs=encoder_output, attention_mask=encoder_attention_mask)
```

### Expected behavior

Calling contrastive search method with batches should return batches of generated ids like beam search.

Comments:
Hi @josh-oo

I believe that issue has been fixed in the latest version as you can see in  [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1q25nXzjuvaYHMTtosH2VVzDuSk877Ta6)

Also I believe in the reproduction of the same error without generate, you shouldn't have `batch_size*top_k`, instead the following should work
```
encoder_output = (torch.rand(2, 7, 768), None) # shape: (batch_size, seq_len, hidden_dim)
encoder_attention_mask = torch.ones(2, 7) # shape: (batch_size, seq_len)
decoder_input_ids = torch.tensor([[0],[0]]) # shape: (batch_size, 1)

model(decoder_input_ids=decoder_input_ids, encoder_outputs=encoder_output, attention_mask=encoder_attention_mask)
```

Hey @josh-oo 👋 

Since the release of v4.24 we have standardized a few internals of contrastive search, which may have fixed existing bugs (like @kiansierra pointed out, thank you for pitching in!)

Hi @kiansierra and @gante , thanks for the quick reply. The new version seems to solve the problem.

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.