Mixed Precision Training and Numerical Stability

Practical rules for fp16, bf16, loss scaling, overflow, underflow, and stable training loops.

Machine Learning · Training

Mixed precision works because many neural-network operations tolerate lower precision while reductions, normalization, and optimizer state often need wider range. The goal is not to make every tensor smaller. The goal is to use the narrow type where it is safe and keep critical accumulations stable.

Floating-point limits

A floating-point format trades precision, exponent range, and storage. fp16 has limited exponent range. bf16 keeps an exponent range similar to fp32 but has fewer mantissa bits.

The practical distinction:

fp16 is more likely to overflow or underflow,
bf16 is less likely to overflow, but has coarser precision,
fp32 is still preferred for many reductions and optimizer states.

If a gradient component is below the smallest representable fp16 value, it can flush to zero and disappear from the update. If an activation or gradient exceeds the representable range, it can become inf and then poison later operations.

Loss scaling

Loss scaling multiplies the loss before backpropagation:

\[\tilde{L} = sL\]

The gradient is scaled too:

\[\nabla_\theta \tilde{L} = s\nabla_\theta L\]

Before the optimizer step, gradients are unscaled:

\[g = \frac{\tilde{g}}{s}\]

This preserves the intended update while moving small gradients into a representable range during backward computation. Dynamic scaling lowers $s$ when overflow is detected and raises it after stable steps.

Autocast boundaries

Autocast chooses operation dtypes. Matrix multiplies and convolutions often run well in lower precision. Softmax, cross entropy, normalization, sums, and other reductions are often safer in fp32.

A safe training skeleton is:

with torch.autocast(device_type="cuda"):
    output = model(input)
    loss = loss_fn(output, target)

scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()

Do not manually call half() on every input and module when using autocast. That bypasses the operation-level policy and can move unstable operations into the wrong dtype.

Failure modes

NaNs are late evidence. The first failure may be an inf gradient several layers upstream, a loss scale that keeps backing off, or a normalization statistic that overflows.

bf16-pretrained models may not fit fp16 range. Running them with fp16 AMP can overflow even if the code is correct.

In-place operations and explicit out= tensors may bypass autocast eligibility in PyTorch. That can leave a hot path slower or a numerically sensitive path narrower than expected.

Practical checks

Log loss scale when using fp16 dynamic scaling.
Check for nonfinite loss, activations, and gradient norms.
Keep optimizer master weights or optimizer state in fp32 unless the optimizer is designed otherwise.
Compare a short fp32 run against mixed precision on the same seed and data order.
Disable autocast around custom reductions until they are tested.
Use bf16 first on hardware where it is fast and available.

Design guidance