Mixed Precision Training

1. Why the Size of a Number Matters

A bit stores either 0 or 1, and a byte contains eight bits. A 32-bit number therefore occupies four bytes, while a 16-bit number occupies two. For one billion values, that is 4 GB versus 2 GB of raw storage, using decimal gigabytes. This counts only those values, not the rest of a training run.

A number's data type tells the computer how to interpret its bits. Integer types, written as INT, represent whole numbers. For example, a signed INT8 uses eight bits to store integers from −128 to 127. It cannot directly represent a fraction such as 0.125. Integer quantization adds a scale, and sometimes an offset, to map those integer codes to real values.

Training commonly uses floating-point types, written as FP, which can represent fractional values as well as very large or very small magnitudes. FP32 is a 32-bit floating-point format; FP16 is a 16-bit one. They store approximations to many numbers, even when the decimal expression looks simple. The question is how much approximation a calculation can tolerate.

2. Sign, Exponent, and Mantissa

In decimal scientific notation, we can write −625 as −6.25 × 10². The sign tells us whether the value is positive or negative, the exponent sets its scale, and the significant digits tell us its value within that scale. Floating-point storage follows a similar idea in binary, using powers of two.

The bits are divided into three fields. The sign uses one bit. The exponent determines the scale, giving the format its range of magnitudes. The fraction supplies the detail within that scale. You will often see this last field called the mantissa; more precisely, it supplies the stored fraction bits of the number's significand.

These fields explain two different limits. Range describes how large or small a magnitude can be represented. Precision describes how finely nearby values can be distinguished. More exponent bits expand the range; more fraction bits bring neighboring representable values closer together. Spending bits on one leaves fewer for the other when the total size is fixed.

3. FP32, FP16, and BF16

FP32 assigns eight bits to the exponent and twenty-three to the fraction. FP16 reduces both, using five exponent bits and ten fraction bits. BF16, short for bfloat16, makes a different choice: it keeps eight exponent bits and uses seven fraction bits. Both smaller formats occupy two bytes, but they preserve different information.

Near 1, the next larger FP16 value is 1.0009765625, while the next larger BF16 value is 1.0078125. FP16 therefore distinguishes finer changes at that magnitude. BF16 has a much wider range: its largest finite value is approximately 3.39 × 10³⁸, close to FP32's 3.40 × 10³⁸. FP16 stops at 65,504.

A value that grows too large can overflow to infinity. At the other end, small values lose precision and may eventually round to zero, a consequence of underflow. BF16's wider exponent range is useful in training because gradients and intermediate results can vary substantially in magnitude. It still has coarser precision than FP16 within their shared normal range. Google's bfloat16 explanation discusses this design tradeoff.

4. Where the Precisions Are Mixed

During a forward pass, a Transformer uses its weights to turn input tokens into predictions. The intermediate results are called activations. A loss measures prediction error, and the backward pass computes gradients: how changes in the weights would affect that loss. Finally, the optimizer uses these gradients to adjust the weights.

Attention projections and feed-forward layers contain large matrix multiplications. These are common places to use FP16 or BF16. Other calculations, such as a loss or a numerically sensitive reduction that adds many values together, can use FP32. Automatic mixed precision selects formats at the operation level rather than assigning one format to the entire model.

The format of a tensor and the precision inside a calculation are also separate choices. A matrix multiplication can read 16-bit inputs, accumulate products in FP32, and write a 16-bit output. The higher-precision accumulator helps preserve the running sum; it cannot restore information already lost when the inputs were rounded. The exact arithmetic depends on the kernel and hardware.

5. Why Weight Updates Need Care

Consider a weight of 1.0 and an optimizer update that subtracts 0.0001. The intended result is 0.9999. If we store that result directly in FP16 or BF16, it rounds back to 1.0. Repeating that same small update while storing only the rounded weight would keep losing the change each time. FP32 retains an approximation close to 0.9999, allowing small updates to build up over time.

The original mixed precision training method keeps an FP32 master copy of the weights while using lower-precision copies for much of the computation. The optimizer updates the master copy, which supplies fresh lower-precision values for later calculations.

In standard PyTorch native AMP with an FP32 model, the model parameters themselves already serve this higher-precision role. Autocast creates the appropriate operation inputs without permanently converting those parameters. Their accumulated .grad buffers are also FP32, even though intermediate backward calculations can use lower precision. A separate master-weight copy is an implementation choice, not an additional copy every AMP setup must allocate.

6. Keeping Small Gradients Alive

FP32 weight storage protects the update, but a gradient can disappear earlier during an FP16 backward calculation. Suppose a gradient should be 0.00000001. As the format explorer showed, that value rounds to zero in FP16. Storing the resulting zero in FP32 afterward cannot recover the missing information.

Loss scaling multiplies the loss by a positive scale before the backward pass. Differentiation multiplies its gradients by the same scale, making small gradients larger while they pass through lower-precision calculations. Before the optimizer uses them, we divide the gradients by that scale in higher precision. This restores their intended magnitude, apart from rounding error.

A scale that is too large can overflow other gradients. Dynamic loss scaling adjusts it over time: when non-finite gradients are detected, the optimizer update is skipped and the scale is reduced; after a run of successful steps, the scale may grow. NVIDIA's mixed precision guide describes the scaling procedure.

BF16 usually does not require loss scaling because its exponent range accommodates much smaller gradients. It still rounds values and can encounter numerical failures. Loss scaling addresses small backward values; it does not repair an invalid forward pass, and it does not give either 16-bit format more fraction bits.

7. Automatic Mixed Precision in PyTorch

PyTorch provides two cooperating tools. torch.autocast chooses operation-specific dtypes within the forward pass and loss calculation. torch.amp.GradScaler handles loss scaling, gradient unscaling, and skipping updates with non-finite gradients when scaling is enabled. Backward runs outside the autocast context. These responsibilities are documented in the AMP reference.

The example below uses one GPU and one batch per update. It expects a causal language model with a Hugging Face-style .logits output and batches containing input_ids, attention_mask, and unshifted labels. Padding and excluded targets must be marked as -100 in the labels. The code pairs each token's prediction with the following token's label.

import torch
import torch.nn.functional as F


def train_epoch(model, loader, optimizer, scaler, device, amp_dtype):
    model.train()
    for batch in loader:
        # Causal prediction uses the labels after the first position.
        if not batch["labels"][:, 1:].ne(-100).any():
            continue

        batch = {name: tensor.to(device) for name, tensor in batch.items()}
        optimizer.zero_grad(set_to_none=True)

        with torch.autocast(device_type=device.type, dtype=amp_dtype):
            logits = model(
                input_ids=batch["input_ids"],
                attention_mask=batch["attention_mask"],
                use_cache=False,
            ).logits
            loss = F.cross_entropy(
                logits[:, :-1, :].reshape(-1, logits.size(-1)),
                batch["labels"][:, 1:].reshape(-1),
                ignore_index=-100,
            )

        scaler.scale(loss).backward()
        scaler.unscale_(optimizer)
        torch.nn.utils.clip_grad_norm_(
            model.parameters(),
            max_norm=1.0,
            error_if_nonfinite=not scaler.is_enabled(),
        )
        scaler.step(optimizer)
        scaler.update()


# Supply a causal language model and a loader before running this setup.
if not torch.cuda.is_available():
    raise RuntimeError("This setup requires a CUDA GPU.")

device = torch.device("cuda")
amp_dtype = (
    torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
)
model = model.to(device=device, dtype=torch.float32)
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5)
scaler = torch.amp.GradScaler("cuda", enabled=(amp_dtype == torch.float16))

train_epoch(model, loader, optimizer, scaler, device, amp_dtype)

The model stays in FP32; there is no model.half() call. The setup selects BF16 when the GPU reports support and otherwise selects FP16. That fallback also needs model compatibility: a checkpoint trained in BF16 may contain values or produce activations outside FP16's range. Converting such a model to an FP16 compute path can overflow even with loss scaling.

The scaler is created once and reused. In the BF16 path, enabled=False makes its scaling operations pass through normally. Gradient clipping, which limits the overall gradient magnitude, happens after unscale_ so its threshold applies to the original gradient. With FP16, scaler.step checks whether the update can proceed. The PyTorch AMP recipe covers this sequence and saving scaler state when resuming training.

A disabled scaler does not perform the non-finite-gradient check for BF16, so the example makes clipping raise an error if the gradient norm is non-finite in that path. Investigate the data and failing operation before continuing training. For a sensitive operation that needs FP32, disable autocast around that region and explicitly convert its floating-point inputs to FP32.

8. Combining It with Gradient Accumulation

Gradient accumulation lets several smaller batches contribute to one optimizer update. When using it with FP16 scaling, every backward pass in that accumulation window must use the same scale. Keep the gradients scaled until all the contributing microbatches have finished.

At the end of the window, unscale once, clip if needed, attempt the optimizer update, and then update the scale. Unscaling midway would mix scaled and unscaled contributions in the same gradient buffer. Changing the scale midway would combine contributions that need different divisors. PyTorch's accumulation example sets these operations at the effective-batch boundary.

The loss still needs the appropriate batch or token averaging described in the previous article. That normalization defines the objective; AMP loss scaling is an additional temporary numerical adjustment that is undone before the optimizer step. If an FP16 update is skipped, a learning-rate schedule defined per optimizer update should not advance as though the weights changed.

9. What Changes in Memory and Speed

A tensor stored in FP16 or BF16 occupies half the raw memory of the same tensor stored in FP32. The complete training run contains many tensors with different roles, however. In our native AMP example, parameters and their gradient buffers remain FP32, and AdamW keeps FP32 running statistics for each parameter. Lower-precision activations can save memory, while temporary casts and other buffers also consume space. Total GPU memory therefore does not automatically fall by half.

Speed comes from the hardware as well as storage. Supported accelerators can execute low-precision matrix multiplications faster and move smaller tensors using fewer bytes. Actual gains depend on matrix shapes, available kernels, and how much time the run spends on those operations. Loading data, communicating between GPUs, or running small operations can limit the overall improvement.

BF16 is often a practical starting point on supported hardware because of its wider range; FP16 can be useful where it is the supported or better-performing format, with scaling and range checks. Either choice changes rounding throughout training. Compare loss behavior and model quality with a trusted baseline, and measure memory and step time for the actual workload. Mixed precision is successful when those lower-cost calculations preserve the training behavior the model needs.