text: 
  
  TITLE
  [t5/t0/mt5 models] faster/leaner custom layer norm
  
  COMMENTS
  6
  
  REACTIONS
  +1: 0
  -1: 0
  laugh: 0
  hooray: 0
  heart: 0
  rocket: 0
  eyes: 0
  
  BODY
  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 