[TF Longformer] Improve Speed for TF Longformer
Author: patrickvonplaten
Pull request: yes

This PR:
- adds a simple test for all tf models to verify that the forward function can be used in graph mode
- optimizes TF Longformer, by removing unnecessary calculation, such as `tf.transpose()` (In contrast to PyTorch, `tf.transpose()` allocates a new tensor and thus should be avoided). This also cleans up the code IMO.

=> These changes lead to a speed-up of 1.03 which is actually not that much...more details in benchmark below.

After a lot of digging TF XLA will not be very easy to implement as a lot of kernels that are highly used in this model `tf.where` are not implemented for XLA (yet). So TF Longformer TPU will not work sadly for the moment @ibeltagy 

### Conclusion

For me the PR was also a good exercise to see whether TF can significantly sped up by removing unnecessary tensor allocations. It seems like it's not really worth it go through all the tf models if the improvement in speed is only around 2,3%.

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

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

```diff
@@            Coverage Diff             @@
##           master    #6447      +/-   ##
==========================================
+ Coverage   78.96%   79.81%   +0.84%     
==========================================
  Files         157      157              
  Lines       28486    28479       -7     
==========================================
+ Hits        22495    22730     +235     
+ Misses       5991     5749     -242     
```


| [Impacted Files](https://codecov.io/gh/huggingface/transformers/pull/6447?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [src/transformers/modeling\_tf\_bert.py](https://codecov.io/gh/huggingface/transformers/pull/6447/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9iZXJ0LnB5) | `98.38% <100.00%> (ø)` | |
| [src/transformers/modeling\_tf\_electra.py](https://codecov.io/gh/huggingface/transformers/pull/6447/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9lbGVjdHJhLnB5) | `98.95% <100.00%> (+73.82%)` | :arrow_up: |
| [src/transformers/modeling\_tf\_longformer.py](https://codecov.io/gh/huggingface/transformers/pull/6447/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9sb25nZm9ybWVyLnB5) | `98.67% <100.00%> (-0.03%)` | :arrow_down: |
| [src/transformers/modeling\_tf\_xlm.py](https://codecov.io/gh/huggingface/transformers/pull/6447/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl94bG0ucHk=) | `18.94% <0.00%> (-74.32%)` | :arrow_down: |
| [src/transformers/modeling\_tf\_flaubert.py](https://codecov.io/gh/huggingface/transformers/pull/6447/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl9mbGF1YmVydC5weQ==) | `24.53% <0.00%> (-63.81%)` | :arrow_down: |
| [src/transformers/modeling\_roberta.py](https://codecov.io/gh/huggingface/transformers/pull/6447/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ19yb2JlcnRhLnB5) | `77.37% <0.00%> (-19.71%)` | :arrow_down: |
| [src/transformers/file\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/6447/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9maWxlX3V0aWxzLnB5) | `82.41% <0.00%> (-0.26%)` | :arrow_down: |
| [src/transformers/modeling\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/6447/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ191dGlscy5weQ==) | `88.05% <0.00%> (+0.55%)` | :arrow_up: |
| [src/transformers/modeling\_tf\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/6447/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9tb2RlbGluZ190Zl91dGlscy5weQ==) | `87.29% <0.00%> (+2.60%)` | :arrow_up: |
| [src/transformers/generation\_tf\_utils.py](https://codecov.io/gh/huggingface/transformers/pull/6447/diff?src=pr&el=tree#diff-c3JjL3RyYW5zZm9ybWVycy9nZW5lcmF0aW9uX3RmX3V0aWxzLnB5) | `85.71% <0.00%> (+2.75%)` | :arrow_up: |
| ... and [4 more](https://codecov.io/gh/huggingface/transformers/pull/6447/diff?src=pr&el=tree-more) | |

------

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


### Speed Benchmarking:

Running this command on the master branch: 
```
python examples/benchmarking/run_benchmark.py --models allenai/longformer-base-4096 --no_memory --sequence_length 512 1024
```

on this env:
```
- transformers_version: 3.0.2
- framework: TensorFlow
- eager_mode: False
- use_xla: False
- framework_version: 2.2.0
- python_version: 3.8.5
- system: Linux
- cpu: x86_64
- architecture: 64bit
- date: 2020-08-14
- time: 10:32:09.525696
- fp16: False
- use_multiprocessing: True
- only_pretrain_model: False
- cpu_ram_mb: N/A
- use_gpu: True
- num_gpus: 1
- gpu: TITAN RTX
- gpu_ram_mb: 24217
- gpu_power_watts: 280.0
- gpu_performance_state: 0
- use_tpu: False
```

gives:
```
====================       INFERENCE - SPEED - RESULT       ====================                         
--------------------------------------------------------------------------------                         
          Model Name             Batch Size     Seq Length     Time in s                                 
--------------------------------------------------------------------------------                         
 allenai/longformer-base-4096        8              512            0.229                                 
 allenai/longformer-base-4096        8              1024           0.463                                 
--------------------------------------------------------------------------------
```


On this branch the speed is improved to:
```
====================       INFERENCE - SPEED - RESULT       ====================                         
--------------------------------------------------------------------------------                         
          Model Name             Batch Size     Seq Length     Time in s                                 
--------------------------------------------------------------------------------                         
 allenai/longformer-base-4096        8              512            0.223                                 
 allenai/longformer-base-4096        8              1024           0.447                                 
--------------------------------------------------------------------------------
```
So we can see an improvement of ca. 3%, which is not that much actually... I guess it's interesting to see what effect removing some unnecessary `tf.transpose()` has in TF, but it might not be worth to go through all `modeling_tf_...` files trying to remove `tf.transpose()` and similar functions.