GPT Neo past_key_values unexpected behaviour
Author: edwinagnew
Pull request: no
Labels: name: WIP

I have been successfully using the GPT2LMHeadModel module for text generation for some time and I recently tried to reuse the code to generate with GPTNeoForCausalLM. Though the documentations appear identical, I get the error "ValueError: not enough values to unpack (expected 2, got 1)"  for the line`output, past = self.model(context, past_key_values=past, use_cache=True).values()` (which works fine for GPT2).

Is this a bug or has the documentation been copied incorrectly? Would appreciate any tips for fixing.

Many thanks

Comments:
I encountered a similar problem when trying to use GPT-Neo with PPLM (https://github.com/uber-research/PPLM). Seems that Neo's `past_key_values` is returning and consuming key-value tensors as well as (I'm guessing) feed-forward tensors:

```python
inputs = tokenizer(prompt, return_tensors='pt')
outputs = model(**inputs)
past = outputs.past_key_values

for idx, p in enumerate(past):
    print(f'{idx}: {tuple(elem.shape for elem in p)}')

# output
# 0: (torch.Size([1, 12, 3, 64]), torch.Size([1, 12, 3, 64]))
# 1: (torch.Size([1, 3, 768]),)
# 2: (torch.Size([1, 12, 3, 64]), torch.Size([1, 12, 3, 64]))
# 3: (torch.Size([1, 3, 768]),)
# 4: (torch.Size([1, 12, 3, 64]), torch.Size([1, 12, 3, 64]))
# 5: (torch.Size([1, 3, 768]),)
# 6: (torch.Size([1, 12, 3, 64]), torch.Size([1, 12, 3, 64]))
# 7: (torch.Size([1, 3, 768]),)
# 8: (torch.Size([1, 12, 3, 64]), torch.Size([1, 12, 3, 64]))
# 9: (torch.Size([1, 3, 768]),)
# 10: (torch.Size([1, 12, 3, 64]), torch.Size([1, 12, 3, 64]))
# 11: (torch.Size([1, 3, 768]),)
```

GPT-2 correctly returns just the key-value tensors:

```python
# 0: (torch.Size([1, 12, 3, 64]), torch.Size([1, 12, 3, 64]))
# 1: (torch.Size([1, 12, 3, 64]), torch.Size([1, 12, 3, 64]))
# 2: (torch.Size([1, 12, 3, 64]), torch.Size([1, 12, 3, 64]))
# 3: (torch.Size([1, 12, 3, 64]), torch.Size([1, 12, 3, 64]))
# 4: (torch.Size([1, 12, 3, 64]), torch.Size([1, 12, 3, 64]))
# 5: (torch.Size([1, 12, 3, 64]), torch.Size([1, 12, 3, 64]))
# 6: (torch.Size([1, 12, 3, 64]), torch.Size([1, 12, 3, 64]))
# 7: (torch.Size([1, 12, 3, 64]), torch.Size([1, 12, 3, 64]))
# 8: (torch.Size([1, 12, 3, 64]), torch.Size([1, 12, 3, 64]))
# 9: (torch.Size([1, 12, 3, 64]), torch.Size([1, 12, 3, 64]))
# 10: (torch.Size([1, 12, 3, 64]), torch.Size([1, 12, 3, 64]))
# 11: (torch.Size([1, 12, 3, 64]), torch.Size([1, 12, 3, 64]))
```

After some more testing, the above seems to be because of local attention layers in GPT-Neo's default configuration. When specifying ```config = GPTNeoConfig(attention_types=[[["global"], 24]])```, I get similar `past_key_values` as in GPT-2:
```python
# 0: (torch.Size([1, 16, 3, 128]), torch.Size([1, 16, 3, 128]))
# 1: (torch.Size([1, 16, 3, 128]), torch.Size([1, 16, 3, 128])) 
# 2: (torch.Size([1, 16, 3, 128]), torch.Size([1, 16, 3, 128])) 
# 3: (torch.Size([1, 16, 3, 128]), torch.Size([1, 16, 3, 128])) 
# 4: (torch.Size([1, 16, 3, 128]), torch.Size([1, 16, 3, 128])) 
# ...
```

I do think the [documentation](https://huggingface.co/transformers/model_doc/gpt_neo.html#transformers.GPTNeoModel.forward) for `past_key_values` should be updated since it currently says: "with each tuple having 2 tensors of shape (batch_size, num_heads, sequence_length, embed_size_per_head)"

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.

Hi @patil-suraj, just checking if there is any progress on this issue or pull request #11630? That PR seems to fix the problem related to my usecase.

The different shape for local attention layers is because of the folding going on in the current implementation.

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.