Question about the choice of use_fast_accum in FP8 Training

[haochengxi]

In the FP8 training setting, I noticed that use_fast_accum is only enabled during the forward pass:

ao/torchao/float8/config.py

Line 174 in 6243040

gemm_config_output: Float8GemmConfig = Float8GemmConfig(use_fast_accum=True)

When benchmarking performance, I've observed that enabling fast accumulation can lead to a 1.1x to 1.2x speedup. I'm curious—was there a specific reason or concern behind the decision to restrict use_fast_accum to only the forward pass?

[vkuzo]

Sorry about the late reply! The fast_accum flag controls the accumulation behavior of the gemm kernel, and enabling it can speed up the gemm at the cost of some accuracy. The default setting we have reflect a perf vs accuracy tradeoff that we have generally found optimal, but you can definitely try using fast_accum more liberally to bias more towards performace.

Note that the gradients are often more spiky in magnitude than activations and weights, so extra accuracy tends to be more important in the backward.