[low-bit optim] Add coat for float8 optimizer

[MirMustafaAli]

This is a Work in Progress PR for #1190.

As a draft PR, I have followed the first piece of advice by @gau-nernst of "extending OptimStateFp8". Have created a separate Dynamic Range Function Instead of creating a different quantize_fp8 method as it will be applied before quantization to achieve larger representation range of float8 datatypes and the class will be storing value k to inverse the it after dequantization.

Requirements:
TBA
Additional Code/logic Added:
TBA
Logic/Code changes to existing codebase:
TBA
Outcome:
TBA
Scope of Usage:
TBA
Example
TBA

Changes:

• Dynamic Range Expansion Function: implementation of formula from the paper
• Created OptimStateFp8WithDynamicRangeExpansion class by extending OptimStateFp8: by referencing the implementation of the OptimStatefp8. I have only overridden the dequantize method
• Implemented aten.copy.default and aten.to_copy.default for OptimStateFp8WithDynamicRangeExpansion:

Benchmarks

Parameters

Parameter	Value
Learning Rate (lr)	0.0001
Automatic Mixed Precision (amp)	bf16
Seed	42
Model	timm/vit_base_patch16_224.augreg_in21k
Optimizer (optim)	AdamWFp8Ao_coat
Compile	False
Profile	False
Project	COAT-benchmarking
Number of Epochs	10
Run Name	AdamWFp8Ao_coat
Full BF16	False
Number of Workers	4
Batch Size	1024
Weight Decay	0
Channels Last	False
Optimizer CPU Offload	None
Cosine LR Scheduler	False
Checkpoint Activations	False

Results

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/ao/1231

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[gau-nernst]

I was thinking you can just add a flag to the current OptimStateFp8, something like dynamic_range_expansion: bool, instead of subclass-ing it.

[MirMustafaAli]

I was thinking you can just add a flag to the current OptimStateFp8, something like dynamic_range_expansion: bool, instead of subclass-ing it.

i have added the flag for optimstatefp8. could you verify its right?

[gau-nernst]

I think this requires a bit more work. You need to verify that you can create an optimizer with this (add test to https://github.com/pytorch/ao/blob/main/test/prototype/test_low_bit_optim.py) as well do some short training runs for sanity checks (using https://github.com/pytorch/ao/blob/main/benchmarks/benchmark_low_bit_adam.py).

I think for merging the PR, we should wait for the official code release to check numeric against them.

If you don't mind, we can discuss more details in GPU-MODE discord group https://discord.gg/gpumode. Just create a thread under torchao and tag me in (@gau.nernst)

[MirMustafaAli]

I understand the situation for merging the PR. Will be glad to work on working on this issue. creating thread in gpumode

[gau-nernst]

Thanks for the update. The PR is coming out nicely. There are some failing CI tests. Can you fix them, including the ruff linter?

Some extra items once that is finished:

• Update doc (link to the paper + usage)
• Run benchmark for sanity check https://github.com/pytorch/ao/blob/main/benchmarks/benchmark_low_bit_adam.py. I'm thinking comparing between BF16 baseline, FP8 optimizer, and FP8 COAT optimizer. Feel free to select a benchmark config suitable for you. And add the benchmark results in this PR description. Ideally, it should show that FP8 COAT is better than FP8 (though we might not observe it)