WIP: FP8 train

[horsebridge]

A draft for FP8 training. It currently depends on fsdp2 and torchao to train with per-tensor FP8 quantization. To enable it, install torchao and set strategy=fsdp2 and fsdp_config.fp8=True. Note: it cannot ensure convergence now

[CLAassistant]

All committers have signed the CLA.

[Yzx835]

A draft for FP8 training. It currently depends on fsdp2 and torchao to train with per-tensor FP8 quantization. To enable it, install torchao and set strategy=fsdp2 and fsdp_config.fp8=True. Note: it cannot ensure convergence now

@horsebridge Hello, have you tested the throughput change in FP8 training? In my experiments, setting fsdp_config.fp8=True results in much slower speed.

[horsebridge]

@Yzx835 Yes, I also observed the slower speed. I think it is an issue in torchao.