Mixed Precision Training

Why Precision Became a Bottleneck

For two decades, deep-learning math was synonymous with a single number format: IEEE 754 single-precision, also called FP32. Every weight, every activation, every gradient was a 32-bit real number. It was simple, numerically safe, and — for AlexNet-sized networks on 2012 hardware — entirely practical.

Then the models grew. GPT-3 has 175 billion parameters. In FP32 the weights alone are $175 \times 10^9 \times 4 \text{ B} = 700 \text{ GB}$, four A100-80GB GPUs full — just to store the model. During training you also need gradients (another 700 GB) and the Adam optimizer state (two FP32 buffers per parameter, another 1.4 TB). We are now past 2.8 TB before any activation has been computed.

At the same time NVIDIA's tensor cores, starting with Volta in 2017, began to offer enormous speed-ups for narrower number formats: roughly 8× for FP16 over FP32 on V100, and 16× on A100. The question was no longer whether to drop precision but how. Mixed precision training is the discipline that answers it without blowing up the numerics.

Anatomy of a Floating-Point Number

An IEEE-754 float is three fields packed into a bit string: a sign bit $s$, an exponent field $e$, and a mantissa field $m$. The value it encodes is $x = (-1)^{s} \cdot (1 + m / 2^{k}) \cdot 2^{e - b}$, where $k$ is the number of mantissa bits and $b$ is the exponent bias. The exponent controls the range (how big or small a number can be), and the mantissa controls the precision (how finely numbers are spaced). The three formats that matter for neural networks split that 32-bit budget differently:

The key insight is that FP16 and BF16 are not the same 16-bit format — they chose opposite sides of the range-vs-precision trade. FP16 kept FP32's precision philosophy and paid the price in range. BF16 kept FP32's range (same 8 exponent bits) and paid the price in precision.

Interactive: The Same Number in Four Formats

Type any number and see exactly which bits it occupies in each format, what value those bits actually decode back to, and where the encoding overflows or underflows. Use the preset buttons to jump to values that highlight specific regimes — for instance, 1e-8 is a typical tiny gradient that underflows FP16 but survives BF16.

The Great 16-Bit Debate: FP16 vs BF16

Before 2018 the default was FP16. Volta and Pascal GPUs had FP16 tensor cores and no native BF16. Mixed precision recipes were built around FP16 — including all the loss-scaling machinery we'll meet shortly.

Starting with Ampere (2020, A100), BF16 became a first-class citizen. And a quiet revolution happened: almost every large-model training run switched from FP16 to BF16. GPT-3, PaLM, Chinchilla, LLaMA, and Gemini were all trained primarily in BF16.

Why? Because gradients are the hardest tensor to keep numerically healthy, and gradients care more about range than precision:

• BF16's range is identical to FP32's. The exponent field is 8 bits in both. A BF16 cast of an FP32 gradient can never underflow or overflow — it can only get rounded.
• FP16's mantissa is precise but its range is tiny. A gradient below ~6×10⁻⁵ becomes subnormal and below ~6×10⁻⁸ vanishes to zero. A gradient above 65504 explodes to ±∞.
• BF16 precision is coarser, but optimizers like Adam effectively apply a running average that washes out per-step quantization noise. You lose a few bits of mantissa and nothing else.

Interactive: Representable Ranges

Rule of thumb: on Ampere-class hardware or newer, default to BF16. Reach for FP16 only when you are targeting older hardware (V100, T4) or have a specific case where FP16's extra precision demonstrably helps — and when you do, budget time for loss-scale tuning.

The Three Numerical Crises of Half-Precision Training

Naive half-precision training — "just cast everything to FP16" — fails in three well-catalogued ways. Each of the fixes we'll meet in the next section exists because of one of these failure modes.

Crisis 1 — Gradient Underflow

FP16 cannot represent positive values smaller than $2^{-24} \approx 5.96 \times 10^{-8}$ (its smallest positive subnormal). Any gradient below that is quantized to exactly zero. The corresponding weight receives no update at all, as if the layer were frozen.

The Micikevicius et al. (2018) Mixed Precision Training paper (ICLR) documents that for many deep networks a substantial fraction of gradient values — often a majority, depending on the architecture and training phase — falls below FP16's smallest representable positive value ($2^{-24} \approx 5.96 \times 10^{-8}$). Their Figure 2 plots the gradient histogram for an SSD-ResNet detector and shows that without scaling the bulk of the distribution is in FP16's underflow region.

Crisis 2 — Activation Overflow

The other end of the range cliff. FP16's maximum is 65504. A self-attention block computing $Q K^{T}$ over long sequences can produce pre-softmax scores in the thousands; squaring large activations in an MSE loss can exceed 65000. Once a value becomes $\pm\infty$, it poisons every downstream tensor (anything × ∞ = ∞; ∞ − ∞ = NaN).

This is why autocast's FP32-only list includes softmax, reductions, and most losses. Operations whose output can grow quadratically — matmul reductions, squares — stay in FP32 precisely to avoid overflow.

Crisis 3 — The Disappearing Weight Update

Suppose we fixed both underflow and overflow. A third, subtler issue remains. Consider a weight $W = 1.0$ and an Adam update $\Delta W = \eta \cdot \hat{m} = 5 \times 10^{-5}$.

The FP16 "gridpoints" near 1.0 are spaced by $2^{-10} \approx 9.77 \times 10^{-4}$. Any update smaller than half of that (≈ 5×10⁻⁴) rounds to the same bin as W itself. So storing $W + \Delta W$ in FP16 gives back exactly W — the update was computed correctly but then thrown away during the write-back.

To see this concretely, drag the "center x" slider and adjust the simulated update magnitude. Watch the green update tick turn red when it falls below half an ULP of FP16.

The Mixed-Precision Recipe

The complete recipe — first published in Micikevicius et al. 2018, Mixed Precision Training — has four ingredients. Each one directly addresses one of the three crises above.

1. FP32 master weights. The optimizer keeps a canonical FP32 copy of every parameter. Fixes Crisis 3.
2. FP16/BF16 forward & backward. At each step, the master weights are cast to the narrow dtype and used for compute. Cuts compute time and activation memory roughly in half.
3. Loss scaling. The loss is multiplied by a scalar $S$ before backward. By the chain rule, every gradient is scaled by $S$. Fixes Crisis 1. (Not needed with BF16.)
4. FP32 "safe list". Operations that are known to overflow or lose precision (softmax, layer norm, losses, reductions) are kept in FP32. Fixes Crisis 2.

Loss Scaling — Shifting the Histogram

Loss scaling is the most distinctive trick of FP16 training, and it is surprisingly simple. Let $L$ be the loss. Instead of backpropagating $L$, we backpropagate $L' = S \cdot L$. By linearity of differentiation, every gradient becomes $\frac{\partial L'}{\partial \theta} = S \cdot \frac{\partial L}{\partial \theta}$ — the gradient histogram is translated to the right on a log axis by $\log_2 S$ bits. Tiny gradients that would have fallen off FP16's lower cliff are now safely in the center of the band. After backprop, we divide each gradient by $S$ to get the true gradient back for the optimizer step.

In practice we never pick $S$ by hand. The dynamic loss scaler starts at a hopeful value (typically 2¹⁵) and adjusts automatically: whenever a gradient overflows to $\pm\infty$, it halves $S$ and skips the step; after many clean steps in a row, it doubles $S$ to stay aggressive.

Watch the feedback loop evolve step by step in the simulator below. Step through a training run and observe $S$ doubling after stretches of clean steps and halving the moment an overflow is injected.

Interactive: Loss Scaling in Action

Drag the slider. At $S = 1$ most of the synthetic gradient distribution falls below FP16's minimum and gets crushed to zero (red on the left). Push $S$ up to 2¹⁵ and essentially the whole distribution lands in the representable band. Push it too far and overflow kicks in on the right. This is exactly the feedback loop PyTorch's GradScaler runs on every step.

FP32 Master Weights

The optimizer owns an FP32 copy of every parameter. At the top of each step, it materializes an FP16 copy for the forward and backward pass. After backward, gradients are unscaled and applied to the master copy, not to the FP16 copy. The FP16 copy is then re-materialized for the next step and discarded.

Memory-wise this sounds expensive: we are carrying an extra FP32 tensor. In practice it is dominated by the Adam state, which was already FP32 and already 6× larger than the weights themselves (two moments × 4 B/param = 8 B/param). Mixed precision is mostly about activations and bandwidth, not parameters.

The Autocast Cast List

Rather than leave the user to tag every op, PyTorch ships a curated cast list. A non-exhaustive but representative split:

The pattern is consistent: ops whose output variance scales with input size (matmul reductions, convolutions) run narrow because their hardware kernels do the accumulation in wider precision internally. Ops that compute transcendentals or take log/exp of small quantities stay wide because a single FP16 intermediate would overflow or underflow.

From Scratch: Mixed Precision in NumPy

NumPy has no autograd, no tensor cores, and no GPU — but it has the exact same IEEE-754 FP16 we are trying to understand. That is all we need to reproduce every crisis and every fix by hand. Click any line to see what the computer is holding after that line.

1import numpy as np
2
3def to_fp16(x):
4    """Quantize x through FP16 and return as FP32 for arithmetic."""
5    return np.float16(x).astype(np.float32)
6
7# ---------- CRISIS 1 — underflow ----------
8grad_fp32 = np.float32(1.0e-8)
9grad_fp16 = to_fp16(grad_fp32)
10# grad_fp32 = 1e-08      grad_fp16 = 0.0   <-- disappeared
11
12# ---------- FIX 1 — loss scaling ----------
13S = np.float32(2**15)                # loss scale = 32768
14scaled    = to_fp16(grad_fp32 * S)   # ≈ 3.28e-4, now in FP16 range
15recovered = scaled / S               # unscale in FP32
16
17# ---------- CRISIS 2 — precision loss ----------
18W    = np.float32(1.0)
19lr   = np.float32(1e-3)
20grad = np.float32(5.0e-2)
21update = lr * grad                   # = 5e-5
22
23W_after_fp32 = W + update            # -> 1.00004995
24W_after_fp16 = np.float16(W + update)
25# cast to FP16 rounds back to exactly 1.0 -- the update was swallowed
26
27# ---------- FIX 2 — FP32 master weight + FP16 compute copy ----------
28W_master = np.float32(1.0)           # source of truth lives in FP32
29for step in range(3):
30    W_fp16   = to_fp16(W_master)     # forward/backward use this copy
31    W_master = W_master - lr * grad  # update done in FP32 (precise)
32
33# W_master after 3 steps: 0.99984998
34# Without master (FP16-only): 1.00000000  <-- no learning at all

Three numerical facts that drop out of the simulation:

• Underflow is total, not graceful. 1e-8 cast to FP16 is not "a tiny positive number" — it is exactly 0.0. Whatever downstream computation depended on it is dead.
• Loss scaling is (almost) lossless. Multiplying by 2¹⁵ and dividing back recovered the gradient within ~3×10⁻⁴ relative error. That is one FP16 rounding, not cumulative decay.
• Master-weight-free FP16 training is strictly broken. Our last print line shows that an FP16-only weight does not move at all over three steps. Full stop, this is not a tuning issue — the number system simply cannot hold the update.

Dynamic Loss Scaling — A From-Scratch Simulator

The GradScaler we relied on in the previous example is not magic — it is roughly 40 lines of Python. Below is a faithful port of PyTorch's growth/backoff algorithm in pure NumPy, followed by a 20-step simulation with one deliberate overflow at step 10 so you can watch $S$ double, halve, and recover in real numbers.

1import numpy as np
2
3class DynamicLossScaler:
4    """Minimal port of PyTorch's GradScaler.update() semantics."""
5
6    def __init__(self, init_scale=2**15, growth_factor=2.0,
7                 backoff_factor=0.5, growth_interval=2000):
8        self.S = float(init_scale)
9        self.growth_factor   = growth_factor
10        self.backoff_factor  = backoff_factor
11        self.growth_interval = growth_interval
12        self.clean_steps     = 0
13
14    def step(self, raw_grad):
15        # Would-be scaled gradient in FP16 range
16        scaled = raw_grad * self.S
17        # Detect overflow as PyTorch does — any non-finite element
18        overflow = not np.all(np.isfinite(scaled))
19        # Also check FP16's representable max
20        if np.any(np.abs(scaled) > 65504):
21            overflow = True
22
23        if overflow:
24            self.S *= self.backoff_factor   # halve on overflow
25            self.clean_steps = 0
26            return None                     # signal: skip optimizer step
27        else:
28            self.clean_steps += 1
29            if self.clean_steps >= self.growth_interval:
30                self.S *= self.growth_factor
31                self.clean_steps = 0
32            return scaled / self.S          # recover true gradient in FP32
33
34# Simulate: 20 steps, one deliberate overflow at step 10
35np.random.seed(0)
36scaler = DynamicLossScaler(init_scale=2**15, growth_interval=5)
37for step in range(20):
38    g = np.random.randn(4).astype(np.float32) * 1e-4
39    if step == 10:
40        g *= 1e6                             # force overflow
41    out = scaler.step(g)
42    action = "skip" if out is None else "ok"
43    print(f"step {step:02d}  S={scaler.S:>10.0f}  clean={scaler.clean_steps}  {action}")

Now the PyTorch counterpart. The only new pieces are the import path and scaler.get_scale() for introspection — every argument of the constructor maps directly onto a field of our NumPy class, and scaler.update() is the exact method that implements the growth/ backoff logic we just traced by hand.

1import torch
2from torch.amp import GradScaler
3
4scaler = GradScaler(device='cuda', init_scale=2.**15,
5                    growth_factor=2.0, backoff_factor=0.5,
6                    growth_interval=2000)
7
8print("scale at start:", scaler.get_scale())
9# The skeleton of a step — see main AMP block above for the full loop:
10# with autocast(...): loss = ...
11# scaler.scale(loss).backward()
12# scaler.step(optimizer)
13# scaler.update()
14# print("scale after step:", scaler.get_scale())

The Same Thing in PyTorch: autocast + GradScaler

Production code doesn't do the casts by hand — PyTorch's AMP (automatic mixed precision) module wires everything up for us. The eight-line pattern below is the canonical training step you'll see in NanoGPT, LLaMA trainers, Megatron-LM, and most modern open-source transformer code. Every step in the skeleton maps directly to something we did by hand above.

1import torch
2import torch.nn as nn
3from torch.amp import autocast, GradScaler
4
5# ---------- Model, data, optimizer (all FP32 by default) ----------
6model     = nn.Linear(512, 512).cuda()
7optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
8loss_fn   = nn.MSELoss()
9
10# GradScaler — dynamic loss scaling controller
11scaler = GradScaler(device='cuda', init_scale=2.**15)
12
13# Synthetic batch
14x = torch.randn(64, 512, device='cuda')
15y = torch.randn(64, 512, device='cuda')
16
17# ---------- One training step with mixed precision ----------
18optimizer.zero_grad()
19
20with autocast(device_type='cuda', dtype=torch.float16):
21    out  = model(x)          # matmul runs in FP16 on tensor cores
22    loss = loss_fn(out, y)   # MSELoss output stays FP32 (safe list)
23
24scaler.scale(loss).backward()    # multiply loss by S, then backprop
25scaler.step(optimizer)           # unscale grads, skip step if inf/NaN
26scaler.update()                  # adjust S for next step (x2 or ÷2)
27
28print(loss.dtype, out.dtype, model.weight.dtype)
29# torch.float32   torch.float16   torch.float32

Interactive: Training Memory Footprint

How much memory do you actually save? The honest answer is "less than you think for weights, a lot for activations." Pick a model size and an activation budget and compare the four training recipes side-by-side. You can see why a 7B-parameter model fits in an 80-GB GPU in BF16 but not in FP32.

Connections to Modern Transformer Systems

Every major efficiency trick in modern transformer architectures is, directly or indirectly, a mixed-precision story. Here is how the concept ripples outward.

Flash Attention — Mixed Precision at the Kernel Level

Flash Attention (Dao et al., 2022) was the watershed optimization for long-context transformers. Its headline trick — tiling the attention computation to keep intermediates in on-chip SRAM instead of writing the full $N \times N$ attention matrix to HBM — is a memory story. Its quieter but equally important trick is mixed precision inside the kernel.

The inputs $Q, K, V$ arrive in FP16 or BF16. Inside the kernel, the matmul $Q K^{T}$ runs on tensor cores at that narrow dtype, but the running max and the softmax denominator are maintained in FP32 accumulators. This is Crisis 2 avoidance at the hardware level: if the softmax denominator were stored in FP16, a single sequence with pre-softmax scores above $\log(65504) \approx 11$ would overflow. Keeping the denominator in FP32 is why Flash Attention gets FP32-quality results at BF16 throughput.

Multi-Head Attention — Where Softmax Stays FP32

The canonical multi-head attention formula is $\text{head}_h = \text{softmax}\!\left(\frac{Q_h K_h^{T}}{\sqrt{d_k}}\right) V_h$. Look at where each dtype lives in a BF16 training setup:

• $Q_h, K_h, V_h$ — BF16 projections of the input (matmul runs on tensor cores).
• $Q_h K_h^{T}$ — BF16 matmul with FP32 accumulator internally, result written back to BF16.
• $\sqrt{d_k}$ — an FP32 constant; division may be promoted.
• softmax — always FP32. The log-sum-exp trick plus the exp/sum would underflow or overflow in BF16's coarse 7-bit mantissa. This is a direct consequence of Crisis 2.
• $\text{softmax}(\cdot) \cdot V_h$ — back to BF16 matmul.

That softmax-in-FP32 choice is exactly what PyTorch's autocast cast-list enforces for you (the LayerNorm that precedes the QKV projections — Section 1 — is what gives us the unit-variance input that makes this softmax-in-FP32 economical). You never wrote it; the safe list did.

Positional Encodings — A Quiet FP32 Holdout

Sinusoidal positional encodings are computed once at startup and added to the token embeddings as $PE(p, 2i) = \sin(p / 10000^{2i/d})$ and $PE(p, 2i+1) = \cos(p / 10000^{2i/d})$. The problematic piece is $10000^{2i/d}$: for $d = 12288$ (GPT-3's hidden size) and $i$ near the top, this exponent reaches $10000 \approx 2^{13.3}$. In FP16 that's close to the max representable value; in BF16 it is fine but the result rounds to ~1-bit precision. Production code computes $PE$ in FP32 and then casts to the activation dtype, so the precision loss only happens once, at the end.

Rotary (RoPE) and ALiBi encodings have the same discipline: compute the rotation/bias in FP32 and apply it to BF16 activations. The common pattern is "construct in the widest precision, apply in the narrowest".

KV-Cache — Inference Memory is a Precision Problem

Autoregressive inference caches the keys and values from every past token so each new token only needs one new attention step. For a model with $L$ layers and $H$ heads of size $d_k$, the KV-cache for a sequence of length $N$ stores $2 \cdot L \cdot H \cdot d_k \cdot N \cdot \text{bytes/elem}$. For LLaMA-70B with a 4K context that is roughly 80 GB in FP16 — larger than the weights themselves. Serving stacks therefore work very hard on KV-cache precision:

• FP16/BF16 is the starting point. Halves the cache compared to FP32.
• FP8 KV-cache is supported by TensorRT-LLM and vLLM on H100. Stores K and V in E4M3 or E5M2 with per-tensor scales. Cuts memory and bandwidth in half again, typically with < 0.5 point loss on downstream benchmarks.
• INT8 KV-cache (vLLM, SmoothQuant) goes further using asymmetric per-channel quantization — more complex but a common choice for cost-constrained serving.

The tradeoff is identical to training's: narrower storage means more tokens per GPU second, at the cost of small numerical drift that teams measure with perplexity and task-level evals.

FP8 and the Scaling Frontier

Hopper (H100, 2022) introduced FP8 tensor cores and two formats: $\text{E4M3}$ (4 exponent, 3 mantissa, max ≈ 448) for forward activations, and $\text{E5M2}$ (5 exponent, 2 mantissa, max ≈ 57344) for gradients, where range matters most. NVIDIA's Transformer Engine library automates the casts and manages per-tensor scale factors the way GradScaler manages a single global scale.

With only 3 mantissa bits in E4M3 the ULP at 1.0 is 0.125 — an enormous precision hit. FP8 therefore demands per-tensor scaling: each tensor gets its own floating-point multiplier, chosen so that the tensor's actual dynamic range lands in the center of FP8's band. The scales themselves are FP32 and computed from recent tensor statistics.

For example, if an activation tensor's observed amax (max absolute value) is 6.4 and we are targeting FP8 E4M3 (max representable ≈ 448), the per-tensor scale factor is $s = 448 / 6.4 = 70$. We store $s \cdot x$ in FP8 and remember $s$ (as FP32) alongside the tensor; at read-time we divide by $s$. NVIDIA's Transformer Engine updates $s$ from the observed amax on every forward pass — an online EMA very similar to BatchNorm's running statistics, but applied per-tensor per-step. The cost is one extra FP32 scale per tensor (4 bytes) and one division per read — trivial compared to the memory and bandwidth saved by storing the bulk tensor in 1-byte FP8.

The upshot: every generation of the mixed-precision story is the same pattern at finer granularity. FP32 training used one precision everywhere. FP16 AMP introduced one global scale (the GradScaler). FP8 training introduces per-tensor scales. The logical endpoint — per-group scales like those used by MXFP8 and fine-grained INT4 quantization in inference — is already shipping.

Tradeoffs at a Glance

Summary

Mixed precision training is the reason a 7-billion-parameter model fits on one GPU and a 175-billion-parameter model fits on a single cluster. But it is not a magic switch — it is a careful division of labor between number formats.

• Compute narrow, store wide. FP16/BF16 on the tensor cores for speed, FP32 master weights and optimizer state for correctness.
• Use loss scaling when range is tight. FP16's 5-bit exponent can't cover typical gradient magnitudes — shift the histogram with $S$. Skip loss scaling with BF16 because its 8-bit exponent already covers FP32's range.
• Keep softmax, norms, and losses in FP32. These ops overflow or lose precision in half formats. The autocast cast-list encodes the hard-won knowledge of which ops are safe where.
• Every modern transformer optimization is a precision story. Flash Attention accumulates softmax in FP32. Multi-head attention keeps its softmax in FP32. KV-cache quantization is precision-traded-for-memory. FP8/Transformer Engine replaces one global scale with per-tensor scales.

If FP32 was the era of "just use floats," the modern era is the era of "use the narrowest representation that survives the math of your op, and not a bit narrower." Mastering mixed precision is mastering exactly that decision on every tensor in your network.