question of blockwise quant fp8 training

[rakkit]

Hi, the blockwise_fp8_training has been there for a while. Is there any reason we dont merge it into float8 folder?

And current moe training only supports FP8_ROWWISE and MXFP8, will FP8_BlockWise be considered to be added into torchao in the near future? (mainly for h100 users)

thanks!

[danielvegamyhre]

Hi @rakkit, some notes about the the fp8 blockwise prototype:

• It only applies to dense models (linear layers), not MoEs yet
• It achieves ~10% TPS improvement vs bf16 in e2e training with torchtitan with PoC integration, so more perf optimization work is necessary.
• We definitely want to add fp8 blockwise support to MoEs and achieve a much higher speedup. I am the only person on the team working full time on training at the moment, and I've been prioritizing mxfp8 for accelerating MOEs on Blackwell GPUs, so it's just a matter of having time to do the work. If you or anyone else wants to contribute I will gladly review the PRs, or provide guidance on what we need to do, my suggestions for design/implementation, etc

[rakkit]

Hi @danielvegamyhre thanks a lot for the explanation.

It only applies to dense models (linear layers), not MoEs yet

Yes, that's why I am asking, cause other quant (row/tensor-wise) for linear layer already in float8 folder

We definitely want to add fp8 blockwise support to MoEs and achieve a much higher speedup. I am the only person on the team working full-time on training at the moment, and I've been prioritizing mxfp8 for accelerating MOEs on Blackwell GPUs, so it's just a matter of having time to do the work. If you or anyone else wants to contribute, I will gladly review the PRs, or provide guidance on what we need to do, my suggestions for design/implementation, etc

Thanks a lot for the explanation. I am happy to work smth there. To make it work in the real world and save the life of writing kernels, would it make sense to start by integrating the deepgemm first?

[danielvegamyhre]

I am happy to work smth there. To make it work in the real world and save the life of writing kernels, would it make sense to start by integrating the deepgemm first?

@rakkit Sure, I was planning to explore deepgemm for this as well. Feel free to work on prototyping this with some tests + benchmarks to see how the performance is.