FP8 training of LoRAs with torchao

[sayakpaul]

@bghira I was able to hack together a dirty script for FP8 LoRA training through torchao: https://gist.github.com/sayakpaul/2743810ed362cdc841b90956e2b155a1

But the results are not there yet. So, opening this discussion up for debugging this together. But some desired things already work:

• Wrapping the transformer into an FP8-compatible module.
• Injecting LoRA config.
• accelerator.prepare() on the transformer.

Training command and other necessary instructions are in the gist itself.

Cc @AmericanPresidentJimmyCarter who might be interested too.

[bghira]

what about el honcho - the big cheese - loading a checkpoint?

[bghira]

can it resume training?

[sayakpaul]

Yeah it can. But it needs to train first :D

[bghira]

[sayakpaul]

I think I miscommunicated it. I mean it is training but we need to ensure the correctness of training results.