Mixed-Precision Training

Floating-point formats covers what bf16 and fp16 are. This entry is about the training loop built around them, which is more than a cast and has one requirement that surprises people.

Why a master copy is needed

The problem is not the forward pass. It is the update.

bf16 has seven mantissa bits, so it resolves about three significant decimal digits. Adding $10^{-7}$ to $10^{-2}$ in bf16 returns $10^{-2}$ unchanged. The update does not round — it is discarded entirely, and the parameter never moves again.

So the optimiser keeps an fp32 master copy, applies the update there, and casts down to bf16 for the next forward pass. The half-precision weights are a derived artefact, regenerated every step.

Against pure fp32’s $16P$. Mixed precision costs weight memory and saves activation memory, and since activations dominate at realistic batch and sequence sizes the total is a large win — but stating it as “halves memory” is wrong and leads to allocation failures.

The precision policy

Not every operation is safe in half precision. The rule is that anything accumulating over many terms stays in fp32.

torch.autocast implements this as a per-operator table. It is worth knowing that the table exists and can be wrong for a custom operator — a hand-written kernel gets no entry, runs in whatever dtype it is handed, and is the usual explanation for a model that trains in fp32 and diverges in bf16.

fp16 and loss scaling

With bf16 the story ends there. With fp16 there is one more piece: five exponent bits do not reach the small gradients, which flush to zero.

Loss scaling multiplies the loss by a large constant before the backward pass so gradients land in range, then divides it out before the optimiser sees them. Dynamic scaling adjusts the constant — double it periodically, halve it and skip the step whenever an overflow appears.

That skipping is the part worth knowing about: a well-tuned fp16 run discards a small fraction of its steps as a matter of routine. bf16 needs none of it, which is the practical reason it took over.

import torch
from torch import nn


def train_step(model: nn.Module, opt, batch, scaler=None) -> float:
    """bf16 needs no scaler; fp16 does. Otherwise identical."""
    opt.zero_grad(set_to_none=True)

    with torch.autocast("cuda", dtype=torch.bfloat16):
        # Only the forward runs under autocast. The backward inherits the
        # dtypes autograd recorded — wrapping it too is a no-op at best.
        loss = model(**batch).loss

    if scaler is None:
        loss.backward()
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        opt.step()
    else:
        scaler.scale(loss).backward()
        scaler.unscale_(opt)                       # before clipping, not after
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        scaler.step(opt)                           # skips on overflow
        scaler.update()

    return loss.item()

unscale_ before clipping is the detail that is easy to miss and quiet when wrong. Clipping a scaled gradient compares its norm against a threshold on a completely different scale, so either everything is clipped or nothing is, and the run simply trains a bit worse without any error.