[t5/t0/mt5 models] faster/leaner custom layer norm
Author: stas00
Pull request: yes
Labels: name: Performance

name: WIP

This PR dynamically replaces `T5LayerNorm` with `apex.normalization.FusedRMSNorm` when the latter is available (it was just [merged into `apex@master`](https://github.com/NVIDIA/apex/pull/1274)). This is similar to how we used to replace the slower at that time `torch.nn.LayerNorm` with `apex.normalization.FusedLayerNorm` until about a year ago for most models.

Unlike highly optimized `torch.nn.LayerNorm` the current "manual" implementation of `T5LayerNorm`, which is known as RMSNorm, is slow under mixed precision mainly due to the explicit up- and down-casting to meet the requirements of fp32 accumulate.

I first tried to rewrite the code to use the optimized pytorch functions but I didn't succeed to make things faster and requested the NVIDIA team to create `apex.normalization.FusedRMSNorm` fused kernel for us, which they kindly did. Huge thanks to @eqy and his team!

Note: For now just 1 model class was modified and once the reviewers are happy, I will replicate this to the rest of the `model_type=t5` models.

### Background info:

- Root MeanSquare Layer Normalization paper: https://arxiv.org/abs/1910.07467
- Here is a good reference for RMSNorm https://github.com/bzhangGo/rmsnorm/blob/master/rmsnorm_torch.py

### Benchmarks

Here are the benchmarks of before and after for 3 different sizes of t5 - it's a nice 6-11% speedup in the ensemble - the larger the model the smaller the speed up. 

I didn't compare `T5LayerNorm` vs `apex.normalization.FusedRMSNorm` directly as it's quite obvious that the latter is much faster to have such an impact on the throughput of the ensemble.

(RTX-3090, apex@master, pt-1.10.2, transformers@master as of 2022-02-04)

```
## t5-small

CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src python examples/pytorch/translation/run_translation.py --model_name_or_path t5-small \
--output_dir output_dir --do_train --label_smoothing 0.1 --logging_strategy no \
--save_strategy no --per_device_train_batch_size 32 --max_source_length 512 \
--max_target_length 512 --num_train_epochs 1 --overwrite_output_dir \
--source_lang en --target_lang ro --dataset_name wmt16 --dataset_config "ro-en" \
--source_prefix "translate English to Romanian: " --warmup_steps 50 \
--max_train_samples 20000 --dataloader_num_workers 2 --bf16


# w/o apex.normalization.FusedRMSNorm

  train_loss               =       2.52
  train_samples            =      20000
  train_samples_per_second =    420.204


# w/ apex.normalization.FusedRMSNorm

  train_loss               =     2.5205
  train_samples            =      20000
  train_samples_per_second =    470.624

Result: 11% speed up!

## t5-base

CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src python examples/pytorch/translation/run_translation.py \
--model_name_or_path t5-base \
--output_dir output_dir --do_train --label_smoothing 0.1 --logging_strategy no \
--save_strategy no --per_device_train_batch_size 16 --max_source_length 512 \
--max_target_length 512 --num_train_epochs 1 --overwrite_output_dir \
--source_lang en --target_lang ro --dataset_name wmt16 --dataset_config "ro-en" \
--source_prefix "translate English to Romanian: " --warmup_steps 50 \
--max_train_samples 10000 --dataloader_num_workers 2 \
--bf16

# w/o apex.normalization.FusedRMSNorm

  train_loss               =     2.2083
  train_samples            =      10000
  train_samples_per_second =    138.157

# w/ apex.normalization.FusedRMSNorm

  train_loss               =     2.2086
  train_samples            =      10000
  train_samples_per_second =    149.883

Result: 8% speed up!

## t5-large


CUDA_VISIBLE_DEVICES=0 PYTHONPATH=src python examples/pytorch/translation/run_translation.py \
--model_name_or_path t5-large \
--output_dir output_dir --do_train --label_smoothing 0.1 --logging_strategy no \
--save_strategy no --per_device_train_batch_size 8 --max_source_length 512 \
--max_target_length 512 --num_train_epochs 1 --overwrite_output_dir \
--source_lang en --target_lang ro --dataset_name wmt16 --dataset_config "ro-en" \
--source_prefix "translate English to Romanian: " --warmup_steps 50 \
--max_train_samples 5000 --dataloader_num_workers 2 --bf16

# w/o apex.normalization.FusedRMSNorm

  train_loss               =     2.0407
  train_samples            =       5000
  train_samples_per_second =     35.374

# w/ apex.normalization.FusedRMSNorm

  train_loss               =     2.0409
  train_samples            =       5000
  train_samples_per_second =     37.392

Result: 6% speed up

```

### Questions:

- docs - how should we document that installing apex should speed up any t5-based model
- the loss is ~2e-4 worse than the original for these short runs - is there a reason to be concerned?

@sgugger, @LysandreJik, @patil-suraj, @patrickvonplaten 

Comments:
Thanks for diving into it! Alongside Sylvain's comments, this LGTM, 

@patil-suraj, thank you for addressing my questions.

I suppose we then should replicate the doc snippet to all t5-based models, right? I'm thinking of adding a new Performance section - not sure if it's better to have it before the classes start so that it's visible or towards the end. I have added one - please have a look. (position in the doc and content please) 

Also do we want to do that copy-syncing mechanism where we designate one model the source and others the copy, but I think it only works for classes, right? i.e. no way to do the same here. I am asking so that the newly added t5-based models won't miss it.

PR looks good to me as well - that's a very nice speed-up @stas00 !

@LysandreJik, could you please have a quick look before I replicate this to many other t5-based models? Thank you!

Yes, please go ahead! Thank you, @stas00!

Actually, there is nothing to replicate, I thought there were derivatives/copy-cats but it appears that there are none.

So just going to merge this once CI completes.