TF2 Mixed Precision, XLA, Distribution
Author: tlkh
Pull request: no

## 🚀 Feature

Hi there, I have benchmarked TF2 with the Transformers library. There are very positive results to be gained from the various TensorFlow 2.0 features:

- Automatic Mixed Precision (AMP)
- XLA compiler
- Distribution strategies (multi-GPU)

Here are the benefits (tested on CoLA, MRPC, SST-2):

- AMP: Between 1.4x to 1.6x decrease in overall time without change in batch size
- AMP+XLA: Up to 2.5x decrease in overall time on SST-2 (larger dataset)
- Distribution: Between 1.4x to 3.4x decrease in overall time on 4xV100
- Combined: Up to 5.7x decrease in overall training time, or 9.1x training throughput

Model quality (measured by validation accuracy) fluctuates slightly. Taking an average of 4 training runs for the single GPU results:

* CoLA: AMP results in slighter lower acc (0.820 vs 0.824) 
* MRPC: AMP results in lower acc (0.823 vs 0.835)
* SST-2: AMP results in slighter lower acc (0.918 vs 0.922)

However, with 4xV100 (4x batch size), interesting AMP can result in better results:

* CoLA: AMP results in higher acc (0.828 vs 0.812)
* MRPC: AMP results in lower acc (0.817 vs 0.827)
* SST-2: AMP results in slightly lower acc (0.926 vs 0.929)

The benchmark script demonstrating the use of this features, and also allowing you to test on your own system, is available [here](https://github.com/NVAITC/benchmarking/blob/master/tf2/bert_dist.py).

Note: on some tasks (e.g. MRPC), dataset is too small. Hence overhead of compiling model with XLA and doing distribution strategy does not speed things up. XLA compile time is also the reason why although throughput can increase a lot (e.g. 2.7x for single GPU), overall (end-to-end) training speed-up is not as fast (as low as 1.4x)

The benefits as seen on SST-2 (larger dataset) is much clear. 

All results can be seen at this [Google Sheet](https://docs.google.com/spreadsheets/d/1538MN224EzjbRL239sqSiUy6YY-rAjHyXhTzz_Zptls/).

## Motivation

I believe documentation and examples for this usage will be very useful to allow the community to train these models much faster on the hardware they might already have (V100/T4 on cloud, RTX GPUs in desktop etc.) 

If possible maybe I could be guided to contribute some examples or documentation!

## Additional context

External material:

* Benchmark Script: https://github.com/NVAITC/benchmarking/blob/master/tf2/bert_dist.py
* Benchmark Results: https://docs.google.com/spreadsheets/d/1538MN224EzjbRL239sqSiUy6YY-rAjHyXhTzz_Zptls/

Testing was performed on an NVIDIA DGX Station with 4x V100 (16GB) with NVLink.

This might also answer part of #1426 (Benchmark Script)

Comments:
Hi @tlkh, thank you for your work on the benchmarks!

We're planning to release some in-depths benchmarks by the end of the week/early next week. We'll add your work to it and we'll notify you once we have set-up an easier way to contribute benchmarks/examples!

This is really great @tlkh.
Do you think you could contribute an improved version of the `run_tf_glue` example with these best practices?
We could include your benchmarks and results in the examples readme.
Also, did you notice the same memory limitation mentioned in #1426?

@thomwolf 

Delighted to contribute! 

I haven't noticed the memory issues in #1426 on V100 (16GB) but I could see if I can replicate them on a Titan V (12GB).

Hey @tlkh as you've probably seen by now, we mentioned your work in the recent [Benchmarking blog post ](https://medium.com/huggingface/benchmarking-transformers-pytorch-and-tensorflow-e2917fb891c2) and added it to our [Benchmark section in our documentation](https://huggingface.co/transformers/benchmarks.html#tf2-with-mixed-precision-xla-distribution-tlkh). Thank you again for your work.