Why Low-Precision Transformer Training Fails: An Analysis on Flash Attention (2510.04212v2)

Abstract: The pursuit of computational efficiency has driven the adoption of low-precision formats for training transformer models. However, this progress is often hindered by notorious training instabilities. This paper provides the first mechanistic explanation for a long-standing and unresolved failure case where training with flash attention in low-precision settings leads to catastrophic loss explosion. Our in-depth analysis reveals that the failure is not a random artifact but caused by two intertwined phenomena: the emergence of similar low-rank representations within the attention mechanism and the compounding effect of biased rounding errors inherent in low-precision arithmetic. We demonstrate how these factors create a vicious cycle of error accumulation that corrupts weight updates, ultimately derailing the training dynamics. To validate our findings, we introduce a minimal modification to the flash attention that mitigates the bias in rounding errors. This simple change stabilizes the training process, confirming our analysis and offering a practical solution to this persistent problem. Code is available at https://github.com/ucker/why-low-precision-training-fails.

• The paper identifies that low-precision training causes numerical instabilities in flash attention due to biased rounding and low-rank representations.
• It uses dynamic normalization adjustments in softmax operations to stabilize training and minimize catastrophic rounding errors.
• Empirical results, including spectral norm and loss curve analyses, validate the proposed modifications for efficient transformer training.

Analysis of Low-Precision Transformer Training Failures

The paper "Why Low-Precision Transformer Training Fails: An Analysis on Flash Attention" explores the intrinsic failures associated with low-precision training of transformers, specifically focusing on flash attention mechanisms within BF16 precision settings. The authors identify the core issues leading to training instabilities and propose solutions to stabilize the process.

Introduction to Training Instabilities

Low-precision arithmetic, such as BF16 and FP8, is widely adopted to enhance computational efficiency by reducing memory usage and accelerating training speeds. However, these formats introduce numerical instability, particularly in attention mechanisms. Flash attention, which reduces memory complexity by processing inputs in blocks iteratively, is prone to catastrophic loss explosions when run in BF16 precision. This paper investigates the causal mechanisms of these failures and offers practical modifications to mitigate them.

Root Causes of Instability

Emergence of Low-Rank Representations

One identified source of instability is the emergence of similar low-rank representations across training steps, leading to biased weight updates. Low-rank structures in the attention scores lead to concentrated errors that compound over time, affecting spectral norms of the model parameters. These errors manifest as biased gradient updates, ultimately derailing the training process.

Figures and Examples:

Figure 1: $,$, and $( )[T]^\top [T]$ at different batch indices and training steps. (c) and (f) show that $( )[T]^\top [T]$ for different tokens and training steps have some similar columns in input features 546 and 678.

Biased Rounding Errors

Another significant source of failure is the biased rounding inherent in BF16 arithmetic operations. When aligning exponents for addition, inaccuracies occur due to the limited precision of BF16, especially during operations like $\bar{}$ in the softmax computation. The paper shows that specific values tend to become exactly 1 during computation, exacerbating the biased errors and contributing to the failure.

Figure 2: Spectral norm across layers and training steps.

Proposed Solutions

The authors propose a minimal yet effective modification to the flash attention mechanism to stabilize training. By dynamically adjusting the normalization factor in the softmax calculation, the introduced rounding errors are minimized. This change ensures that elements of the affected matrices remain less than 1, preventing the conditions that lead to catastrophic rounding errors.

Implementation Strategy

The modification involves checking for repeated maximum values during softmax operation and dynamically adjusting the normalization constant to ensure the maximum value of the normalized matrix $\bar{}$ is less than 1. This principled approach directly tackles the biased rounding error while retaining the benefits of flash attention's efficiency.

Figure 3: Analysis of $\bar{}$.

Conclusion

This paper successfully identifies and provides a mechanistic explanation for the failures in low-precision training of transformers utilizing flash attention. By addressing the root causes rooted in both the training dynamics and arithmetic precision, it offers a pathway for stable training. The proposed modifications, validated through empirical results, present a solid foundation for further research into improving the robustness of large-scale neural network training in low-precision settings. Future investigations could extend these findings to other architectures and precision formats.

Figure 4: Loss curves of two independent runs of GPT-2 training with flash attention in BF16 reported in the Github issue of nanoGPT.