FEAT / Trainer: Add adamw 4bit optimizer

[SunMarc]

What does this PR do ?

This PR adds the 4-bit optimizer from torchao library into HF Trainer. For now, it requires the main branch of torchao and torch >=2.3 (maybe we can wait a bit before merging). For those who wants to try, you can pass optim="adamw_torch_4bit" in TrainingArguments.

Since we already have the 8-bit optimizer from bnb that works well, i'm not adding it.

Related thread : https://x.com/marksaroufim/status/1809398186198593566

cc @muellerzr as you might be interested

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[muellerzr]

Nice! LG2M, cc @msaroufim :)

[msaroufim]

There's also an AdamWFp8 btw and it's the fastest one we've found when the HW supports it https://github.com/pytorch/ao/tree/main/torchao/prototype/low_bit_optim#benchmarks

Also cc @gau-nernst this is very exciting!

[amyeroberts]

Thanks for adding!

[SunMarc]

There's also an AdamWFp8 btw and it's the fastest one we've found when the HW supports it https://github.com/pytorch/ao/tree/main/torchao/prototype/low_bit_optim#benchmarks

Nice ! I'll add it in a separate PR !

[msaroufim]

Heads up @SunMarc we just released torchao 0.4! https://github.com/pytorch/ao/releases/tag/v0.4.0

[SunMarc]

Nice ! I'll merge it as soon as we merge the torchao quantization PR in transformers as there is some overlap !

[fzyzcjy]

Hi, is it OK if I PR to add the 8bit counterpart? (see #34893 for details) Thanks!