Gradient checkpointing BERT & ALBERT poc
Author: LysandreJik
Pull request: yes
Labels: name: wontfix

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     
--------------------------------------------------------------------------------
```

Comments:
# [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.