"Why Low-Precision Transformer Training Fails: An Analysis on Flash Attention"

[AustenMan]

I'm curious if the authors have seen this paper, and whether a toggle could be added for this stability patch they proposed?

https://arxiv.org/pdf/2510.04212

[tridao]

Thanks for sharing, I just skimmed it. Quite interesting and I'll need to read more carefully later.
The proposed method (Sec 4) doesn't look too hard to implement. We welcome contributions.
To modify the softmax computation:

• Cute-DSL version:

  flash-attention/flash_attn/cute/softmax.py

  Line 51 in 35384ec

  def online_softmax(

• CUDA C++ version:

  flash-attention/csrc/flash_attn/src/softmax.h

  Line 137 in 35384ec

  __forceinline__ __device__ void softmax_rescale_o(Tensor0 &acc_s, Tensor1 &acc_o, float softmax_scale_log2) {

[tridao]

Briefly looking at the code, looks like the authors are using a Pytorch reimplementation of FA, which I worry may not have the same numerical behavior as the FA implementation here. E.g. here we keep accumulator in fp32, use fused multiply add (fma) instead of separate mul and add, etc. Those details all affect numerical error.
https://github.com/ucker/why-low-precision-training-fails/blob/67811bd77435a5bcd6a5cb56ccd8a7be8d328a6b/attention.py#L22