Stable FP4 training in Small Transformer Language Models

Abstract

Training transformers in NVFP4 can reduce memory use and computational cost, but low precision attention gradients can destabilize training and cause blowups from which the model does not recover. We identify one source of this instability as a violation of a fundamental constraint in the attention backward pass: each row of the gradient with respect to the attention logits (dS) should sum to zero. We introduce a simple projection that restores this constraint after low-precision computation. In controlled experiments with small Transformers, the projection enables stable training when attention gradients are computed in NVFP4, while the uncorrected NVPF4 baseline experiences irreversible loss blowups. Unlike methods that retain a high precision attention backward pass, our approach directly corrects the low-precision gradient. These results suggest that preserving exact mathematical constraints can make FP4 Transformer training more reliable.