GPTNeo: RuntimeError: shape mismatch when using past_key_values to go forward more than one token
Author: sboparen
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: 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`.


Comments:
Hi @sboparen 

Right now the caching is implemented such that when `past_key_values` are passed current token length must be 1.
This is due to the local attention layer which uses dynamic block length. This is a known limitation and I'm working on it at the moment.

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.

Unstale

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.