GPTNeo: RuntimeError: shape mismatch when using past_key_values to go forward more than one token
## 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: 4.6.0.dev0
- Platform: Linux-5.11.11-arch1-1-x86_64-with-glibc2.33
- Python version: 3.9.2
- PyTorch version (GPU?): 1.8.1 (False)
- Tensorflow version (GPU?): not installed (NA)
- Using GPU in script?: No
- Using distributed or parallel set-up in script?: No

### Who can help
gpt_neo: @LysandreJik, @patil-suraj 
<!-- 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: @LysandreJik
- blenderbot, bart, marian, pegasus, encoderdecoder,  t5: @patrickvonplaten, @patil-suraj
- longformer, reformer, transfoxl, xlnet: @patrickvonplaten
- fsmt: @stas00
- funnel: @sgugger
- gpt2: @patrickvonplaten, @LysandreJik
- rag: @patrickvonplaten, @lhoestq
- tensorflow: @LysandreJik

Library:

- benchmarks: @patrickvonplaten
- deepspeed: @stas00
- ray/raytune: @richardliaw, @amogkam
- text generation: @patrickvonplaten
- tokenizers: @LysandreJik
- trainer: @sgugger
- pipelines: @LysandreJik

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
- research_projects/bert-loses-patience: @JetRunner
- research_projects/distillation: @VictorSanh

 -->

## Information

Model I am using (Bert, XLNet ...): GPTNeo

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)

## To reproduce

My motivation is to use past caching with backtracking, e.g. we already computed for `a b c d e` but now we want to compute for `a b c F G`. Ideally we would be able to use the past values and then go forward once with ` F G`. I have this working with GPT2 but with GPTNeo I ran into a crash which I narrowed down to the steps below.

Steps to reproduce the behavior:

1. Run the following script. It also uses small GPT2 to show an example of things working as expected.

```
#!/usr/bin/env python3
import torch
from transformers import *

for model_class, path in [
    (GPT2LMHeadModel, "gpt2"),
    (GPTNeoForCausalLM, "EleutherAI/gpt-neo-1.3B"),
]:
    tokenizer = GPT2Tokenizer.from_pretrained(path)
    tokens = tokenizer.encode(
        "one two three four five six seven eight nine ten",
    )
    model = model_class.from_pretrained(path)
    for k in range(len(tokens)):
        # First do all but k tokens.
        output = model.forward(
            input_ids=torch.tensor(tokens[: len(tokens) - k], dtype=torch.long),
            past_key_values=None,
        )
        # Then the rest.
        if k > 0:
            output = model.forward(
                input_ids=torch.tensor(tokens[len(tokens) - k :], dtype=torch.long),
                past_key_values=output.past_key_values,
            )
        top_logit, top_token = sorted(
            [(v, i) for i, v in enumerate(output.logits[-1, :].float().tolist())],
            reverse=True,
        )[0]
        print(f"{path} {k} OK {tokenizer.decode([top_token])!r} {top_logit}")
```

Here is what I get:

```
gpt2 0 OK ' eleven' -66.31873321533203
gpt2 1 OK ' eleven' -66.31869506835938
gpt2 2 OK ' eleven' -66.31873321533203
gpt2 3 OK ' eleven' -66.31871795654297
gpt2 4 OK ' eleven' -66.3187255859375
gpt2 5 OK ' eleven' -66.3187484741211
gpt2 6 OK ' eleven' -66.31873321533203
gpt2 7 OK ' eleven' -66.31874084472656
gpt2 8 OK ' eleven' -66.31873321533203
gpt2 9 OK ' eleven' -66.31874084472656
EleutherAI/gpt-neo-1.3B 0 OK ' eleven' 0.025278091430664062
EleutherAI/gpt-neo-1.3B 1 OK ' eleven' 0.02527904510498047
Traceback (most recent call last):
  File "/home/sboparen/2021/desk04/bug/./doit.py", line 22, in <module>
    output = model.forward(
  File "/home/sboparen/2021/desk04/bug/transformers/models/gpt_neo/modeling_gpt_neo.py", line 959, in forward
    transformer_outputs = self.transformer(
  File "/usr/lib/python3.9/site-packages/torch/nn/modules/module.py", line 889,
in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/sboparen/2021/desk04/bug/transformers/models/gpt_neo/modeling_gpt_neo.py", line 843, in forward
    outputs = block(
  File "/usr/lib/python3.9/site-packages/torch/nn/modules/module.py", line 889,
in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/sboparen/2021/desk04/bug/transformers/models/gpt_neo/modeling_gpt_neo.py", line 550, in forward
    attn_outputs = self.attn(
  File "/usr/lib/python3.9/site-packages/torch/nn/modules/module.py", line 889,
in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/sboparen/2021/desk04/bug/transformers/models/gpt_neo/modeling_gpt_neo.py", line 492, in forward
    outputs = self.attention(
  File "/usr/lib/python3.9/site-packages/torch/nn/modules/module.py", line 889,
in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/sboparen/2021/desk04/bug/transformers/models/gpt_neo/modeling_gpt_neo.py", line 420, in forward
    query = self._split_seq_length_dim_to(query, 1, 1, self.embed_dim)
  File "/home/sboparen/2021/desk04/bug/transformers/models/gpt_neo/modeling_gpt_neo.py", line 225, in _split_seq_length_dim_to
    return torch.reshape(tensors, split_dim_shape + (hidden_size,))
RuntimeError: shape '[1, 1, 1, 2048]' is invalid for input of size 4096
```

<!-- 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

The script should finish without error and continue to print `OK ' eleven' 0.02527...` for all values of `k`.