Feed forward chunking
Author: Pradhy729
Pull request: yes

Official PR for #5928 

Comments:
@patrickvonplaten - here's an initial implementation I have. My first step is to get the model to work with chunked feed forward - and it works! I still need to run the benchmark test to find out the benefits in terms of memory.

However, I see a problem. The new architecture causes some of the nn.Module weights and bias parameter-names to change - which would be a problem with loading existing pretrained weights from checkpoints. 
For example:
`bert.encoder.layer.0.intermediate.dense.weight` --> becomes `bert.encoder.layer.0.feed_forward.dense.dense.weight`

See the failing tests for more details. Any thoughts/ideas on how to get around this?


# [Codecov](https://codecov.io/gh/huggingface/transformers/pull/6024?src=pr&el=h1) Report
> Merging [#6024](https://codecov.io/gh/huggingface/transformers/pull/6024?src=pr&el=desc) into [master](https://codecov.io/gh/huggingface/transformers/commit/175cd45e13b2e33d1efec9e2ac217cba99f6ae58&el=desc) will **decrease** coverage by `0.31%`.
> The diff coverage is `100.00%`.

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

```diff
@@            Coverage Diff             @@
##           master    #6024      +/-   ##
==========================================
- Coverage   79.44%   79.12%   -0.32%     
==========================================
  Files         148      148              
  Lines       27193    27198       +5     
==========================================
- Hits        21604    21521      -83     
- Misses       5589     5677      +88     
```


| [Impacted Files](https://codecov.io/gh/huggingface/transformers/pull/6024?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/transformers/configuration\_reformer.py](https://codecov.io/gh/huggingface/transformers/pull/6024/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9jb25maWd1cmF0aW9uX3JlZm9ybWVyLnB5) | `100.00% <ø> (ø)` | |
| [src/transformers/configuration\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/6024/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9jb25maWd1cmF0aW9uX3V0aWxzLnB5) | `96.57% <100.00%> (+0.02%)` | :arrow_up: |
| [src/transformers/modeling\_bert.py](https://codecov.io/gh/huggingface/transformers/pull/6024/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19iZXJ0LnB5) | `88.49% <100.00%> (+0.09%)` | :arrow_up: |
| [src/transformers/modeling\_tf\_bert.py](https://codecov.io/gh/huggingface/transformers/pull/6024/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9iZXJ0LnB5) | `69.31% <0.00%> (-26.18%)` | :arrow_down: |
| [src/transformers/tokenization\_roberta.py](https://codecov.io/gh/huggingface/transformers/pull/6024/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fcm9iZXJ0YS5weQ==) | `76.71% <0.00%> (-21.92%)` | :arrow_down: |
| [src/transformers/tokenization\_utils\_base.py](https://codecov.io/gh/huggingface/transformers/pull/6024/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fdXRpbHNfYmFzZS5weQ==) | `86.43% <0.00%> (-7.42%)` | :arrow_down: |
| [src/transformers/tokenization\_transfo\_xl.py](https://codecov.io/gh/huggingface/transformers/pull/6024/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fdHJhbnNmb194bC5weQ==) | `38.73% <0.00%> (-3.76%)` | :arrow_down: |
| [src/transformers/tokenization\_utils\_fast.py](https://codecov.io/gh/huggingface/transformers/pull/6024/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fdXRpbHNfZmFzdC5weQ==) | `92.14% <0.00%> (-2.15%)` | :arrow_down: |
| [src/transformers/tokenization\_openai.py](https://codecov.io/gh/huggingface/transformers/pull/6024/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fb3BlbmFpLnB5) | `82.57% <0.00%> (-1.52%)` | :arrow_down: |
| [src/transformers/tokenization\_bert.py](https://codecov.io/gh/huggingface/transformers/pull/6024/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fYmVydC5weQ==) | `91.07% <0.00%> (-0.45%)` | :arrow_down: |
| ... and [4 more](https://codecov.io/gh/huggingface/transformers/pull/6024/diff?src=pr&el=tree-more) | |

------

[Continue to review full report at Codecov](https://codecov.io/gh/huggingface/transformers/pull/6024?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/6024?src=pr&el=footer). Last update [175cd45...406d621](https://codecov.io/gh/huggingface/transformers/pull/6024?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


hh

Hey @Pradhy729, thanks a lot for continuing the PR. I made a couple of changes: fix the docs and added tests for all models, whereas only Reformer and Bert tests are on for now. 

 Would be great if @LysandreJik @sgugger @thomwolf @sshleifer can review. 

This PR shows how `feed_forward_chunking` can be employed for all models. Feed forward chunking is explained here: https://huggingface.co/blog/reformer#2-chunked-feed-forward-layers in combination with some benchmarking. It can give good memory improvements for certain model architectures.
For Bert a test is added showing that the model gives equal results. This function can easily be added to other models, the same way it was done for BERT. There is no real drawback in implementing this IMO. 

**To-Do after review is positive**:
1. Add feed forward chunking to more models. @Pradhy729, feel free to add it to as many models as you want. The rest can also be added in a new PR or we open a "first good issue" for it.
2. Add feed forward chunking for language modeling loss. Chunking of feed forward layers in the attention block is often not really helpful to save memory - only if the model has very few attention heads. On the other hand, a real bottleneck is often the last word embedding layer. When training the loss does not have to be calculated in one huge batch (over time dim), but can be chunked the same way it is done here for Feed forward layers. This is not even really implemented in Reformer yet and would definitely require a new PR.

Great! Thanks @patrickvonplaten 
Will wait for reviewers and start working on the others.