Gradient checkpointing BERT & ALBERT poc
Proof of concept for gradient checkpointing in PyTorch, using a model-agnostic approach. The POC is done for BERT and ALBERT.

Pros:

- Model agnostic, only a few lines to add to models to be able to use this functionality
- Reinforces the model layer API, adding `get_layers()` (name to be discussed) alongside `get_input_embeddings()` and `get_output_embeddings()`

Cons:

- The checkpoint API can only handle positional arguments, pytorch tensors or None only. This means that:
    - The `output_hidden_states` must be cast to a tensor in the model
    - Models that pass keyword arguments to their layers need to pass positional arguments (see GPT-2 for example, which uses keyword arguments [here](https://github.com/huggingface/transformers/blob/b45e65efa0fbff2611ddd68e14fa75cacef3fe08/src/transformers/modeling_gpt2.py#L488-L493)).


If you think this is a cool API, I'll go ahead and implement this for the remaining models. @patrickvonplaten @thomwolf @julien-c @sgugger @ibeltagy

Here are the results using the benchmarking script:

```py
from transformers import PyTorchBenchmark, PyTorchBenchmarkArguments, BertConfig
args = PyTorchBenchmarkArguments(models=["bert-base-cased"], batch_sizes=[8], sequence_lengths=[8, 32, 128, 512], no_inference=True, training=True)
config_base = BertConfig.from_pretrained("bert-base-cased", gradient_checkpointing=False)
benchmark = PyTorchBenchmark(args, configs=[config_base])
benchmark.run()
```
Result (only relevant info):

```
====================        TRAIN - SPEED - RESULTS         ====================
--------------------------------------------------------------------------------
          Model Name             Batch Size     Seq Length     Time in s   
--------------------------------------------------------------------------------
       bert-base-cased               8               8             0.028     
       bert-base-cased               8               32            0.029     
       bert-base-cased               8              128            0.072     
       bert-base-cased               8              512            0.296     
--------------------------------------------------------------------------------

====================        TRAIN - MEMORY - RESULTS        ====================
--------------------------------------------------------------------------------
          Model Name             Batch Size     Seq Length    Memory in MB 
--------------------------------------------------------------------------------
       bert-base-cased               8               8              2419     
       bert-base-cased               8               32             2481     
       bert-base-cased               8              128             2985     
       bert-base-cased               8              512             8233     
--------------------------------------------------------------------------------
```

```py
from transformers import PyTorchBenchmark, PyTorchBenchmarkArguments, BertConfig
args = PyTorchBenchmarkArguments(models=["bert-base-cased"], batch_sizes=[8], sequence_lengths=[8, 32, 128, 512], no_inference=True, training=True)
config_base = BertConfig.from_pretrained("bert-base-cased", gradient_checkpointing=True)
benchmark = PyTorchBenchmark(args, configs=[config_base])
benchmark.run()
```

Result (only relevant info):

```
====================        TRAIN - SPEED - RESULTS         ====================
--------------------------------------------------------------------------------
          Model Name             Batch Size     Seq Length     Time in s   
--------------------------------------------------------------------------------
       bert-base-cased               8               8             0.049     
       bert-base-cased               8               32             0.05     
       bert-base-cased               8              128            0.109     
       bert-base-cased               8              512            0.473     
--------------------------------------------------------------------------------

====================        TRAIN - MEMORY - RESULTS        ====================
--------------------------------------------------------------------------------
          Model Name             Batch Size     Seq Length    Memory in MB 
--------------------------------------------------------------------------------
       bert-base-cased               8               8              2385     
       bert-base-cased               8               32             2403     
       bert-base-cased               8              128             2465     
       bert-base-cased               8              512             3969     
--------------------------------------------------------------------------------
```
# [Codecov](https://codecov.io/gh/huggingface/transformers/pull/5415?src=pr&el=h1) Report
> Merging [#5415](https://codecov.io/gh/huggingface/transformers/pull/5415?src=pr&el=desc) into [master](https://codecov.io/gh/huggingface/transformers/commit/87716a6d072b2b66415ce43086c73b04e63fe0fe&el=desc) will **increase** coverage by `0.01%`.
> The diff coverage is `100.00%`.

[![Impacted file tree graph](https://codecov.io/gh/huggingface/transformers/pull/5415/graphs/tree.svg?width=650&height=150&src=pr&token=9qOlN6Hb1c)](https://codecov.io/gh/huggingface/transformers/pull/5415?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master    #5415      +/-   ##
==========================================
+ Coverage   77.69%   77.71%   +0.01%     
==========================================
  Files         140      140              
  Lines       24334    24343       +9     
==========================================
+ Hits        18906    18917      +11     
+ Misses       5428     5426       -2     
```


| [Impacted Files](https://codecov.io/gh/huggingface/transformers/pull/5415?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/transformers/configuration\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/5415/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9jb25maWd1cmF0aW9uX3V0aWxzLnB5) | `95.71% <100.00%> (-0.69%)` | :arrow_down: |
| [src/transformers/modeling\_albert.py](https://codecov.io/gh/huggingface/transformers/pull/5415/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19hbGJlcnQucHk=) | `80.86% <100.00%> (+0.12%)` | :arrow_up: |
| [src/transformers/modeling\_bert.py](https://codecov.io/gh/huggingface/transformers/pull/5415/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19iZXJ0LnB5) | `88.46% <100.00%> (+0.80%)` | :arrow_up: |
| [src/transformers/modeling\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/5415/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ191dGlscy5weQ==) | `87.12% <100.00%> (+0.21%)` | :arrow_up: |
| [src/transformers/modeling\_tf\_openai.py](https://codecov.io/gh/huggingface/transformers/pull/5415/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9vcGVuYWkucHk=) | `20.78% <0.00%> (-74.20%)` | :arrow_down: |
| [src/transformers/tokenization\_auto.py](https://codecov.io/gh/huggingface/transformers/pull/5415/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fYXV0by5weQ==) | `95.55% <0.00%> (-2.23%)` | :arrow_down: |
| [src/transformers/modeling\_openai.py](https://codecov.io/gh/huggingface/transformers/pull/5415/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19vcGVuYWkucHk=) | `79.72% <0.00%> (-1.38%)` | :arrow_down: |
| [src/transformers/generation\_tf\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/5415/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9nZW5lcmF0aW9uX3RmX3V0aWxzLnB5) | `86.68% <0.00%> (+0.25%)` | :arrow_up: |
| [src/transformers/generation\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/5415/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9nZW5lcmF0aW9uX3V0aWxzLnB5) | `97.10% <0.00%> (+0.28%)` | :arrow_up: |
| ... and [8 more](https://codecov.io/gh/huggingface/transformers/pull/5415/diff?src=pr&el=tree-more) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/huggingface/transformers/pull/5415?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/huggingface/transformers/pull/5415?src=pr&el=footer). Last update [87716a6...b7e417a](https://codecov.io/gh/huggingface/transformers/pull/5415?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).
<|||||>I really like the API, I think it's fine if we enforce all attention layers to use positional arguments and wrap the output attentions bool into a tensor. 

Can we test how much memory is saved here for `bert-base-uncased` layers 6 - 18 for example? 
Should be quite easy to do now with the benchmark utils.<|||||>Benchmarked with the script and updated the PR @patrickvonplaten!<|||||>Awesome, it looks we can gain quite a lot of memory :-)<|||||>@LysandreJik , another problem is that `torch.utils.checkpoint.checkpoint` expects the function to return a tuple of `Variable`s. This won't work with forward functions that return other types as in [here](https://github.com/huggingface/transformers/blob/master/src/transformers/modeling_bart.py#L321).<|||||>This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.
<|||||>@LysandreJik , any plans to resurrect this? <|||||>Yes it's on my TODO (probably in ~2 weeks), and will be for most models (with some exceptions like BART and T5, which need a lot of plumbing to work with this POC)<|||||>This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.
