[WIP] Lm loss feed forward chunking
Author: Pradhy729
Pull request: yes
Labels: name: wontfix

Final word embedding layer in LM loss calculation presents a bottleneck as the entire time dim - can be chunked similar to feed forward layers.

Comments:
# [Codecov](https://codecov.io/gh/huggingface/transformers/pull/6366?src=pr&el=h1) Report
> Merging [#6366](https://codecov.io/gh/huggingface/transformers/pull/6366?src=pr&el=desc) into [master](https://codecov.io/gh/huggingface/transformers/commit/6e8a38568eb874f31eb49c42285c3a634fca12e7&el=desc) will **decrease** coverage by `0.96%`.
> The diff coverage is `n/a`.

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

```diff
@@            Coverage Diff             @@
##           master    #6366      +/-   ##
==========================================
- Coverage   79.34%   78.38%   -0.97%     
==========================================
  Files         148      148              
  Lines       27196    27196              
==========================================
- Hits        21579    21317     -262     
- Misses       5617     5879     +262     
```


| [Impacted Files](https://codecov.io/gh/huggingface/transformers/pull/6366?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/transformers/modeling\_tf\_mobilebert.py](https://codecov.io/gh/huggingface/transformers/pull/6366/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9tb2JpbGViZXJ0LnB5) | `24.55% <0.00%> (-70.09%)` | :arrow_down: |
| [src/transformers/generation\_tf\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/6366/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9nZW5lcmF0aW9uX3RmX3V0aWxzLnB5) | `86.71% <0.00%> (-0.26%)` | :arrow_down: |
| [src/transformers/tokenization\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/6366/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fdXRpbHMucHk=) | `90.40% <0.00%> (+0.40%)` | :arrow_up: |
| [src/transformers/tokenization\_bert.py](https://codecov.io/gh/huggingface/transformers/pull/6366/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fYmVydC5weQ==) | `91.51% <0.00%> (+0.44%)` | :arrow_up: |
| [src/transformers/tokenization\_openai.py](https://codecov.io/gh/huggingface/transformers/pull/6366/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fb3BlbmFpLnB5) | `84.09% <0.00%> (+1.51%)` | :arrow_up: |
| [src/transformers/tokenization\_utils\_fast.py](https://codecov.io/gh/huggingface/transformers/pull/6366/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fdXRpbHNfZmFzdC5weQ==) | `94.28% <0.00%> (+2.14%)` | :arrow_up: |
| [src/transformers/tokenization\_transfo\_xl.py](https://codecov.io/gh/huggingface/transformers/pull/6366/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fdHJhbnNmb194bC5weQ==) | `42.48% <0.00%> (+3.75%)` | :arrow_up: |
| [src/transformers/tokenization\_utils\_base.py](https://codecov.io/gh/huggingface/transformers/pull/6366/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fdXRpbHNfYmFzZS5weQ==) | `93.84% <0.00%> (+7.41%)` | :arrow_up: |
| [src/transformers/tokenization\_roberta.py](https://codecov.io/gh/huggingface/transformers/pull/6366/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy90b2tlbml6YXRpb25fcm9iZXJ0YS5weQ==) | `98.63% <0.00%> (+21.91%)` | :arrow_up: |
| [src/transformers/modeling\_tf\_gpt2.py](https://codecov.io/gh/huggingface/transformers/pull/6366/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9ncHQyLnB5) | `94.72% <0.00%> (+22.87%)` | :arrow_up: |

------

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


@patrickvonplaten quick question here: Did you mean to chunk the projection onto vocab_size as in:
https://github.com/huggingface/transformers/blob/6e8a38568eb874f31eb49c42285c3a634fca12e7/src/transformers/modeling_bert.py#L527-L530
or the transformation that happens before that:
https://github.com/huggingface/transformers/blob/6e8a38568eb874f31eb49c42285c3a634fca12e7/src/transformers/modeling_bert.py#L506-L510 


This one is a bit harder actually. I meant the both the calculation in https://github.com/huggingface/transformers/blob/6e8a38568eb874f31eb49c42285c3a634fca12e7/src/transformers/modeling_bert.py#L795 and https://github.com/huggingface/transformers/blob/6e8a38568eb874f31eb49c42285c3a634fca12e7/src/transformers/modeling_bert.py#L1015 should be put into one chunk function. Here a lot of memory can be saved because currently go from the `last_hidden_state` tensor of size `[batch_size, seq_len, hidden_size]` to a `[batch_size, seq_len, vocab_size]` logit tensor and then reduce it to `[1]` loss scalar. Note that `vocab_size` is much larger than `hidden_size` and often is the bottleneck of a model. We don't need to compute `[batch_size, seq_len, vocab_size]` though if we apply chunked "loss" calculation from `last_hidden_states` to `loss` directly. Here we could greatly reduce memory consumption. But definitely leave this PR for later, we need to more carefully think about possible design changes as the two lines in question (linked above) are in different spots in the code. When we have the other `feed_forward_chunking` implemented we can take a look at this again :-) 

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.