Mixed precision training: autocasting, loss scaling and numerical limits

What it is

Mixed precision training combines floating-point formats within one training run. Expensive matrix operations may use FP16 or BF16, while sensitive operations retain FP32. Automatic mixed precision selects formats according to operation and device rules; it is not equivalent to converting every tensor to half precision. Benefits depend on the model, hardware and workload. [1][2]

How it works

Autocasting chooses operation precision. Loss scaling addresses a separate issue: very small FP16 gradients can become zero. Multiplying the loss by a scale enlarges the backward gradients; removing that scale before the optimizer update preserves their intended magnitude. A dynamic scaler detects non-finite gradients, skips affected optimizer updates and adjusts its scale. It cannot make every numerically unstable model valid. [1][2]

Practical workflow

Establish an FP32 reference with fixed data and evaluation rules. For ordinary eager PyTorch training, wrap the forward pass and loss computation in torch.autocast, then perform backward outside that context. With FP16, use torch.amp.GradScaler where appropriate. Do not manually convert the entire model merely to enable autocasting. [1]

If clipping gradients, unscale them before clipping. Call the scaler’s step and update methods in the documented order, then clear gradients for the next update. Save scaler state with model and optimizer checkpoints. [2] With accumulation, keep one scale across the whole effective batch. Compiled execution and custom operations need their own version-specific precision checks.

Illustrative example

Consider a classifier whose convolutions dominate runtime but whose final loss needs greater numerical range. Autocasting can run eligible convolutions in lower precision while evaluating a supported sensitive loss in FP32. If a step produces infinite gradients, the scaler can skip that update; this is a diagnostic event, not evidence that training has converged. The example describes behavior, not a measured speedup.

Limits and verification

FP16 and BF16 have different numerical ranges. A model trained for BF16 may overflow in FP16, whose largest finite value is 65,504. Increasing the loss scale cannot fix an overflowing forward computation. BF16 commonly avoids FP16-style gradient scaling, but still introduces rounding and requires validation. [1]

Compare validation metrics and loss trajectories with the FP32 reference, inspect non-finite values and skipped updates, and measure peak memory and throughput after warm-up. Small or input-bound models may gain little. [2] Report precision, hardware, software versions and observed differences. This article reports no local experiment.

Sources

[1] PyTorch: Automatic Mixed Precision package [2] PyTorch: Automatic Mixed Precision recipe