Add AnyPrecisionAdamW optimizer
Author: atturaioe
Pull request: yes

# What does this PR do?
Add `AnyPrecisionAdamW` optimizer from `torchdistx`


Fixes # (issue)
#18827

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?


## Who can review?
@stas00 


Comments:
Hi, @stas00. I want to ask you whether should I add `anyprecision_adamw` specific arguments to the `trainings_args.py` or use the default ones in `trainer.py`. I'll be working on tests.

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

I'd say let's add a generic `--optim-args` optional arg which can then supply options to any future optimizer - i.e. it'd pair with `--optim`.

I'm trying to remember if we already have the plumbing for parsing in place - I think the `--debug` flag has it. edit: no, not that one. I remember writing it, but can't remember which one uses it. it's there somewhere - so many options.

But something like `--optim-args "key1:val1; key2:val2; ..."`

so here it'd be `--optim anyprecision_adamw --optim-args "use_kahan_summation=true; momentum_dtype=bf1oat16; ..."`

and we would convert any dtypes into actual `torch.foo` dtype using `getattr(torch, momentum_dtype)`

@atturaioe, this is just another variation - perhaps `--optim-args` can support just the exact syntax as python function sig?

```
--optim anyprecision_adamw --optim-args "use_kahan_summation=True, momentum_dtype=torch.bf1oat16; ..."
```

so `,` separator and perhaps writing out the dtypes exactly as they are in python and converting them on the fly to an actual class name. Same for booleans.

Perhaps it'd be easier to mimick the signature. Not sure. Let's see what you think is better.


Yeah! But should I parse the `--optim-args` into `dict` or something like that right in the `trainer.get_optimizer_cls_and_kwargs`?

Yes, that's exactly right:

https://github.com/huggingface/transformers/blob/d842f2d5b9bd4e361644c332bf9dc7f9b064f581/src/transformers/trainer.py#L1094


Is it any good? 
I didn't quite understand about converting dtypes on the fly(using `eval`?).

eval would be unsafe. here is a quick proof of concept:

```
python -c "import torch; x = 'torch.float16'; print(getattr(torch, x.split('.')[1]))"
```

Just pasting @lessw2020's comment from https://github.com/huggingface/transformers/pull/18961#discussion_r970160808 so that it doesn't get hidden by github once resolved and we will want to revisit this down the road and support other configs:

> 1 - For mixed precision - you could either
> a - run with the current defaults (M=fp32, Var = BF16, Kahan = False) and that would provide the memory and speed improvements from the Variance in BF16. That works nicely, and you can make that all work 'automatically' per above control options.
> b - you could also go all BF16 (M=BF16, Var = BF16, Kahan = False) because you will still get high precision weight updates with the master weights being in fp32. This is not as well tested though, but is something we are going to enable in FSDP soon by moving the working weight gradients to BF16, meaning you only have FP32 weights, nothing else.
> 
> To your question - having the weights in BF16 (via model.to) will only work if Kahan summation is active. If you don't run it with Kahan, then you are exactly right, you will hit weight stagnation and it will not be performant.
> The addition of Kahan is what makes it all work nicely.
> 
> Re: mark as experimental and tune as users run with it - that sounds like a great idea. I would just go ahead and use the current defaults then (M=FP32, Var = BF16, Kahan = False) as it's plug and play into FP32 or BF16 mixed precision.
> I'm working on a video tutorial now actually for this optimizer. Maybe we can add to the video once this PR is in, and show people how to run it with the manual change of model.to() and setting the defaults directly to get people comfortable with running in pure BF16.


https://github.com/pytorch/torchdistx/issues/68

@atturaioe, I'm back from vacation - what support do you need to finish this PR?

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_18961). All of your documentation changes will be reflected on that endpoint.

Hi @stas00, hope you had great time!
So the problem here that the `momentum_dtype` and `variance_dtype` set to different `dtypes` (`float32/bfloat16`) don't get cast dynamically, in the optimizer's `step()`, unless they're both of the same `dtype`. But of course I can set them both to same `dtype`, so the tests will pass.
Please correct me if I misunderstood something here.

Let's perhaps start with using the same dtype only and deal with that unusual case down the road should someone actually want to use it?

This commit changes the default params to the `float32` since there are 2 options for them to be the same dtype:
1 - all of them `float32`
2 - all of them `bfloat16` - won't pass tests since we have to move the `model.to(torch.bfloat16) ` while running tests

That's probably good enough as the initial integration. We can iterate to test the other variations once it becomes part of pytorch-core. 

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_18961). All of your documentation changes will be reflected on that endpoint.

ok, so as it has been awhile since this was created please rebase to main and flip the Draft mode to ready and we can then ask Sylvain to have a last look and merge. 

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_18961). All of your documentation changes will be reflected on that endpoint.

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_18961). All of your documentation changes will be reflected on that endpoint.

Thank you guys for helping/guiding me through this PR!