[HF Trainer] [new optimizer] add `AnyPrecisionAdamW` (bf16)
### Feature request

pytorch just merged https://github.com/pytorch/torchdistx/pull/52, which adds  `AnyPrecisionAdamW` (bf16-support, and future new dtypes)

we should add it to our HF Trainer arsenal

This is open to the community - it shouldn't be too difficult to add by just checking the existing optimizers. Here are some pointers to start unraveling:

https://github.com/huggingface/transformers/blob/e88e9ff045347c9d92d85806a6987dc7ebcbdd5b/src/transformers/training_args.py#L393-L394
and
https://github.com/huggingface/transformers/blob/e88e9ff045347c9d92d85806a6987dc7ebcbdd5b/src/transformers/training_args.py#L94-L106

the key of course is the documentation and tests. checking the existing tests and working from there is what's needed.


One would start looking at mimicking the integration of other optimizers, 

So in this case it'd follow the path of `adamw_torch` , as it's the nearest similar optimizer.

it might help to look at the previous PRs that added new optimizers, e.g. find the PR that added `adamw_bnb_8bit` - that could be a good model to copy from. And you can see the scope of work that needs to be done. Except this one should be simpler than `adamw_bnb_8bit` as it just plugs in a core pytorch optimizer, that's why I said `adamw_torch` is another good model.

Please remember that this requires pytorch-nightly as this new feature hasn't made it yet into pytorch-1.13. So you will need to install it from https://pytorch.org/get-started/locally/ (Choose Preview (Nightly))

Thank you!

Hello, I'll like to be assigned to this issue. <|||||>Yes, please, @Zeesky-code - once you have a working PR please tag me there.

Thank you!<|||||>No longer working on this issue and it's now open to the community.
Thank you :)

> Hello, I'll like to be assigned to this issue.

<|||||>Hi, may I take it if it's available?<|||||>Yes, please. But please read the OP first and see if you understand what needs to be done. Thank you!<|||||>@stas00, where do I import `AnyPrecisionAdamW` from? I tried installing `pip3 install --pre torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/nightly/cpu`<|||||>Oh! my bad! it's not pt-nightly but `torchdistx`

```
$ git clone https://github.com/pytorch/torchdistx
$ cd torchdistx/
$ grep -Ir AnyPrecisionAdamW
src/optimizers/anyprecision_optimizer.py:# AnyPrecisionAdamW: a flexible precision AdamW optimizer
src/optimizers/anyprecision_optimizer.py:class AnyPrecisionAdamW(Optimizer):
src/optimizers/anyprecision_optimizer.py:                        "AnyPrecisionAdamW does not support sparse gradients"
```

We can probably try something like this:
```
try:
    from optimizers.anyprecision_optimizer import AnyPrecisionAdamW
except:
    raise ValueError("please install https://github.com/pytorch/torchdistx")
```<|||||>also please note that the import is about to move and once this PR is merged https://github.com/pytorch/torchdistx/pull/60 please update your clone. Thank you!

so it'll become:

```
from torchdistx.optimizers.anyprecision_optimizer import AnyPrecisionAdamW
```<|||||>Hi @stas00 and @atturaioe - 
Just wanted to drop in here to say thanks for integrating AnyPrecision!  
Also wanted to let you know I'm adding a bfloat16 check internal to the anyPrecisionAdamW optimizer as a safety mechanism, and working on the documentation as well. 
Please let me know if you hit specific integration issues or questions, but for now the very short documentation preview is that there are two primary use cases for AnyPrecision currently:
a - successfully training *entirely* in pure BF16 - you can do this b/c kahan summation ensures high precision updates to the weights. Without that you will hit 'weight stagnation' where BF16 can't keep up like FP32 over the training cycle.  This gives you faster training with lower memory requirements and generally meets or even exceeds full FP32 results (some regularization effect). 
Has been tested up to 11B param size models. 

Basic process:
~~~
# init model
my_model = build_model(config)

# move model to all BF16
my_model.to(torch.bfloat16)

# run AnyPrecision in pure BF16 with Kahan - pass in usual adamW options in ... below:
optimizer = AnyPrecisionAdamW(my_model.parameters(),..., 
                               momentum_dtype=torch.bfloat16, variance_dtype=torch.bfloat16, 
                               use_kahan_summation=True)  
~~~

b - Training with AdamW variance state in BF16 - results in both memory savings and training speed up, and in limited testing up to 1B, matches mixed precision results after second epoch.  (variance state (the variance of the variance) typically rapidly declines after second epoch or so, which was the intuition that fp32 precision probably is not needed after that).   
~~~
optimizer = AnyPrecisionAdamW(my_model.parameters(),..., momentum_dtype=torch.float32, variance_dtype=torch.bfloat16, use_kahan_summation=False)
~~~

Hope that the above is helpful for now.  Will have more detailed docs etc. this week but please let me know if any questions in the interim and thanks again for the integration work here!

<|||||>That's very helpful, Less - thank you for sharing these use cases and the details!

I will leave to @atturaioe the stage to ask questions as he has been performing the heavy lifting on this task.<|||||>I should add - setting momentum_dtype and variance_dtype to torch.float32 and use_kahan_summation=False, brings AnyPrecision to the traditional AdamW optimizer so you can quickly compare using BF16, pure or variance only, for your training. <|||||>> setting momentum_dtype and variance_dtype to torch.float32 and use_kahan_summation=False, brings AnyPrecision to the traditional AdamW optimizer so you can quickly compare using BF16, pure or variance only, for your training.

awesome, that would make a good quality test then.

Let's continue the discussion in the PR https://github.com/huggingface/transformers/pull/18961 so it's more "actionable" :)<|||||>@stas00 Hi! I'd like to pick up this issue if no one else is working on it at the moment.<|||||>Oh, thank you for bringing this up, @mollerup23  - and wanting to contribute! 

This has already been resolved in https://github.com/huggingface/transformers/pull/18961

We just forgot to close this issue.