Why architecture matters here
The case for mixed precision is arithmetic intensity: a modern accelerator’s tensor cores deliver several times more FLOPs in bf16/fp16 than fp32, and at the memory wall, 2 bytes per element moves twice the tensor per second that 4 bytes does. For a model that costs millions of accelerator-hours, the format decision is a budget decision — the same loss curve at half the cost, or a larger model for the same cost.
Why not just flip everything to 16 bits? Because training error compounds where inference error does not. Consider a weight of magnitude 1.0 receiving a gradient update of 1e-4 with learning rate 1e-3: the update is 1e-7. In fp16, the gap between representable numbers near 1.0 is about 1e-3 — the update rounds to exactly zero, every step, forever. The weight is frozen while the loss curve looks merely “slow.” Meanwhile activations in attention can spike past 65,504 and become inf, one inf becomes NaN in the next softmax, and NaN propagates to every parameter in one backward pass. These are not tail risks; they are the default behavior of pure-fp16 transformer training.
Architecture matters because each failure has a targeted, nearly-free fix — master weights defeat update rounding, loss scaling defeats gradient underflow, fp32 islands defeat activation overflow — and the fixes compose into a recipe that is now so standard it hides behind one flag in PyTorch AMP or DeepSpeed. Understanding what the flag does is what lets you debug the run at 3 a.m. when the loss scale is collapsing and distinguish “numerics problem” from “learning-rate problem” — two incidents with identical symptoms and opposite remedies.
The architecture: every piece explained
FP32 master weights. The optimizer’s copy of every parameter stays in fp32, together with Adam’s first and second moments. Each step casts masters down to bf16/fp16 for compute; the update w ← w − lr·m̂/(√v̂+ε) is applied to the fp32 master, where a 1e-7 increment is representable. The cost is memory — 4 bytes master + 8 bytes moments per parameter, dwarfing the 2-byte compute copy — which is exactly why ZeRO and FSDP shard optimizer state first: it is the biggest, least-touched tensor in the system.
Loss scaling (fp16 only). Gradient distributions in deep transformers concentrate at magnitudes 1e-6–1e-3 — substantially below fp16’s normal floor. Multiplying the loss by a scale factor L (say 2¹⁶) multiplies every gradient by L via the chain rule, sliding the whole distribution into representable range; gradients are divided by L in fp32 before clipping and the optimizer step. Dynamic scaling automates the choice: on any overflow (an inf/nan in the unscaled gradients) skip the step and halve L; after ~2,000 clean steps, double it. The scale factor thus rides just under the overflow ceiling, maximizing headroom against underflow.
bf16 as the modern default. With fp32’s exponent range, bf16 needs no loss scaling — a whole failure class and its telemetry deleted. The price is mantissa: ~3 decimal digits versus fp16’s ~3.3 and fp32’s ~7.2. In practice gradient noise from batching exceeds bf16 rounding noise, so large-model pretraining standardized on bf16; fp16 persists on hardware without bf16 support and in some fine-tuning stacks.
Precision islands. Operations that reduce over many elements or exponentiate stay in fp32 regardless: softmax (exp overflows; the max-subtraction trick assumes precision), layernorm statistics (variance of thousands of elements), the loss itself, and often the residual stream and gradient all-reduce, whose long sums accumulate rounding linearly. Matmul accumulation happens in fp32 inside the tensor core by design. The compiled recipe — “autocast” — is precisely a per-op dispatch table implementing this taxonomy.
End-to-end flow
One fp16 training step, end to end. (1) The runtime casts fp32 master weights to fp16 compute copies (or reuses cached casts). (2) Forward pass: matmuls run on tensor cores — fp16 inputs, fp32 accumulators, fp16 outputs; autocast routes softmax and layernorm through fp32 islands; activations are stashed in fp16 for the backward pass, halving activation memory. (3) The loss is computed in fp32 and multiplied by the current scale, say L = 65,536. (4) Backward pass: every gradient emerges scaled by L. A raw gradient of 3e-7 — unrepresentable in fp16 — travels as 0.0197, comfortably normal. (5) Gradients are upcast to fp32, unscaled (divide by L), and scanned for inf/nan.
(6a) Clean step: global-norm clipping in fp32, Adam updates the fp32 masters, the good-step counter increments; at 2,000 it doubles L to 131,072 and resets. (6b) Overflow step: some activation spiked, a scaled gradient hit inf. The step is skipped — no optimizer update, no scheduler tick — L halves to 32,768, and training continues. One skipped step per few thousand is the mechanism working; the loss curve never sees it.
(7) In data-parallel training the all-reduce of gradients runs in fp32 (or bf16 with care): summing across hundreds of replicas is exactly the long-reduction shape where half-precision rounding compounds, and a biased gradient mean is a silent quality tax rather than a crash. (8) Checkpoints persist the fp32 masters, optimizer moments, and the current loss scale — resuming at the default scale of 65,536 after the run had settled at 8,192 buys an immediate overflow cascade. The bf16 variant of this step deletes stages 3, 5’s unscale, and 6b’s scale management entirely — the recipe collapses to masters + islands + fp32 reductions, which is why bf16 won.