FA3 had some bugs

Researchers found two bugs in FlashAttention 3 (FA3) that caused gradient norm to grow a thousandfold and loss to end 0.2 nats above FP32 attention when pretraining a 450M-parameter transformer on 50B tokens, with training healthy for the first 25B tokens and no NaN appearing. One bug, carried over from FA2, is an fma rounding issue in the exp(x_i * scale - max_scaled) computation, and the second is that FA3 rounds dS to BF16 before multiplying it, so the attention-score gradient no longer sums to exactly zero. The issues surfaced because the team was using Muon rather than AdamW, whose weight decay and common QK-norm or QK-clipping practices had masked the problem.

For reasons best known to themselves OpenAI blessed the world with several hundred new mathematical proofs. This inspired elation of some mathematicians, despair of many others, and confusion of a great many people whose only encounter with zeta moments involved Sean Connery and a laser grid. Lean, the language the proofs are written in, originally started for proving software so you might think we were using it for that too? But, surprisingly, we are still finding bugs https://arxiv.org/abs/2609.34272 in possibly the most widely used kernel of the last few years, FlashAttention 3: When we used FlashAttention-3 to pretrain a 450M-parameter transformer on 50B tokens, however, we ran into a problem: training was healthy for 25B tokens, then the gradient norm grew a thousandfold and the loss ended 0.2 nats above FP32 attention, without a single NaN. There were two issues. One was an issue I remember from 2024 in FA2 https://github.com/pytorch/pytorch/issues/121558 . Horace’s comment on that issue sums it up very nicely: We’re computing exp x i scale - max scaled . Now, max scaled = max x i scale . The idea here is that before we take the exponent, we normalize everything down to 0 so that exp x doesn’t blow up massively. Now, for the largest value of x i , max scaled = x i scale . So, x i scale - max scaled == 0 , right? Unfortunately, no, due to fma. max scaled is actually equal to r ound fp32 x i scale . But in fma, x i scale never gets rounded So with fma, we are computing x i scale - round fp32 x i scale , which is not equal to 0 This was an opt-in thing for FA2, and just never fixed for FA3. Turns out though, there was another bug: dS , which is the gradient with respect to the attention scores, should sum to exactly zero: its expresses how the keys differ from each other, not where they are absolutely. FA3 rounds dS to BF16 before multiplying it, so it doesn’t quite sum to zero. Early in training that’s mostly noise, but as the keys get larger and attention spikier it can swamp the real gradient. Part of the reason this hadn’t come up was that applying QK-norm or doing QK-clipping is quite common, and AdamW’s weight decay tended to avoid the explosion too. Muon, however, is more exposed, and that’s what the team were using when they noticed this. Quasi-Riemann is one thing, but reduced precision floats? Quite another.