low precision training upcoming feature tracker

[vkuzo]

This is a running list of planned features for low precision training. As features are completed we plan to delete them from this list, to keep things simple.

MX

pytorch/pytorch

• [in progress] improving torch.compile performance for to_mx cast across dim1: request for faster inductor kernels for blockwise reduction across dim1 -> write pytorch#149982

pytorch/torchao

• [in progress] performance tracker: MX single node performance tracker #1768

float8

performance

• [in progress] optimize torch.compile performance for float8 tensorwise scaling/casting kernels
  • [fixed behind a flag, off by default] Improve the Inductor generated kernel for the pattern of output1 = pointwise(intput); output2 = transpose(output1) pytorch#130015
  • [planned] [Inductor] Fusion of Tiled Point-Wise and Reduction Operators pytorch#128063
• [in progress] ensure that float8 rowwise scaling is performant with TP and async TP [Async TP] Fuse all-gather-matmuls for float8 rowwise training pytorch#149990

distributed

• [planned] verify integration with PP

new features

• [planned] MoE support: [roadmap/tracker] Low precision MoE training #2147
• [planned] float8 SDPA (priority TBD)

ecosystem

• [in progress] enable TP in torchtune's float8 integration: TP + FP8 + Compile: metadata error torchtune#2682

other

• [planned] expose float8 training via the quantize_ API
• [planned] migrate torchao.float8 code to torchao.quantization for better unification with the rest of torchao, in a BC-preserving way: [wip] unification of torchao.float8 with the rest of torchao #894

[airMeng]

Has SDPA supported FP8 already like flash-attention https://github.com/Dao-AILab/flash-attention/blob/main/hopper/benchmark_flash_attention_fp8.py does?

[vkuzo]

Has SDPA supported FP8 already like flash-attention https://github.com/Dao-AILab/flash-attention/blob/main/hopper/benchmark_flash_attention_fp8.py does?

we would be interested in accepting a contribution for this! To me, the priority is not clear as we have not seen convincing evidence that SDPA + float8 can achieve good accuracy numerically.