[Trainer] fused torch `AdamW` is added in 2.0
Author: stas00
Pull request: no

### Feature request

there is a faster pytorch version of AdamW, it's the fused one. it was added in Feb-23 and will be part of pt-2.0.

`AdamW(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8…, fused=True).`

there is also `Adam` (not `AdamW`) which also has the fused version since pt-1.13 but we don't expose this one.

Now the question is this:

should we add `--optim adamw_fused_torch` and allow it only for pt-2.0+ or silently switch `--optim adamw_torch` to the fused version when pt-2.0+ is used?

cc: @sgugger 


Comments:
Good question! It might also revive the question of switching the optimizer from the HF implementation to the PyTorch one. Maybe we could add `adamw_fused_torch` as an option and then use for the default value of optim:
- `adamw_hf` on PyTorch < 2.0 (as before)
- `adamw_fused_torch` on PyTorch >= 2.0 so that users get the nice speed-up

What do you think?

You know me, I'm all for progress so I'd vote for your proposal +1.

The problem is that `adam_hf` != `adamw_*torch` algorithmically, so I will let you decide if you're OK with such a change of the default. I guess since `--adamw_hf` is still available - it's only a matter of communication to the community of the change of the default.

In a way perhaps pt-2.0 release allows us to change things as well.

-----------

Now practically let's implement `adamw_fused_torch`, then I can benchmark that it is faster while keeping loss not worse than `adamw_torch` and if all goes well then make it the default. Of course, it can be the same PR. or in 2 PRs

I can work on it, unless you prefer to do it.


Discussed with Lysandre and he's also fine changing the default with the PyTorch 2.0 release (we will highlight it in the release notes as a breaking change so users are aware). I think users upgrading to PyTorch 2.0 will expect some differences anyway. Also note that we have used the PyTorch AdamW in all the Accelerate example and no one raised an issue of convergence problems or differences with the Trainer examples in the past year.

As for the practicalities, go ahead with a PR if you want to do it. We just need to have it in main by tomorrow evening for the release :-) 

I would like to add from the PyTorch side that fused AdamW is still in its nascent stage and has had recent fixes regarding grad scaling interaction on float16s which unfortunately were too recent to be included in PT 2.0 (https://github.com/pytorch/pytorch/pull/95847). If most hf models will not be using auto-mixed precision, this may not be an issue, but I did want to call out the risk and add a +1 for the safer "add as an option for now so people can enroll for faster implementations, but don't make fused adamw the default yet".  

oh, thank you for the heads up, @janeyx99 - most models are using AMP. can the fix be pushed into 2.0.1?

so should we not allow its use if AMP is used?

edit: further discussion with Jane - only fp16/AMP is affected.

---------------------
OK, so @sgugger - we have to recall the plan on making it the default.

I already implemented the default change in https://github.com/huggingface/transformers/pull/22144 so roll back and make an issue to switch to it in pt-2.0.1 or pt-2.1.0? but perhaps it's a good idea to let it steep for a few months anyway.

and need to deal with f16/AMP too

The issue I'm talking about is https://github.com/pytorch/pytorch/issues/95781, and the fix has already landed last week and will go into whatever next release we have. And the repro given in the issue was used with fp16/amp and interacts with grad scaler. I agree it may be wise to let the implementation bake in, but it would be great to get people trying it out to harden its implementation.

Thanks for chiming in @janeyx99 . Let's wait a bit to change the default then!

ok, so https://github.com/huggingface/transformers/pull/22144 adds `--optim adamw_torch_fused` - but will assert on fp16/AMP w/ pt-2.0.0 and is already programmed to allow it for pt-2.0.1.

Making it a default will happen at pt-2.1.0 release or later.

I confirmed that the fp16 issue has been solved in today's nightly, the last fix was here https://github.com/pytorch/pytorch/pull/97415, discussed here https://github.com/pytorch/pytorch/issues/96755

As you can see the loss now matches for fp16/amp:

| Variation                 |   Train<br>samples<br>per<br>second |   Diff<br>% |   Train<br>loss |
|:--------------------------|------------------------------------:|------------:|----------------:|
| --optim adamw_torch_fused |                              387.10 |           3 |            2.66 |
| --optim adamw_torch       |                              377.61 |           0 |            2.66 |
| --optim adamw_apex_fused  |                              389.49 |           3 |            2.66 |

so our future proofing (allowing `--fp16` with `--adamw_torch_fused`) should work once pytorch makes a new release.