Backward-State Policy Is Part of the Learning Algorithm

2026-09-30

AI Summary

This study addresses the frequently overlooked backward-pass state strategies in low-precision training, which introduce gradient computation errors undetectable by conventional checks. We demonstrate that backward states are fundamental to learning algorithms and propose a specification that designates read values at each point of use. Furthermore, we transcend evaluation paradigms limited to copy precision and final loss by establishing a reference standard grounded in actual forward-pass gradients. Leveraging PyTorch, Transformer Engine, and FP8 emulation, we conduct single-operator verification and mean gradient consistency analysis to successfully predict and rectify gradient deviations caused by state reuse. All experimental results align with theoretical expectations, effectively ensuring the stability of FP8 training.

Abstract

Low-precision training rounds tensors that the backward pass reads again, often for several gradients; each use can read the forward's rounded value, the original, or a new random rounding. This backward-state policy looks like a memory and precision detail, settled by copy accuracy and final loss. We argue that it is part of the learning algorithm, and that neither check shows whether it is right. Copy accuracy does not decide the outcome: in three pairs of 390M runs with an emulated FP8 backward, training fails when attention's backward reuses the forward's rounded output and succeeds with a new rounding from the same distribution. Even the most accurate copy, the original itself, can be wrong by our reference: the gradient of the forward pass as it actually ran, with gradients passed through rounding unchanged. For example, a normalization output stored in low precision feeds two gradients: the gain's gradient needs the original, but the next layer's weight gradient needs the rounded value that layer multiplied. Final loss, the other check, does not rule out the error of reading the original for both: it persists in models trained with such a store, while planned loss comparisons stay within a margin fixed in advance. We therefore derive from this reference which value each use must read, or which substitute gives the same gradient on average with the forward held fixed, and check these per-use requirements on single operators, without training. In three tests using PyTorch and Transformer Engine, the requirements predicted beforehand whether reuse changes what the backward computes on average relative to an independent copy, and every prediction held. Backward-state policy is thus part of the learning algorithm: it should be specified and checked use by use, not settled by copy accuracy and final loss.

Problem

• low-precision training
• backward-state policy
• gradient computation
• stochastic rounding
• FP8

Innovation

• Low-precision training
• Backward-state policy
• Stochastic rounding
• Gradient correctness
• FP8