Generate: Passing encoder_hidden_states in with num_beams > 1 raises an error
Author: michalk8
Pull request: no

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: both 4.12.5 and 4.13.0.dev0
- Platform: Linux
- Python version: 3.8.12
- PyTorch version (GPU?): 1.10, no GPU
- Tensorflow version (GPU?): N/A
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No

### Who can help
@patrickvonplaten
<!-- Your issue will be replied to more quickly if you can figure out the right person to tag with @
 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- ALBERT, BERT, XLM, DeBERTa, DeBERTa-v2, ELECTRA, MobileBert, SqueezeBert: @LysandreJik
- encoder-decoder models (For example, BlenderBot, BART, Marian, Pegasus, T5, ByT5): @patrickvonplaten, @patil-suraj
- Longformer, Reformer, TransfoXL, XLNet, FNet: @patrickvonplaten
- FSMT: @stas00
- Funnel: @sgugger
- GPT-2, GPT: @patrickvonplaten, @LysandreJik
- RAG, DPR: @patrickvonplaten, @lhoestq
- TensorFlow: @Rocketknight1
- JAX/Flax: @patil-suraj @patrickvonplaten 
- TAPAS, LayoutLM, LayoutLMv2, LUKE, ViT, BEiT, DEiT, DETR, CANINE: @NielsRogge
- GPT-Neo, GPT-J, CLIP: @patil-suraj
- Wav2Vec2, HuBERT, SpeechEncoderDecoder: @patrickvonplaten, @anton-l

If the model isn't in the list, ping @LysandreJik who will redirect you to the correct contributor.

Library:

- Benchmarks: @patrickvonplaten
- Deepspeed: @stas00
- Ray/raytune: @richardliaw, @amogkam
- Text generation: @patrickvonplaten
- Tokenizers: @LysandreJik
- Trainer: @sgugger
- Pipelines: @Narsil
- Speech: @patrickvonplaten, @anton-l
- Vision: @NielsRogge, @sgugger

Documentation: @sgugger

Model hub:

- for issues with a model, report at https://discuss.huggingface.co/ and tag the model's creator.

HF projects:

- datasets: [different repo](https://github.com/huggingface/datasets)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Examples:

- maintained examples (not research project or legacy): @sgugger, @patil-suraj

For research projetcs, please ping the contributor directly. For example, on the following projects:

- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->

## Information

Model I am using (Bert, XLNet ...): GPT2 (LM head)

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)

I have a non LM encoder (a set Transformer, to be precise), where the last dim is fed into GPT2 LM head. During the training, seems to work, as well as when calling generate, apart from the case when more than 1 beam is specified.
This seems to be resolved by stacking the passed hidden states along the batch dim, but I am unsure whether this is a correct fix.
Also if there's a better way of approaching this, I'd be happy to hear.

## To reproduce

Steps to reproduce the behavior:

```python
from transformers import GPT2Config, GPT2LMHeadModel
import numpy as np
import torch

class CustomGPT2LMHeadModel(GPT2LMHeadModel):
    def prepare_inputs_for_generation(self, input_ids: torch.LongTensor, past=None, **kwargs):
        res = super().prepare_inputs_for_generation(input_ids, past=past, **kwargs)
        res["encoder_hidden_states"] = kwargs["encoder_hidden_states"]
        return res

config = GPT2Config()
config.add_cross_attention = True
config.n_embd = 120

model = CustomGPT2LMHeadModel(config)

batch_size = 32
n_beams = 3
input_ids = torch.ones((batch_size, 1)).long() * config.bos_token_id
ehs = torch.from_numpy(np.random.normal(size=(batch_size, 16, config.n_embd))).float()
ehs_stacked = torch.cat([ehs] * n_beams, axis=0)

output = model.generate(input_ids=input_ids,
                        encoder_hidden_states=ehs,  # works with ehs_stacked or n_beams=1
                        num_beams=n_beams)
```

Traceback:
<details>

```pytb
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
/tmp/ipykernel_151115/3265814751.py in <module>
----> 1 output = model.generate(input_ids=input_ids,
      2                         encoder_hidden_states=ehs,
      3                         num_beams=n_beams)

~/.miniconda3/envs/cellrank/lib/python3.8/site-packages/torch/autograd/grad_mode.py in decorate_context(*args, **kwargs)
     26         def decorate_context(*args, **kwargs):
     27             with self.__class__():
---> 28                 return func(*args, **kwargs)
     29         return cast(F, decorate_context)
     30 

~/.miniconda3/envs/nsode/lib/python3.8/site-packages/transformers/generation_utils.py in generate(self, input_ids, max_length, min_length, do_sample, early_stopping, num_beams, temperature, top_k, top_p, repetition_penalty, bad_words_ids, bos_token_id, pad_token_id, eos_token_id, length_penalty, no_repeat_ngram_size, encoder_no_repeat_ngram_size, num_return_sequences, max_time, max_new_tokens, decoder_start_token_id, use_cache, num_beam_groups, diversity_penalty, prefix_allowed_tokens_fn, output_attentions, output_hidden_states, output_scores, return_dict_in_generate, forced_bos_token_id, forced_eos_token_id, remove_invalid_values, synced_gpus, **model_kwargs)
   1068                 input_ids, expand_size=num_beams, is_encoder_decoder=self.config.is_encoder_decoder, **model_kwargs
   1069             )
-> 1070             return self.beam_search(
   1071                 input_ids,
   1072                 beam_scorer,

~/.miniconda3/envs/nsode/lib/python3.8/site-packages/transformers/generation_utils.py in beam_search(self, input_ids, beam_scorer, logits_processor, stopping_criteria, max_length, pad_token_id, eos_token_id, output_attentions, output_hidden_states, output_scores, return_dict_in_generate, synced_gpus, **model_kwargs)
   1805             model_inputs = self.prepare_inputs_for_generation(input_ids, **model_kwargs)
   1806 
-> 1807             outputs = self(
   1808                 **model_inputs,
   1809                 return_dict=True,

~/.miniconda3/envs/cellrank/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1100         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1101                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1102             return forward_call(*input, **kwargs)
   1103         # Do not call functions when jit is used
   1104         full_backward_hooks, non_full_backward_hooks = [], []

~/.miniconda3/envs/nsode/lib/python3.8/site-packages/transformers/models/gpt2/modeling_gpt2.py in forward(self, input_ids, past_key_values, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, encoder_hidden_states, encoder_attention_mask, labels, use_cache, output_attentions, output_hidden_states, return_dict)
   1042         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
   1043 
-> 1044         transformer_outputs = self.transformer(
   1045             input_ids,
   1046             past_key_values=past_key_values,

~/.miniconda3/envs/cellrank/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1100         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1101                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1102             return forward_call(*input, **kwargs)
   1103         # Do not call functions when jit is used
   1104         full_backward_hooks, non_full_backward_hooks = [], []

~/.miniconda3/envs/nsode/lib/python3.8/site-packages/transformers/models/gpt2/modeling_gpt2.py in forward(self, input_ids, past_key_values, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, encoder_hidden_states, encoder_attention_mask, use_cache, output_attentions, output_hidden_states, return_dict)
    885                 )
    886             else:
--> 887                 outputs = block(
    888                     hidden_states,
    889                     layer_past=layer_past,

~/.miniconda3/envs/cellrank/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1100         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1101                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1102             return forward_call(*input, **kwargs)
   1103         # Do not call functions when jit is used
   1104         full_backward_hooks, non_full_backward_hooks = [], []

~/.miniconda3/envs/nsode/lib/python3.8/site-packages/transformers/models/gpt2/modeling_gpt2.py in forward(self, hidden_states, layer_past, attention_mask, head_mask, encoder_hidden_states, encoder_attention_mask, use_cache, output_attentions)
    415             residual = hidden_states
    416             hidden_states = self.ln_cross_attn(hidden_states)
--> 417             cross_attn_outputs = self.crossattention(
    418                 hidden_states,
    419                 attention_mask=attention_mask,

~/.miniconda3/envs/cellrank/lib/python3.8/site-packages/torch/nn/modules/module.py in _call_impl(self, *input, **kwargs)
   1100         if not (self._backward_hooks or self._forward_hooks or self._forward_pre_hooks or _global_backward_hooks
   1101                 or _global_forward_hooks or _global_forward_pre_hooks):
-> 1102             return forward_call(*input, **kwargs)
   1103         # Do not call functions when jit is used
   1104         full_backward_hooks, non_full_backward_hooks = [], []

~/.miniconda3/envs/nsode/lib/python3.8/site-packages/transformers/models/gpt2/modeling_gpt2.py in forward(self, hidden_states, layer_past, attention_mask, head_mask, encoder_hidden_states, encoder_attention_mask, use_cache, output_attentions)
    334             attn_output, attn_weights = self._upcast_and_reordered_attn(query, key, value, attention_mask, head_mask)
    335         else:
--> 336             attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)
    337 
    338         attn_output = self._merge_heads(attn_output, self.num_heads, self.head_dim)

~/.miniconda3/envs/nsode/lib/python3.8/site-packages/transformers/models/gpt2/modeling_gpt2.py in _attn(self, query, key, value, attention_mask, head_mask)
    191 
    192     def _attn(self, query, key, value, attention_mask=None, head_mask=None):
--> 193         attn_weights = torch.matmul(query, key.transpose(-1, -2))
    194 
    195         if self.scale_attn_weights:

RuntimeError: The size of tensor a (96) must match the size of tensor b (32) at non-singleton dimension 0

```

</details>

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior



<!-- A clear and concise description of what you would expect to happen. -->
No error is raised.

Comments:
Hey @michalk8,

In general, we try not to use the issue tracker too much for questions of non-official, costumized models. Could you maybe use the forum: https://discuss.huggingface.co/ instead the next time. 

To answer your question, the approach of doing:

```python
ehs_stacked = torch.cat([ehs] * n_beams, axis=0)
```

is definitely in the right direction! The idea is that if the user passes `encoder_hidden_states` she/he is also responsible for making sure the dimensions are correct. We sadly can't wrap complex use cases such this one into `generate` as it would quickly make the function unreadable (it already is a bit tbh). 

I think instead of stacking it, you actually have to interleave it between the batch size. The dimension order in generate is:
[ - batch_size 
    - beam_size
      - sequence_length
        - ...
 ]

In your case, you have batch size 32 and you will need to expand the encoder_hidden_states between each batch for the first dimension. You could probably use this code here: 
https://github.com/huggingface/transformers/blob/8395f14de6068012787d83989c3627c3df6a252b/src/transformers/generation_utils.py#L479

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/master/CONTRIBUTING.md) are likely to be ignored.

Hello, I want to ask is there any other way to pass `encoder_hidden_states` parameter when using `generate()` function, instead to create a new class inheriting `GPT2LMHeadModel` ? 

actually , if I directly use `GPT2LMHeadModel.generate( inputs = ..., encoder_hidden_states = ...)` , `encoder_hidden_states ` does not passed .