[low-bit optim] Add COAT optimizer

[gau-nernst]

Paper: https://arxiv.org/abs/2410.19313
Code: https://github.com/NVlabs/COAT (not available yet)

Seems like we already have most of the building blocks. The only new logic is "dynamic range expansion"

We can start implementing it first, then wait for the official code release for numeric checks.

[MirMustafaAli]

I would like to work on this @gau-nernst!!

[gau-nernst]

@MirMustafaAli Go ahead and submit a PR 😄. Let me know if you face any problems.

[MirMustafaAli]

@gau-nernst Which section should i look for to implement "dynamic range expansion"?. According to my understanding of repo it must be in float8 folder as it's aimed at hopper architecture utilizing type float8. Any pointers, PR's and reference methods which i can follow would be very much helpful for me.

[gau-nernst]

You can park it under torchao/prototype/low_bit_optim. The float8/ folder is more for training stuff (fp8 matmul).

You can extend our current OptimStateFp8. See https://github.com/pytorch/ao/blob/000a49026459dd1dadf5ca34322d98e7b1680250/torchao/prototype/low_bit_optim/subclass_fp8.py. I think we only need to change the quantize_fp8 function.

Another option is to create a separate optimizer. See https://github.com/pytorch/ao/blob/000a49026459dd1dadf5ca34322d98e7b1680250/torchao/prototype/low_bit_optim/adam.py. You can wrap all of the logic in a functional way (see single_param_adam() in the link above), and add boilerplate code for the optimizer (i.e. init optim states, call the functional optim step with torch.compile...)

[MirMustafaAli]

Thanks!! Will work on your advice.