Mimic `adamw_torch_4bit` and have `adamw_torch_8bit`

[fzyzcjy]

Feature request

Hi thanks for the lib! Currently there is adamw_torch_4bit, but I hope to mimic it to have a adamw_torch_8bit that uses 8bit torchao adamw.

The reason is that, I would like to use deepspeed cpu offload for the optimizer, and also use 8bit adamw. However, the 8bit one in current hf transformers does not support cpu, so I need to use the torchao one.

Motivation

Your contribution

yes, willing to PR

[Rocketknight1]

cc @muellerzr for deepspeed/accelerate!

[muellerzr]

A PR for this would be great 🤗 cc @SunMarc

[fzyzcjy]

Thanks! I will do that later.

[SunMarc]

Feel free to add it ! Let me know if you need any help

[fzyzcjy]

Thanks! I will firstly mimic the 4bit one and see whether it works.

[fzyzcjy]

PR created: #34993