Update `past_key_values` in GPT-2
Author: forest1988
Pull request: yes

# What does this PR do?

It seems GPT-2 and BartDecoder has a different style of `past_key_values`.
Advised by @patrickvonplaten,
I opened this PR to change GPT-2's cache format from a single tensor to a tuple of 2 tensors.
Once this problem is solved, it is expected that `past_key_values` in GPT-2 will be handled in the same way as in Bart.

Sorry there remain some errors. This PR is [WIP].
I would appreciate your advice on how to update `generation_utils.py`.
Can I modify `_reorder_cache` so that past is replaced from Tuple[torch.Tensor] to Tuple[Tuple[torch.Tensor]],
or should I consider other output variations, output.mem and outputs.past_buckets_states?

Fixes #9391

From patrickvonplaten:

This PR cleans the `_reorder_cache` logic. Now `_reorcher_cache` defaults to an erroneous `NotImplementedError` in `generation_utils.py` forcing the model to implement its corresponding `_rerorder_cache` it the `modeling_...py` file itself. This is cleaner as `_reorder_cache` strongly differs from model to model. In addition, this PR makes sure that `gradient_checkpointing` can only be used if the model is in training mode and makes sure that `use_cache` is disabled when training and `gradient_checkpointing` is enabled to prevent errors.


## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [x] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/master/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/master/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?

GPT2: @LysandreJik, @patrickvonplaten


Comments:
CircleCI error messages says as below.

In `run_tests_torch`:
``` 
-- Docs: https://docs.pytest.org/en/stable/warnings.html
=========================== short test summary info ============================
FAILED tests/test_modeling_gpt2.py::GPT2ModelTest::test_beam_sample_generate
FAILED tests/test_modeling_gpt2.py::GPT2ModelTest::test_beam_search_generate
FAILED tests/test_modeling_gpt2.py::GPT2ModelTest::test_beam_search_generate_dict_outputs_use_cache
FAILED tests/test_modeling_gpt2.py::GPT2ModelTest::test_gpt2_gradient_checkpointing
FAILED tests/test_modeling_gpt2.py::GPT2ModelTest::test_group_beam_search_generate
==== 5 failed, 4202 passed, 1775 skipped, 744 warnings in 216.47s (0:03:36) ====

Exited with code exit status 1
CircleCI received exit code 1
```

In `run_tests_flax`:
```
FAILED tests/test_modeling_gpt2.py::GPT2ModelTest::test_beam_sample_generate
FAILED tests/test_modeling_gpt2.py::GPT2ModelTest::test_beam_search_generate
FAILED tests/test_modeling_gpt2.py::GPT2ModelTest::test_beam_search_generate_dict_outputs_use_cache
FAILED tests/test_modeling_gpt2.py::GPT2ModelTest::test_gpt2_gradient_checkpointing
FAILED tests/test_modeling_gpt2.py::GPT2ModelTest::test_group_beam_search_generate
==== 5 failed, 4172 passed, 1805 skipped, 751 warnings in 282.27s (0:04:42) ====

Exited with code exit status 1
CircleCI received exit code 1
```


Is there a difference between `past_key_value` and `layer_past`?　I understand that they both represent the contents of `past_key_values`, the past of each layer, but are they different?

I first thought it might be a difference between the Causal language model and the Seq2Seq language model, but it seems that both `past_key_value` and `layer_past` are used in `modeling_bart.py`.

And as for the contents of `layer_past`, should it be named `past_state`, as the following part of `modeling_bart.py` shows?

https://github.com/huggingface/transformers/blob/236cc365aff2512ef773c6b1786555dab6fb182f/src/transformers/models/bart/modeling_bart.py#L1236-L1244

I've updated `generation_utils.py`, and it seems `mems` in transfo_xl and xlnet causes a new error.

```
=========================== short test summary info ============================
FAILED tests/test_modeling_gpt2.py::GPT2ModelTest::test_gpt2_gradient_checkpointing
FAILED tests/test_modeling_transfo_xl.py::TransfoXLModelTest::test_beam_sample_generate
FAILED tests/test_modeling_transfo_xl.py::TransfoXLModelTest::test_beam_sample_generate_dict_output
FAILED tests/test_modeling_transfo_xl.py::TransfoXLModelTest::test_beam_search_generate
FAILED tests/test_modeling_transfo_xl.py::TransfoXLModelTest::test_beam_search_generate_dict_output
FAILED tests/test_modeling_transfo_xl.py::TransfoXLModelTest::test_group_beam_search_generate
FAILED tests/test_modeling_transfo_xl.py::TransfoXLModelTest::test_group_beam_search_generate_dict_output
FAILED tests/test_modeling_xlnet.py::XLNetModelTest::test_beam_sample_generate
FAILED tests/test_modeling_xlnet.py::XLNetModelTest::test_beam_sample_generate_dict_output
FAILED tests/test_modeling_xlnet.py::XLNetModelTest::test_beam_search_generate
FAILED tests/test_modeling_xlnet.py::XLNetModelTest::test_beam_search_generate_dict_output
FAILED tests/test_modeling_xlnet.py::XLNetModelTest::test_group_beam_search_generate
FAILED tests/test_modeling_xlnet.py::XLNetModelTest::test_group_beam_search_generate_dict_output
=== 13 failed, 4194 passed, 1775 skipped, 743 warnings in 205.38s (0:03:25) ====

Exited with code exit status 1
CircleCI received exit code 1
```

https://github.com/huggingface/transformers/blob/236cc365aff2512ef773c6b1786555dab6fb182f/src/transformers/models/xlnet/modeling_xlnet.py#L581-L607

It seems `mems` is something similar to `past_key_values`. 
Is there any difference between these two elements with different names?
Also, is it safe to change `mems` from `List[torch.Tensor]` to `Tuple[Tuple[torch.Tensor]]`?

Hey @forest1988, 

You're PR looks very nice! Yes, it is expected that `XLNet` and `TransfoXL` fail actually since they also have been using the "default" `_reorder_cache` function of `modeling_utils.py`. Could you do the following changes to correct this:

1) Copy that old `_reorder_cache` (the one before you did your changes) function that was in `generation_utils.py` to both `modeling_xlnet.py` and `modeling_transfo_xl.py` file so that those have the same function as before? 
2) Copy the current `_reorder_cache` function of `generation_utils.py` into `modeling_gpt2.py`?
3) Add a default `_reorder_cache` function to `generation_utils.py` that looks as follows:

```python
def _reorder_cache(self, past, beam_idx):
    raise NotImplementedError(...)
```

I've just updated `torch.utils.checkpoint.checkpoint` check in `modeling_gpt2.py`, referring to `modeling_bart.py`.

This way it's much cleaner and correct :-) The reason I'm proposing this change is that the `_reorder_cache` function is so different for each model that there should be **no** default function. A default function could confuse people that want to add a new model in a way that they think it works out of the box, but in most cases it just doesn't. A clear error message such as:


```python
def _reorder_cache(self, past, beam_idx):
    raise NotImplementedError(f"Make sure that a `_reorder_cache` function is correctly implemented in {self.__class__.__module__} to enable beam search for {self.__class__}")
```
```


I think this should solve the problems, let me know if you need more help :-) 

Thank you for your advice! I'll update `_reorder_cache` soon and commit it.

Hi @patrickvonplaten,

Thanks to your kind advice, I could solve the problem of `_reorder_cache` in `GPT-2`, `XLNet`, `TransfoXL` (, and `CTRL`).
Referring to `modeling_bart.py`, in which `_reorder_cache` is placed in `ConditionalGeneration` Model, I added `_reoder_cache` in `LMHead` Models in each Causal Language Models.

The last one remaining bug is:
```
FAILED tests/test_modeling_gpt2.py::GPT2ModelTest::test_gpt2_gradient_checkpointing
```

I think I should modify `test_gpt2_gradient_checkpointing` so that it has `use_cache=False`, or reconsider my previous update and re-modify the usage of `checkpoint` in modeling_gpt2.
> I've just updated torch.utils.checkpoint.checkpoint check in modeling_gpt2.py, referring to modeling_bart.py.
>


All checks have passed!
I appreciate all your help.

However, in the documentation of `_reorder_cache`, there are references to both `past_key_values` and `mems` regardless of which object is used.
I think we can fix that and only mention the one we use, or we can leave the reference to both to show that the aim of the function is the same.
If there is a need to modify it, please let me know.


Hi @patrickvonplaten,

> I hope it's fine for you that I went into the PR to do some final fixes. Thanks a lot for cleaning this up :-)

Of course! Thank you for adding fixes to make this PR more valuable!

Awesome, merging - great job @forest1988 !

Thank you for your advice and encouraging comments!
It’s my pleasure to have opened this PR!