Mixed-precision training

Your 24 GB GPU is almost large enough for the model. Almost.

The forward pass fits in fp32, but the activations push the run over the limit at batch size 32. The GPU spends much of its time multiplying matrices, yet those multiplications use four-byte numbers when the hardware has dedicated two-byte units sitting there idle. At 3 a.m., the practical choices are unpleasant: buy a larger GPU, reduce the batch, or wait twice as long.

Mixed-precision training is the third option made less unpleasant. It uses 16-bit arithmetic for operations that tolerate rounding, while retaining 32-bit arithmetic where a small numerical error could become a training failure.

The word mixed matters. This is not “convert the model to half precision and hope.” It is a controlled compromise between speed, memory, and numerical range.

The two 16-bit formats

A floating-point number stores three things:

• a sign, which says positive or negative;
• an exponent, which controls how large or tiny the number can be;
• a fraction, also called the mantissa, which controls how finely nearby values are represented.

The exponent matters when gradients become extremely small or activations become huge. The fraction matters when you need fine distinctions between nearby values.

The fraction-bit count does not include the leading one implied by normalised binary numbers. So fp16 has roughly 11 significant bits and bf16 roughly 8. Around 1, adjacent fp16 values are about 0.0009765625 apart; adjacent bf16 values are about 0.0078125 apart.

That is the bargain:

• fp16 is more precise for ordinary-sized values, but its narrow exponent range makes overflow and underflow common.
• bf16 is much less precise, but its exponent range is approximately the same as fp32’s. It can represent very small gradients and very large activations that fp16 cannot.

A number you can see

Here is the range problem in NumPy:

import numpy as np

x = np.array([65000.0, 1e-7, 1e-30], dtype=np.float32)
print("fp32:", x)

y = x.astype(np.float16)
print("fp16:", y)
print("fp16 squared:", y * y)

fp32: [6.5e+04 1.0e-07 1.0e-30]
fp16: [6.5e+04 1.2e-07 0.0e+00]
fp16 squared: [inf  0.  0.]

The first value is close to fp16’s ceiling. Squaring it produces about 4.2e9, above 65504, so the result becomes inf. The second survives conversion, but squaring it produces roughly 1.4e-14, below fp16’s smallest subnormal, so it becomes zero. The third is already too small and becomes zero during conversion.

Bfloat16 rounds these values more coarsely than fp16, but its exponent range is broad enough to retain values in both situations. It can still overflow if a computation exceeds its maximum.

Why 16-bit arithmetic is faster

A matrix multiplication performs enormous numbers of multiply-and-add operations. On recent GPUs, Tensor Cores can perform them on fp16 or bf16 inputs using specialised hardware. Two-byte inputs also move through memory faster and occupy half the storage of four-byte inputs.

That creates two wins:

1. Throughput: eligible matrix multiplications and convolutions can use faster low-precision hardware.
2. Capacity: activations and temporary buffers take less memory, leaving room for a larger batch or sequence.

Neither is automatic. A small model may spend more time launching kernels than multiplying numbers; unsupported operations may remain fp32; and half the tensor bytes does not mean half the training memory. Speed depends on tensor shapes, GPU generation, kernels, and whether the run is compute-bound. Treat 2x as a possibility, not a guarantee. Benchmark the complete training step.

What autocast actually does

PyTorch’s autocast is a context manager that chooses a dtype for each eligible operation inside a block. It may send a convolution or matrix multiplication through bf16 while keeping a reduction or another sensitive operation in fp32.

It does not permanently change model parameters. In the standard pattern, the model remains stored in fp32, and autocast changes selected forward-pass computations. The exact policy is operation- and device-dependent: autocast is not a promise that every operation becomes 16-bit or that every loss is fp32.

For bf16, the usual loop is:

import torch

from torch.amp import autocast, GradScaler

device = "cuda"
model = model.to(device)

for x, y in loader:
    x, y = x.to(device), y.to(device)
    optimizer.zero_grad(set_to_none=True)

    with autocast(device_type="cuda", dtype=torch.bfloat16):
        pred = model(x)
        loss = loss_fn(pred, y)

    loss.backward()
    optimizer.step()

There is no GradScaler here. Bfloat16’s wide exponent range usually makes gradient underflow rare enough that dynamic loss scaling is unnecessary.

For fp16, use the scaler:

scaler = GradScaler("cuda")

for x, y in loader:
    x, y = x.cuda(), y.cuda()
    optimizer.zero_grad(set_to_none=True)

    with autocast(device_type="cuda", dtype=torch.float16):
        pred = model(x)
        loss = loss_fn(pred, y)

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

The scaler multiplies the loss before backpropagation. If the loss is multiplied by 65536, a gradient of 1e-8 becomes roughly 6.6e-4 during the backward pass, where fp16 can represent it. scaler.step divides gradients back down before updating parameters. If unscaled gradients contain inf or nan, it skips the update and lowers the scale for later iterations.

Scaling changes the numerical route, not the intended gradient. It guards against underflow; it cannot repair an exploding model, an invalid custom operation, or an excessive learning rate.

scaler.scale(loss).backward()
scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
scaler.step(optimizer)
scaler.update()

Call unscale_ once for that optimizer between backward and step, whenever you inspect, clip, or modify gradients.

The memory calculation people get wrong

Suppose a model has one billion parameters and uses AdamW. Ignoring allocator overhead:

• fp32 parameters: 1B × 4 bytes = 4 GB;
• fp32 gradients: another 4 GB;
• Adam’s first and second moments: 1B × 4 × 2 = 8 GB.

That is about 16 GB before activations, temporary buffers, CUDA workspaces, and data. If autocast stores large intermediate activations in bf16, it may save several gigabytes. It does not automatically convert existing fp32 parameters, gradients, and Adam states into two-byte objects.

This explains a common surprise: the forward pass succeeds with mixed precision, then optimizer.step() fails with out-of-memory because the optimizer is allocating state for the first time. The symptom does not mean autocast failed; persistent training state may dominate the budget.

With standard AMP, parameters remain fp32, so Adam’s moment tensors are normally fp32 too. If you explicitly store parameters in bf16, use an optimizer with fp32 master weights, or use a sharded system, the accounting changes. Inspect actual parameter, gradient, and optimizer-state dtypes.

For an activation-heavy transformer, mixed precision may make the difference between batch size 8 and 16. If activations remain the bottleneck, activation checkpointing trades recomputation for lower memory. If optimizer state is the bottleneck, changing activation dtype will not solve it.

Which dtype should you choose?

Bf16 does not have better accuracy by definition; it has worse local precision than fp16. It often wins because training cares more about avoiding catastrophic range loss than about representing nearby values with extreme precision.

As of August 2026, bf16 is the sensible default for most modern GPU training. FP8 is also used on supported datacenter hardware, particularly for large transformers, but it needs scaling metadata and carefully implemented kernels. Use the framework and hardware path that supports it rather than replacing every tensor’s dtype by hand.

Failure modes you can diagnose

The run is still out of memory

If the error appears during the forward pass, activations or temporary workspaces are probably the main pressure. If it appears on the first optimizer.step, optimizer state is a prime suspect. If it appears only after several steps, investigate saved tensors, caches, or a memory leak.

Check peak allocated and reserved memory and inspect parameter and optimizer-state dtypes. Reduce sequence length or per-device batch size, use gradient accumulation, choose a lighter optimizer, or add checkpointing. Do not assume autocast halves every memory category.

The loss becomes nan or inf

With fp16, this can come from an intermediate exceeding 65504 or gradients underflowing before scaling. Confirm that the loop uses scaler.scale, scaler.step, and scaler.update. Check gradients after unscaling:

scaler.unscale_(optimizer)

bad = False
for parameter in model.parameters():
    if parameter.grad is not None and not torch.isfinite(parameter.grad).all():
        bad = True
        break

if bad:
    print("non-finite gradient")

Then find the first operation that produced the bad value. A custom exponential, unnormalised attention score, excessive learning rate, or invalid input can fail in fp32 too. Bfloat16 may avoid fp16 range failures, but it cannot make explosive code safe.

If one operation is genuinely sensitive, compute only that operation in fp32. For example, a custom probability calculation can consume logits.float() outside the autocast block while matrix multiplications remain low precision.

Training runs but accuracy degrades

A model can converge in fp32 and stall in bf16 because of rounding rather than range. Repeated additions, small probability differences, normalisation statistics, and custom reductions may lose meaningful information. Manually forcing .half() can make this worse by overriding autocast’s safer choices.

Remove manual casts, keep parameters in their normal dtype, and force only the proven trouble spot to fp32. Compare a short, fixed batch sequence under fp32 and bf16 with the same seed. Look for the first operation where values diverge.

The run is not faster

Check that the dtype is supported natively by the GPU and that matrix dimensions and batch shapes allow efficient kernels. Measure the whole step, including input transfer and data loading. Mixed precision helps most when large convolutions or matrix multiplications dominate; it is not a universal speed switch.

Manual casts cause dtype errors

Errors such as “expected mat1 and mat2 to have the same dtype” often result from casting an input, one layer, or a saved tensor independently. Use one autocast boundary around the forward pass and loss, and avoid scattered .half() or .bfloat16() calls. If a library operation requires a particular dtype, make that boundary explicit and document why.

A production checklist

Before comparing runs, record the GPU, dtype, batch shape, sequence length, gradient accumulation count, and parameter dtype. Compare peak memory, examples per second, and validation metrics after the same number of processed examples.

Use the training loop as the surrounding structure: zero gradients deliberately, run the forward pass under autocast, scale only for fp16, unscale before clipping, and step the optimizer exactly once. Keep the optimizer step outside autocast.

For a new model, begin with a tiny fixed batch. Run a few fp32 steps, then bf16 or fp16 steps. Check that losses and gradients are finite, parameters change, and metrics move as expected. A short numerical smoke test is cheaper than discovering a silent precision bug after a long run.

In one breath

Mixed precision is selective low-precision arithmetic, not an all-or-nothing conversion. Two-byte fp16 and bf16 values reduce memory and bandwidth for eligible activations and let modern GPUs use fast matrix hardware. The formats trade off range and precision: fp16 has more fraction precision but a maximum finite value of 65504; bf16 has much less fraction precision but approximately fp32’s exponent range.

Use bf16 autocast by default on supported modern hardware. Use fp16 with GradScaler when hardware or workload requires it. Scaling protects small gradients from underflow, while autocast chooses operations that can tolerate rounding. Neither prevents an invalid model or guarantees that total memory is cut in half.

Remember the optimizer: a one-billion-parameter fp32 model with AdamW can spend about 16 GB on parameters, gradients, and two moment tensors before activations. Measure the real bottleneck, then fix that bottleneck.

Quick check