Repository: pytorch
Issue: [RFC] APEX style fused optimizers in PyTorch
Labels: module: performance,module: optimizer,triaged,module: mta

<!--
If you have a question or would like help and support, please ask at our
[forums](https://discuss.pytorch.org/).

If you are submitting a feature request, please preface the title with [feature request].
If you are submitting a bug report, please fill in the following details.
-->

Having APEX style fused optimizers in pytorch by implementing, e.g., `_fused_sgd(Tensor[] params, Tensor[] grads, Tensor[] momentums, ...)` and expose this function via `torch.optim.fused.FusedSGD`.

# Motivation
We've developed `foreach` and `_multi_tensor` optimizers which use `foreach` functions to reduce the overhead of 
calling a bunch of small CUDA kernels.
As `foreach` handles one or more lists of tensors and optional scalar or a list of scalars, for dense contiguous tensors,
they are more efficient than the native optimizers (`torch.optim` optimizers) by the factor of the number of parameters in a model **without** much loss of flexibility and hackability (see https://github.com/pytorch/pytorch/blob/master/torch/optim/_multi_tensor/_functional.py for the implementation).

If we're sure that tensors involved in parameter update suffice the following conditions, [APEX style fused optimizers](https://github.com/NVIDIA/apex/tree/master/apex/optimizers)
should be more efficient than the `_multi_tensor` optimizers because each fused optimizer kernel takes list of tensors and do all the required math in it.

conditions
- tensors are dense and contiguous
- tensors are on the same device
- tensors have the same dtype

I'd love to know what you think.

# Benchmark

## huggingface/transformers
### Environment
- 2 x A100
- Container: `nvcr.io/nvidia/pytorch:22.02-py3`
- transformers source code https://github.com/crcrpar/transformers/tree/4.17.0-mtoptim.

### GPT2-medium
base command is `python examples/pytorch/language-modeling/run_clm.py --model_name_or_path gpt2-medium --dataset_name wikitext --dataset_config_name wikitext-2-raw-v1 --logging_strategy no --save_strategy no --do_train --per_device_train_batch_size 4 --num_train_epochs 1 --warmup_steps 8 --block_size 512 --fp16`.


| Optimizer     | samples/sec |
| :---        |    ---: |
| `torch.optim.AdamW`                        | 21.268 |
| `torch.optim.multi_tensor.AdamW`   | 21.289 |
| `apex.optimizers.FusedAdam`           | 22.325 |

### T5-base
base command is `python examples/pytorch/translation/run_translation.py --model_name_or_path t5-base --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 --source_lang en --target_lang ro --dataset_name wmt16 --dataset_config ro-en --source_prefix "translate English to Romanian:"  --warmup_steps 50  --dataloader_num_workers 2  --max_train_samples 40000 --fp16`

| Optimizer | samples/sec |
| :---------- | -------------: |
| `torch.optim.AdamW`                       | 202.8     |
| `torch.optim._multi_tensor.AdamW` | 203.282 |
| `apex.optimizers.FusedAdam`           | 210.165 |


## NVIDIA/DeepLearningExamples BERT

Goal: See the performance benefit of taking `torch.cuda.amp.GradScaler` as an argument of `Optimizer.step`.

Currently APEX has two implementations of LAMB optimizer, one is [`apex.optimizers.FusedLAMB`](https://github.com/NVIDIA/apex/blob/master/apex/optimizers/fused_lamb.py) and the other is [`apex.optimizers.FusedMixedPrecisionLamb`](https://github.com/NVIDIA/apex/pull/1237). The latter can take `torch.cuda.amp.GradScaler` as an argument and unscale gradients inside the CUDA kernel as the optimizer class has the attribute of `_step_supports_amp_scaling`.
<!-- 
(which lets `scaler.step(optimizer)` select [this path](https://github.com/pytorch/pytorch/blob/954c7e2a770f73377f753f9a80f60dfddc5b1da2/torch/cuda/amp/grad_scaler.py#L324-L331)
-->

### Environment
- code https://github.com/NVIDIA/DeepLearningExamples/tree/f3dbf8a69522d69c63c4508769bd8137658786a1/PyTorch/LanguageModeling/BERT
- container: nvcr.io/nvidia/pytorch:22.02-py3
- A100 (40GB) x 8

### Settings
- 1000 steps
- batch size 32

### Results

| optimizer | GradScaler compatibility | sequences/sec |
|:-----------:|:-----------------------------:|:------------------:|
| `FusedLAMB`                      | ❌ | 887.864 |
| `FusedMixedPrecisionLamb` | ✔️ | 893.424 |

I also compared how long `scaler.step(optimizer)` itself takes on one A100 using nvtx ranges:

| optimizer | duration (CPU) | duration (CUDA) |
|:-----------:|------------------:|--------------------:|
| `FusedLAMB`                       | 17.444 ms | 17.242 ms |
| `FusedMixedPrecisionLamb` |   8.668 ms | 15.587 ms |



<!--
### Pros & Cons of `FusedMixedPrecisionLamb`
- allows for apex.amp O2 style using `torch.cuda.amp.autocast` and `torch.cuda.amp.GradScaler`
- 
With this optimizer, NVIDIA/DeepLearningExamples BERT implements sync free, APEX O2 style AMP training where the optimizer updates FP16 `nn.Module`'s parameters using FP32 param copy.
- maintains `lr` and `step` as a device `torch.Tensor` -> to implement a sync free 
-->

<!--
```
python3 -m torch.distributed.launch --nproc_per_node=8 /workspace/bert/run_pretraining.py --input_dir=/workspace/bert/data/lddl/pretrain/phase1/unbinned/parquet/ --output_dir=/workspace/bert/results/gc_new_fused_lamb_no_lrsched_bs32_seq128_second --config_file=bert_configs/large.json --vocab_file=vocab/vocab --train_batch_size=32 --max_seq_length=128 --max_predictions_per_seq=20 --max_steps=1000 --warmup_proportion=0.2843 --num_steps_per_checkpoint=20000 --learning_rate=6e-3 --seed=12439 --fp16 --gradient_accumulation_steps=1 --allreduce_post_accumulation --do_train --json-summary /workspace/bert/results/dllogger.json --disable_progress_bar --num_workers=4 --disable_jit_fusions
```

### `apex.optimizers.FusedLAMB`

e2e_train_time : 311.09122371673584  training_sequences_per_second : 887.8643095044812  final_loss : 8.012489318847656  raw_train_time : 288.04401445388794

### `apex.optimizers.FusedMixedPrecisionLAMB`

e2e_train_time : 306.36698150634766  training_sequences_per_second : 893.4239103300664  final_loss : 7.99187707901001  raw_train_time : 286.82464957237244

```
python3 -m torch.distributed.launch --nproc_per_node=8 /workspace/bert/run_pretraining.py --input_dir=/workspace/bert/data/lddl/pretrain/phase1/unbinned/parquet/ --output_dir=/workspace/bert/results/gc_new_fused_lamb_no_lrsched_bs32_seq128_second --config_file=bert_configs/large.json --vocab_file=vocab/vocab --train_batch_size=32 --max_seq_length=128 --max_predictions_per_seq=20 --max_steps=1000 --warmup_proportion=0.2843 --num_steps_per_checkpoint=20000 --learning_rate=6e-3 --seed=12439 --fp16 --gradient_accumulation_steps=1 --allreduce_post_accumulation --do_train --json-summary /workspace/bert/results/dllogger.json --disable_progress_bar --num_workers=4 --disable_jit_fusions
```
-->

---

rel: https://github.com/pytorch/pytorch/issues/38655, https://github.com/pytorch/pytorch/issues/58833,https://github.com/pytorch/pytorch/issues/71274

cc @VitalyFedyunin @ngimel @vincentqb @jbschlosser @albanD @crcrpar @mcarilli @ptrblck