Support for GrokAdamW

[mtasic85]

Hi there,

It would be great to have support for GrokAdamW optimizer but with low bit quantization.

You can check reference implementation: https://github.com/cognitivecomputations/grokadamw

It has shown promising results already.

[msaroufim]

Hi @mtasic85 for new optimizers generally our take is it's probably not too hard to implement them from scratch using the AdamW code as a reference implementation. If this is something you're interested in lmk we can walk you through how to do this

[p4arth]

Attempted to implement this here using the AdamW code as suggested by @msaroufim. Was able to get a working implementation, still need to check for correctness and optimizations if any.