Mixed precision and loss scaling

What it is and when you reach for it

Mixed precision training runs the compute-heavy operations of a network, mostly matrix multiplies and convolutions, in a 16-bit floating point format while keeping a 32-bit copy of the weights for the optimizer to update. The wins are concrete: half the bytes per activation and gradient means roughly double the effective memory and memory bandwidth, and on every GPU since Volta the tensor cores execute half-precision matmuls at several times the fp32 rate. The technique was formalized by Micikevicius et al. (2017) with three ingredients: half-precision compute, an fp32 master copy of the weights, and loss scaling to keep fp16 gradients from flushing to zero. The arrival of bfloat16 on TPUs and on NVIDIA GPUs from Ampere onward made the third ingredient optional, and the modern default for large-model training is bf16 compute with fp32 master state and no loss scaler at all. You reach for mixed precision whenever you train anything nontrivial on modern hardware, which is to say always; the only real decisions left are fp16-plus-scaler versus bf16, and which operations to hold back in fp32.

The math

Three bit layouts

A binary float is a sign bit, an exponent field, and a mantissa (fraction) field, encoding ±1.mantissa × 2^{exponent−bias}. The three formats that matter split their bits like this:

        sign  exponent      mantissa            range (normal)     rel. precision
fp32    s     eeeeeeee      mmmmmmmmmmmmmmmmmmmmmmm
        1     8 bits        23 bits             ~1.2e-38..3.4e38   2^-24 ~ 6e-8

fp16    s     eeeee         mmmmmmmmmm
        1     5 bits        10 bits             ~6.1e-5..65504     2^-11 ~ 4.9e-4

bf16    s     eeeeeeee      mmmmmmm
        1     8 bits        7 bits              ~1.2e-38..3.4e38   2^-8  ~ 3.9e-3

The exponent field sets dynamic range: how far from 1.0 a value can wander before it overflows to infinity or underflows to zero. The mantissa sets precision: how many significant bits each value carries. fp16 spent its bits on precision (10 mantissa bits, about 3 decimal digits) and starves the exponent: the largest finite fp16 is 65504, the smallest normal is about 6.1×10^-5, and with subnormals the absolute floor is 2^-24 ≈ 6×10^-8. bfloat16 is simply fp32 with the bottom 16 mantissa bits sliced off: it keeps the full 8-bit exponent and therefore fp32's entire range, at the cost of 7 mantissa bits, about 2 to 3 decimal digits. bf16 trades precision for range; fp16 traded range for precision, and range is the thing deep learning actually needs.

Why fp16 underflows gradients and bf16 does not

Weights and activations in a trained network cluster within a few orders of magnitude of 1, comfortably inside any of these formats. Gradients do not. Backpropagated activation gradients shrink as they pass through layers scaled by small weights and saturating nonlinearities, and the mixed-precision paper's measurements show large fractions of gradient values with magnitudes below 2^-24, fp16's absolute floor. Every one of those becomes exactly zero when stored in fp16: not imprecise, gone. The model keeps training but the small-gradient parameters silently stop learning, and the loss curve degrades or diverges from the fp32 baseline. bf16 has no such cliff: its floor is fp32's ~10^-38, some thirty orders of magnitude of headroom below any gradient a sane model produces. What bf16 gives up instead is resolution: 1.0 + 0.003 rounds back to 1.0 in bf16, since the next representable value above 1 is 1 + 2^-7 ≈ 1.0078. Networks are robust to that kind of rounding in the forward and backward compute; they are not robust to entire gradient populations flushing to zero. This asymmetry is the whole reason bf16 training needs no loss scaling and fp16 training does.

Loss scaling, derived

Backpropagation is linear in the output gradient: every chain-rule product that produces ∂L/∂θ starts from ∂L/∂L = 1, so training on S·L instead of L multiplies every intermediate and final gradient by exactly S. That gives a free translation knob in log space:

grad(S * L) = S * grad(L)          exact, by linearity of backprop

fp16 exponent axis (powers of 2):
  -24        -14                 0                +15
   |----------|------------------|-----------------|
   ^ subnormal floor                       65504 ^
      [ gradient histogram ]  --- multiply by S=2^16 -->  [ shifted histogram ]
      (partly below floor,                                (fully inside range)
       flushes to zero)

Choose S so the gradient histogram slides up out of the underflow region without its largest values crossing 65504, compute the backward pass in fp16, then divide the gradients by S in fp32 before they touch gradient clipping or the optimizer. The unscale must happen in fp32: dividing by S in fp16 would re-underflow exactly the values the scale was protecting.

Static scaling picks a constant, typically 2¹⁵ or 2¹⁶, and works when you know the gradient range. Dynamic scaling removes the tuning with a bang-bang control loop. Start high, S = 2¹⁶. After each backward pass, check the gradients for inf or NaN. If any appear, the scale pushed some gradient past 65504: skip the optimizer step entirely (the gradients are garbage) and halve S. If a run of N consecutive steps (2000 in torch.amp.GradScaler's defaults) completes cleanly, double S. The loop hunts for the largest scale that does not overflow and tracks it as gradient magnitudes drift during training; the occasional skipped step and halving is the deliberate cost of probing the ceiling. This is why fp16 training logs show the scale sawtoothing, and why a scale that halves repeatedly without recovering signals real divergence rather than scaler mischief.

Master weights in fp32

Floating point addition loses the smaller addend once the magnitude gap exceeds the mantissa width: in fp16, x + δ returns x whenever |δ| < |x| · 2^-11. Optimizer updates are exactly this shape. A typical weight is ~10^-1; a typical per-step update with a small learning rate is 10^-5 or below; the ratio 10^-4 sits under fp16's 2^-11 ≈ 4.9×10^-4 threshold, so applying the update to an fp16 weight does nothing at all. Training proceeds, loss plateaus, nothing is NaN, and the weights simply are not moving. The fix is to keep the authoritative copy of every weight in fp32, whose 2^-24 threshold gives eleven more bits of headroom, apply updates there, and cast fresh 16-bit copies for each forward pass. In bf16 the problem is even more acute (2^-8 threshold), which is why the bf16 recipe still keeps fp32 master weights: bf16 removes the need for loss scaling, not the need for fp32 accumulation of small updates. The same logic keeps Adam's moment EMAs in fp32: they too are long-running accumulations of small increments.

Which ops stay in fp32, and why

Autocast is not "run everything in 16 bits"; it is a per-op policy. Matmuls and convolutions go to 16-bit inputs because tensor cores accumulate their dot products in fp32 internally, so the cheap format costs little accuracy where the flops are. Everything whose failure mode is a bad reduction stays in fp32: sums and means over many elements (a sum of 4096 squared activations around 10 reaches ~409,600, already past fp16's 65504), variance computations inside LayerNorm and BatchNorm (cancellation between large nearly-equal terms), softmax (exp has a brutal dynamic range; see the softmax page for the subtract-max machinery), log, pow, and the loss itself. This is precisely the split torch.autocast hardcodes in its op lists, and the reason a manual implementation must cast the loss computation back up. Even inside fused fp16 attention kernels the softmax statistics and accumulators are carried in fp32.

fp8, briefly

The same trade pushed further gives fp8, in two flavors: e4m3 (4 exponent bits, 3 mantissa bits, more precision) for weights and activations, and e5m2 (more range) for gradients. Eight bits is too narrow for a single global loss scale, so fp8 training keeps per-tensor scaling factors updated from a rolling history of each tensor's observed maximum, which is the delayed-scaling machinery NVIDIA's Transformer Engine implements on Hopper and Blackwell. Conceptually it is loss scaling gone local: one scale per tensor instead of one for the whole backward pass. The inference-side story, fp8 KV caches and quantized serving, is covered in the TensorRT-LLM chapter.

Implementation, twice

First the manual version: a two-layer MLP written so every cast is visible. The PyTorch tab is a faithful expansion of what torch.autocast plus GradScaler do: fp16 compute copies of fp32 master weights, fp32 loss, scaled backward, inf/NaN check, unscale, conditional step, and the grow/backoff loop with PyTorch's default constants. The JAX tab is the bf16 policy written as explicit dtype discipline on a params pytree: params fp32, compute bf16, output and loss fp32; since the casts are part of the traced graph, gradients arrive in fp32 automatically, and no scaler is needed because the format is bf16.

import torch

torch.manual_seed(0)
dev = "cuda"

# fp32 master weights: the only copy the optimizer ever touches.
w1 = (torch.randn(4096, 1024, device=dev) * 0.02).requires_grad_()
w2 = (torch.randn(1024, 4096, device=dev) * 0.02).requires_grad_()
opt = torch.optim.AdamW([w1, w2], lr=3e-4)   # states in fp32 too

scale, good_steps = 2.0 ** 16, 0             # GradScaler's defaults
GROWTH_INTERVAL = 2000

for step in range(1000):
    x = torch.randn(32, 1024, device=dev)
    y = torch.randn(32, 1024, device=dev)

    # Forward in fp16 compute copies. The .half() casts are inside the
    # graph, so autograd casts gradients back to the fp32 leaves.
    h = torch.relu(x.half() @ w1.half().t())
    pred = h @ w2.half().t()
    loss = (pred.float() - y).pow(2).mean()  # loss math in fp32 (a reduction)

    (loss * scale).backward()                # every grad now carries factor S

    with torch.no_grad():
        grads = [w1.grad, w2.grad]
        if any(not g.isfinite().all() for g in grads):
            scale, good_steps = scale * 0.5, 0   # overflow: back off, skip step
        else:
            for g in grads:
                g.div_(scale)                # unscale in fp32, never in fp16
            opt.step()                       # fp32 update on master weights
            good_steps += 1
            if good_steps == GROWTH_INTERVAL:
                scale, good_steps = scale * 2.0, 0   # probe a higher ceiling
    opt.zero_grad(set_to_none=True)

import jax
import jax.numpy as jnp
import optax

key = jax.random.PRNGKey(0)
k1, k2 = jax.random.split(key)

# fp32 params pytree: the master copy, owned by the optimizer.
params = {"w1": jax.random.normal(k1, (1024, 4096)) * 0.02,
          "w2": jax.random.normal(k2, (4096, 1024)) * 0.02}

def forward(params, x):
    # Explicit dtype policy, jmp's "params=float32,compute=bfloat16,
    # output=float32" spelled out: cast down at the boundary, compute
    # in bf16, cast the output back up.
    p = jax.tree.map(lambda t: t.astype(jnp.bfloat16), params)
    h = jax.nn.relu(x.astype(jnp.bfloat16) @ p["w1"])
    return (h @ p["w2"]).astype(jnp.float32)

def loss_fn(params, x, y):
    return jnp.mean((forward(params, x) - y) ** 2)   # reduction in fp32

opt = optax.adamw(3e-4)
opt_state = opt.init(params)   # moment EMAs in fp32, like the params

@jax.jit
def train_step(params, opt_state, x, y):
    # The astype calls are traced ops, so their cotangents cast back up:
    # grads arrive as an fp32 pytree matching params. bf16 has fp32's
    # exponent range, so no loss scaling is required anywhere.
    loss, grads = jax.value_and_grad(loss_fn)(params, x, y)
    updates, opt_state = opt.update(grads, opt_state, params)
    return optax.apply_updates(params, updates), opt_state, loss

Then the idiomatic version. In PyTorch, torch.autocast replaces the hand-placed casts with the per-op policy and torch.amp.GradScaler owns the scale, the inf check, the conditional step, and the growth loop; switching the dtype to bf16 makes the scaler unnecessary (it is a no-op when disabled, so the code can stay shape-identical). In JAX, jmp packages the same discipline as a declarative policy object plus a loss-scale type for the fp16 case.

import torch

use_bf16 = torch.cuda.is_bf16_supported()      # Ampere and newer
amp_dtype = torch.bfloat16 if use_bf16 else torch.float16
# GradScaler is a transparent no-op when disabled, so one loop serves both.
scaler = torch.amp.GradScaler("cuda", enabled=not use_bf16)

for x, y in loader:
    with torch.autocast("cuda", dtype=amp_dtype):
        pred = model(x)                # matmuls in 16-bit per the op policy
        loss = loss_fn(pred, y)        # reductions autocast keeps in fp32
    scaler.scale(loss).backward()

    scaler.unscale_(opt)               # expose true grads BEFORE clipping
    torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)

    scaler.step(opt)                   # skips the step if inf/NaN was found
    scaler.update()                    # backoff on overflow, growth after 2000
    opt.zero_grad(set_to_none=True)

import jax
import jax.numpy as jnp
import jmp

policy = jmp.get_policy("params=float32,compute=bfloat16,output=float32")

def forward(params, x):
    params, x = policy.cast_to_compute((params, x))
    h = jax.nn.relu(x @ params["w1"])
    return policy.cast_to_output(h @ params["w2"])

# fp16 on older accelerators: jmp adds the scaler side as a pytree value
# threaded through the step, mirroring GradScaler's skip/adjust loop.
loss_scale = jmp.DynamicLossScale(jnp.float32(2 ** 15))

def train_step(params, opt_state, loss_scale, x, y):
    def scaled_loss(p):
        return loss_scale.scale(loss_fn(p, x, y))
    grads = jax.grad(scaled_loss)(params)
    grads = loss_scale.unscale(grads)

    finite = jmp.all_finite(grads)              # the inf/NaN check
    loss_scale = loss_scale.adjust(finite)      # halve or grow the scale
    updates, new_opt_state = opt.update(grads, opt_state, params)
    new_params = optax.apply_updates(params, updates)
    # Commit the step only if grads were finite, else keep the old trees.
    params, opt_state = jmp.select_tree(
        finite, (new_params, new_opt_state), (params, opt_state))
    return params, opt_state, loss_scale

Using it on a real shape of problem

On a transformer block of model dimension 1024 with batch 32 and sequence length 512, switching the forward and backward compute from fp32 to bf16 roughly halves activation memory and, on an Ampere-or-newer GPU, speeds the matmul-bound steps by anywhere from 1.5x to 3x depending on how memory-bound the model is; exact numbers are hardware-dependent and worth measuring on your own stack rather than assuming. The functional check matters more than the speed check: run a few hundred steps in full fp32 and in mixed precision from the same seed and overlay the loss curves. They will not match step for step (different rounding), but they should track within noise; a mixed-precision curve that detaches from the fp32 curve is the signal to hunt for an op that needed fp32, a missing master copy, or, in fp16, a scale stuck at a low value. In fp16 runs, also log scaler.get_scale(): healthy training shows it climbing to a plateau and sawtoothing gently, with a handful of skipped steps early on.

Applications

The standard bf16 LLM recipe

Every major LLM pretraining stack today runs the same configuration: bf16 compute for the model, fp32 master weights and optimizer state, fp32 for gradient all-reduce across data-parallel ranks, no loss scaler. That last choice, reducing in fp32, exists because summing gradients over hundreds of ranks is one more long accumulation with the same failure mode as everything else in this page. This is literally torchtitan's mixed-precision policy: FSDP's MixedPrecisionPolicy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32), sharded bf16 parameters for compute with fp32 communication and fp32 sharded optimizer state. At the small end, nanoGPT is a clean single-file example of the autocast idiom above: one torch.autocast context with bf16 where supported and a GradScaler that is enabled only in the fp16 fallback. AdamW sits inside all of these recipes as the fp32 half of the bargain: the moments and master weights it keeps are the reason the 16-bit rounding never compounds across steps.

When fp16 plus a scaler is still the answer

bf16 arithmetic requires Ampere (A100, RTX 30-series) or newer on NVIDIA, so the fp16-plus-GradScaler path remains the working recipe on V100s, T4s, GTX 10/16 and RTX 20 series cards, which is a large fraction of cloud spot capacity, Colab-class free tiers, and installed research hardware. fp16 also retains a genuine edge in inference and in some fine-tuning regimes: its 10 mantissa bits are three more than bf16's, so when the value range is controlled, fp16 is the more accurate 16-bit format, which is why inference engines often prefer it even on hardware with bf16.

Debugging NaNs

Mixed precision changes the meaning of a NaN. In fp16 training, first look at the scale: infs in the scaled gradients are the scaler's normal probing behavior, handled by the skip-and-halve loop, and are only pathological when the scale collapses repeatedly toward 1. A genuine fp16 overflow in the forward pass (attention logits are the classic site, which is one reason attention softmax runs in fp32) shows up as NaN in the loss itself, before the scaler is involved. In bf16, precision almost never produces NaNs; a bf16 NaN is nearly always real mathematics, a log of zero, a division by zero, an exploding update, and should be debugged as such, with torch.autograd.set_detect_anomaly(True) to find the producing op on a small repro. The most common false accusation in both regimes is blaming the format for a learning rate or initialization problem that fp32 would have exposed a few hundred steps later.

Against the real libraries

torch.amp adds, over the manual loop above: per-op cast policies covering the whole operator surface (including ops this page never mentioned, with a cache so each weight is cast once per step), nesting and disabling of autocast regions, CPU and other-backend autocast, a GradScaler that handles multiple optimizers and gradient accumulation correctly, and state-dict serialization of the scale so training resumes cleanly. The manual version breaks in exactly the places these features cover: forget one cast, or unscale twice, and you get silent wrongness rather than an error. jmp packages the JAX discipline: Policy objects parsed from strings, cast_to_compute / cast_to_param / cast_to_output applied to whole pytrees, and NoOpLossScale / StaticLossScale / DynamicLossScale types that make the fp16 machinery a swappable value instead of control flow. NVIDIA Transformer Engine operates a level below both: fused transformer layers whose internals pick fp8/fp16/bf16 per tensor, with the per-tensor delayed-scaling recipes fp8 needs, integrated into Megatron, JAX, and PyTorch stacks.

The from-scratch loop is enough when you control the whole model and want the casts auditable, in teaching and research code, and when debugging what autocast is actually doing to a misbehaving op. Verification is direct: run the manual fp16 loop and the autocast/GradScaler loop on the same model, seed, and data, and compare per-step losses and the evolution of the scale; the trajectories should agree closely at the start (divergence grows with steps as rounding compounds), and the scaler's skip-and-halve events should occur at the same steps if the manual inf check is placed correctly.

Traps and misconceptions

"Just call model.half()." Converting every parameter and buffer to fp16 is not mixed precision, it is low precision: no fp32 master copy, so small updates vanish into the 2^-11 rounding threshold, and norms and losses compute their reductions in fp16. It can appear to work for a while, which makes it worse. Mixed precision is defined by what it keeps in fp32, not by what it casts down.

"Loss scaling fixes overflow." Backwards: scaling multiplies gradients up and can only cause overflow, which the dynamic loop then detects and backs off from. What scaling fixes is underflow. If your forward pass overflows fp16 (activations or attention logits past 65504), no scale value helps, because the damage happens before the loss exists; the fixes are fp32 islands for the offending ops, or bf16.

"Clip gradients, then unscale." Clipping scaled gradients compares them against a threshold that is effectively max_norm × S with S changing over training: a random, drifting clip. The correct order is scaler.unscale_(opt) first, clip second, scaler.step third, and the scaler guards against double-unscaling for exactly this workflow.

"bf16 is just a safer fp32." bf16 matches fp32's range, not its precision: 2 to 3 significant decimal digits. Long accumulations, EMAs, and anything else that adds many small numbers into a large one must still be carried in fp32, which is why optimizer state, gradient reductions, and softmax accumulators stay there even in all-bf16-compute stacks. Pure-bf16 experiments that skip fp32 accumulation typically show slow, quiet convergence degradation rather than dramatic failure.

"Autocast converted my model." Autocast is a context, not a conversion: weights remain fp32, and ops are cast at call time inside the region only. Computing the loss outside the region, calling explicit .half() inside it, or running custom ops that lack autocast rules all silently sidestep the policy. The symptom is usually dtype-mismatch errors at best and quiet fp32 fallback (no speedup) at worst; when in doubt, print dtypes at the boundaries.