[trainer] add `--optim adamw_torch_fused` for pt-2.0+
Author: stas00
Pull request: yes

This PR implement the discussion of https://github.com/huggingface/transformers/issues/22141 to 
1. add support for the fused version of torch's AdamW. via `--optim adamw_torch_fused`
2. due it being too new, untested and a known bug the fix of which didn't make it into pt-2.0.0 - I did not make `--optim adamw_torch_fused` the default for pt>=2.0 - but prepared a place holder to do that for pt-2.1 instead.
3. added an assert for `--fp16` as apparently it's buggy with fp16/AMP https://github.com/huggingface/transformers/issues/22141#issuecomment-1467013132 fixed in https://github.com/pytorch/pytorch/issues/95781 but didn't make it into pt-2.0 cut off - this should be fixed in pt-2.0.1. But should already work with pt-2.1.0 nightly - except I think it's broken still (reported on pytroch slack).

**Bottom line: for pt-2.0 `--optim adamw_torch_fused` will become available for any use except `--fp16` which will automatically be re-enabled upon pt-2.0.1 release, which probably will happen a month later. And we want to give `--optim adamw_torch_fused` time before making it a default.**

## Quality and Speed Comparison Benchmarks:

```
PYTHONPATH=src CUDA_VISIBLE_DEVICES=0 python scripts/benchmark/trainer-benchmark.py --base-cmd ' \
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 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 --fp16 \
' --target-metric-key train_samples_per_second --repeat-times 1 --variations '--optim adamw_torch_fused|--optim adamw_torch|--optim adamw_apex_fused' --report-metric-keys train_loss --base-variation '--optim adamw_torch'
```

and then adding `--bf16` and `--fp16` for the non-fp32 benchmarks.

### bf16 

| Variation                 |   Train<br>samples<br>per<br>second |   Diff<br>% |   Train<br>loss |
|:--------------------------|------------------------------------:|------------:|----------------:|
| --optim adamw_torch_fused |                              382.06 |           9 |            2.22 |
| --optim adamw_torch       |                              350.22 |           0 |            2.22 |
| --optim adamw_apex_fused  |                              386.81 |          10 |            2.22 |

### fp16

| Variation                 |   Train<br>samples<br>per<br>second |   Diff<br>% |   Train<br>loss |
|:--------------------------|------------------------------------:|------------:|----------------:|
| --optim adamw_torch_fused |                              389.41 |           0 |            2.66 |
| --optim adamw_torch       |                              389.37 |           0 |            2.55 |
| --optim adamw_apex_fused  |                              399.27 |           3 |            2.53 |

it's easy to see fp16 is broken - bad loss

### fp32

| Variation                 |   Train<br>samples<br>per<br>second |   Diff<br>% |   Train<br>loss |
|:--------------------------|------------------------------------:|------------:|----------------:|
| --optim adamw_torch_fused |                              107.98 |           3 |            2.21 |
| --optim adamw_torch       |                              105.14 |           0 |            2.21 |
| --optim adamw_apex_fused  |                              108.20 |           3 |            2.21 |


*** Setup:

```
Datetime    : 2023-03-13 15:17:47

Software:
transformers: 4.27.0.dev0
torch       : 2.1.0.dev20230312+cu117
cuda        : 11.7
python      : 3.8.16

Hardware:
1 GPUs      : NVIDIA A100 80GB PCIe, 79.21GB
```

Fixes: https://github.com/huggingface/transformers/issues/22141

Comments:
_The documentation is not available anymore as the PR was closed or merged._

Arg, indeed let's wait a bit to make it the default then and leave the default untouched for now.

I made additional fixes, please see the OP for details. Of course, I'm open to suggestions to handle these multiple nuances differently.

its give me the following error

`ValueError: --optim adamw_torch_fused with --fp16 requires PyTorch>2.0`

That's correct. the new fused version is broken for fp16/amp in pt-2.0. It's fixed in pt-nightly and will be fully available in 2.0.1 and/or 2.1.0.