Precision and mixed precision

Problem

Matrix units on modern GPUs run several times faster on 16-bit and 8-bit inputs than on 32-bit ones, and half the bytes per value means half the memory traffic for every memory-bound operation (see how a GPU runs a model). Training entirely in a 16-bit format, though, fails in specific ways: small gradients round to zero, large activations overflow to infinity, and a weight update smaller than the gap between two neighbouring 16-bit numbers is lost entirely. Mixed precision keeps the speed by running the matmul-heavy parts in a low-precision format while keeping the parts that need range or accumulation in fp32. Doing that well needs a clear picture of what each format can and cannot represent.

Intuition

A floating-point number is scientific notation in base 2: a sign, an exponent that says where the binary point goes, and a mantissa that holds the significant digits. A fixed bit budget is split between the two. Exponent bits buy range, how large and how small a number can get before it overflows to infinity or underflows to zero. Mantissa bits buy precision, how many significant digits survive rounding at any scale. fp16 and bf16 both use 16 bits; fp16 spends them on precision (10 mantissa bits), bf16 on range (8 exponent bits, the same as fp32). That one choice decides almost everything else on this page.

Each square is one bit at the same width in every row. bf16 is fp32 with the low 16 mantissa bits dropped, so converting between them is a shift and a rounding step; tf32 is not a storage format at all, but the precision matrix units use when they are allowed to round fp32 inputs.

Mechanics

How a float is read

With e exponent bits and m mantissa bits, the stored exponent field is offset by a bias of 2^e−1 − 1 so that both tiny and huge numbers are representable. A normal number is read as

value = (−1)^sign × 1.mantissa₂ × 2^{field − bias}     spacing just above 1 = 2^−m (machine epsilon)

The all-ones exponent field is reserved for infinity and NaN in the IEEE formats, and the all-zeros field for zero and subnormal numbers, which drop the implicit leading 1 to fill the gap between the smallest normal number and zero with evenly spaced, less precise values. Between any two powers of two there are exactly 2^m representable numbers, so the relative spacing stays near 2^−m everywhere in the normal range while the absolute spacing doubles with each power of two. Converting a number into a format rounds it to the nearest representable value, ties to the one with an even last bit (round-to-nearest-even).

The formats

• fp32 (torch.float32): the default. About 7 decimal digits, range from about 10⁻³⁸ to 10³⁸.
• tf32: fp32's exponent with fp16's mantissa. NVIDIA matrix units since the Ampere generation can multiply fp32 inputs at this precision while still accumulating and storing fp32. PyTorch controls it with torch.set_float32_matmul_precision: "highest" (the default) keeps full fp32; "high" allows tf32 or a two-bf16 split; "medium" allows bf16.
• fp16 (torch.float16, IEEE half): 3 digits of precision but a range that stops at 65504 and, for normal numbers, 6.10e-05.
• bf16 (torch.bfloat16, "brain float"): fp32's range with 2 to 3 digits of precision. PyTorch converts fp32 to bf16 with round-to-nearest-even on the top 16 bits.
• fp8 e4m3 (torch.float8_e4m3fn): more precision, less range. The "fn" means finite: it has no infinities, gives the all-ones exponent back to ordinary numbers (keeping only S.1111.111 as NaN), and so reaches 448 instead of 240. PyTorch's conversion saturates: anything at or past the overflow threshold becomes ±448.
• fp8 e5m2 (torch.float8_e5m2): fp16's exponent with 2 mantissa bits, IEEE-like with infinities; overflow rounds to infinity.

Range versus precision

The two failure modes are different. A range failure turns a number into ±infinity (overflow) or 0 (underflow), which is catastrophic: one infinity in a gradient poisons the whole update. A precision failure perturbs a number by up to half a spacing, which is usually harmless noise in a neural network, except in one place: adding a small number to a large one. If the small number is less than half the spacing at the large one, the sum rounds back to the large number exactly. That is what happens to a weight update in a low-precision weight, and it is why the weights themselves stay fp32 (master weights, below). Training has values of very different scales in different places: activations and weights near 1, gradients that can be many orders of magnitude smaller, and sums of thousands of terms inside every matmul. Each needs a different trade.

Solid bars are the normal range, dashed bars the subnormals. The hollow dots are gradients fp16 cannot hold; loss scaling moves every gradient right by the same factor, and the whole set fits.

autocast and its op lists

torch.autocast(device_type, dtype=None, enabled=True, cache_enabled=None) is a context manager that changes which dtype operations run in, without changing the model. Inside it, each eligible op is intercepted at PyTorch's dispatcher (the Autocast dispatch key) and its floating-point CUDA inputs are cast according to a per-op cast policy. The default dtype is torch.float16 for CUDA and torch.bfloat16 for CPU. fp64 tensors are never cast, and ops on no list fall through and run in whatever dtype their inputs already have. The CUDA registration in autocast_mode.cpp sorts ops into five policies:

The pattern follows the formats. Matmuls and convolutions are compute-bound, run on matrix units that are fastest in 16 bits, and accumulate their dot products in fp32 internally, so they go low. Exponentials, logs, normalizations, reductions and losses can overflow or need to sum many terms accurately, so they go to fp32. One op is banned outright: binary_cross_entropy raises an error under CUDA autocast, because it usually follows a sigmoid whose output loses too much precision in fp16; the error points to binary_cross_entropy_with_logits, which fuses the two and is on the fp32 list.

autocast also keeps a weight cast cache: the first time an fp32 leaf tensor that requires grad (a parameter) is cast to the low dtype inside an autocast region, the cast copy is kept and reused by later ops in the same region, so a weight used by several matmuls is converted once per forward pass. The cache is cleared when the outermost autocast region exits. Its documentation also says to wrap only the forward pass and the loss, not backward(): backward ops run in the same dtype autocast chose for their forward ops anyway.

Loss scaling

In fp16, the gradients of a deep network often include values below 5.96e-08, fp16's smallest subnormal, which round to zero, and many more in the subnormal range where they keep only a few bits. Loss scaling multiplies the loss by a factor S before backward(). By the chain rule every gradient is multiplied by S too, which shifts the whole gradient distribution up into fp16's range (the second diagram). Before the optimizer uses them, the gradients are divided by S again, in fp32, where the small values are representable.

A fixed S is hard to choose: too small and gradients still underflow, too large and some overflow to infinity. GradScaler picks it dynamically. It starts at init_scale, by default 2¹⁶ = 65,536. Each iteration scaler.step(optimizer) unscales the gradients and checks them for infs and NaNs: if any are found, it skips optimizer.step() so the weights are untouched, and update() multiplies the scale by backoff_factor (0.5). After growth_interval (2000) consecutive iterations without infs it multiplies the scale by growth_factor (2.0). The scale settles near the largest value that does not overflow, and skipped steps become rare. unscale_ and update keep the inf flags and the scale on the GPU without a CPU-GPU sync; step reads the flags back with .item() to decide whether to call the optimizer.

bf16 needs none of this. Its exponent is fp32's, so every magnitude down to fp32's smallest normal number, 1.18e-38, is representable in bf16, far below where gradients live. What bf16 gives up is precision, and scaling cannot add mantissa bits. Its documentation for update() notes the converse risk: bf16-pretrained models are often incompatible with fp16 training, because their activations use a range fp16 does not have.

Bar height is the scale used at each step: green steps run the optimizer, amber steps overflow and are skipped. The dashed line is the largest scale at which 0.6 × scale still fits under fp16's 65504. With the default growth_interval of 2000 the same probe costs one skipped step in 2001 instead of one in 4.

Master weights and fp32 accumulation

Micikevicius et al. (2018) named the three techniques that made fp16 training match fp32 accuracy: loss scaling, an fp32 master copy of the weights, and accumulating fp16 products into fp32. The master copy exists for the precision failure above: an update lr × g is often far smaller than the weight it is added to, and in fp16 the spacing near 1.0 is 9.77e-04. Any update below half of that vanishes. Keeping the weights and the optimizer's state in fp32, applying the update there, and using a low-precision copy only for the forward and backward matmuls keeps every small update.

In PyTorch's autocast recipe this is automatic: the parameters stay fp32 (model.half() is exactly what the autocast documentation says not to do), autocast's weight cache provides the fp16 or bf16 copy for each forward pass, and the optimizer updates the fp32 parameters. The third technique lives inside the matmul kernels, which accumulate in fp32 even when their inputs are 16-bit. Pure-bf16 training, with parameters stored in bf16, saves the master copy's memory but has to deal with lost updates some other way (stochastic rounding or compensated summation), and is a separate recipe from autocast.

fp8

fp8 doubles matrix-unit throughput again on hardware that supports it, but with only 3 or 2 mantissa bits and a narrow range, a single global scale is not enough. fp8 training recipes keep a scale per tensor (or per block of a tensor), computed from recent maximum absolute values, and use e4m3 where precision matters (weights and activations in the forward pass) and e5m2 where range matters (gradients in the backward pass), as the FP8 format paper proposes. PyTorch provides the dtypes and conversions; the scaled matmuls and their recipes live in separate libraries and are beyond this page.

The PyTorch API

Worked example

Implementation

import torch

model = MyModel().cuda()                      # parameters stay float32: the master weights
optimizer = torch.optim.AdamW(model.parameters(), lr=3e-4)

# float16: autocast plus dynamic loss scaling.
scaler = torch.amp.GradScaler("cuda")
for x, y in loader:
    optimizer.zero_grad(set_to_none=True)
    with torch.autocast("cuda", dtype=torch.float16):
        loss = loss_fn(model(x), y)           # matmuls in fp16, softmax/norms/loss in fp32
    scaler.scale(loss).backward()             # outside autocast
    scaler.unscale_(optimizer)                # so clipping sees true gradient norms
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    scaler.step(optimizer)                    # skipped if any gradient is inf/NaN
    scaler.update()

# bfloat16: the same autocast, no scaler.
for x, y in loader:
    optimizer.zero_grad(set_to_none=True)
    with torch.autocast("cuda", dtype=torch.bfloat16):
        loss = loss_fn(model(x), y)
    loss.backward()
    optimizer.step()

# Let plain float32 matmuls use TF32 on GPUs that support it.
torch.set_float32_matmul_precision("high")

torch.finfo(torch.float16).max                # 65504.0
torch.finfo(torch.bfloat16).eps               # 0.0078125

Save scaler.state_dict() with the model and optimizer so a resumed fp16 run does not restart its scale from 2¹⁶. Inside an autocast region an op's output dtype follows its policy, so a tensor coming out of a matmul is fp16 and one coming out of a softmax is fp32. Mixing an autocast output with an fp32 tensor in an op outside autocast needs an explicit .float().

Tradeoffs

• fp16 versus bf16. fp16 has 3 more mantissa bits and needs loss scaling, occasional skipped steps and care with large activations. bf16 needs none of that and is the usual choice where the hardware supports it, at the cost of noisier values.
• Speed versus where it comes from. Matrix-unit throughput gains apply only to the compute-bound ops autocast lowers. Memory-bound ops gain only from halved bytes, and ops on the fp32 list gain nothing, so the end-to-end speedup is smaller than the matmul speedup.
• Memory. Activations saved for backward are halved, but fp32 master weights and fp32 optimizer state stay, plus the cached low-precision weight copies. The savings are mostly in activations.
• tf32. A free speedup for fp32 code with fp16-level precision inside matmuls. That is fine for training most networks and wrong for code that needs full fp32 matmuls, which is why PyTorch's default is "highest".
• fp8. Another 2× in matmul throughput, but per-tensor scaling, saturation and format choice per tensor role add engineering that 16-bit training avoids.

Connections

References

• Micikevicius, Narang, Alben, Diamos, Elsen, Garcia, Ginsburg, Houston, Kuchaiev, Venkatesh and Wu. Mixed Precision Training. ICLR, 2018. arXiv:1710.03740.
• Micikevicius et al. FP8 Formats for Deep Learning. 2022. arXiv:2209.05433 (the e4m3 and e5m2 formats; PyTorch's fp8 headers cite it).
• Kalamkar et al. A Study of BFLOAT16 for Deep Learning Training. 2019. arXiv:1905.12322.
• Henry, Tang and Heinecke. Leveraging the bfloat16 Artificial Intelligence Datatype For Higher-Precision Computations. 2019. arXiv:1904.06376 (the two-bf16 split behind "high" matmul precision).
• IEEE. IEEE Standard for Floating-Point Arithmetic (IEEE 754-2019): binary32 and binary16.
• PyTorch v2.15.0: torch/amp/autocast_mode.py, torch/amp/grad_scaler.py, aten/src/ATen/autocast_mode.cpp, aten/src/ATen/autocast_mode.h.