PyTorch Mixed Precision Training with torch.cuda.amp

When you train a PyTorch model by default, every tensor operation runs in FP32 precision. Python PyTorch mixed precision training changes part of that computation to FP16, which reduces memory footprint and can accelerate training on GPUs with Tensor Cores. The implementation in PyTorch is built around two components: torch.cuda.amp.autocast and torch.cuda.amp.GradScaler. This article explains how they work together, how to integrate them into a training loop, and where the approach breaks down.

What Mixed Precision Actually Changes

Mixed precision does not mean converting the entire model to FP16. It means running selected operations in FP16 while keeping others in FP32. The reason is numerical stability: some operations, like reductions and softmax, lose too much precision in FP16, while matrix multiplications and convolutions benefit from the faster FP16 math units on modern GPUs.

autocast handles this selection automatically. When you wrap a forward pass in the autocast context, PyTorch chooses the precision for each operation based on a built-in per-op policy. For example, torch.matmul and torch.nn.Conv2d run in FP16 when the inputs are FP32, while torch.softmax and torch.nn.LayerNorm are forced to FP32 regardless of input type. You do not manually cast tensors; the context manager decides.

This matters because casting everything to FP16 manually is error-prone. If you call .half() on your model and inputs, operations like batch normalization and softmax may produce unstable results. The autocast policy exists precisely to avoid those cases.

The Two Core Components: autocast and GradScaler

The API, introduced in PyTorch 1.6, consists of two pieces:

• torch.cuda.amp.autocast is a context manager (also usable as a decorator) that enables automatic precision selection for the wrapped code.
• torch.cuda.amp.GradScaler is a helper that scales the loss before backpropagation and rescales the gradients afterward.

GradScaler exists because gradients in FP16 can underflow to zero. FP16 has a limited exponent range; very small gradient values become zero, and the model stops learning. Scaling the loss by a large factor before the backward pass shifts gradients into a representable range. After the backward pass, the gradients are divided back by the same factor before the optimizer step.

The scaling factor is dynamic. GradScaler starts with a default initial scale of 65536.0 and adjusts it during training: it increases the scale when no inf or NaN gradients occur, and decreases it when they do.

A Minimal Training Loop with Mixed Precision

Here is the minimal change to a standard PyTorch training loop:

import torch
import torch.nn as nn
from torch.cuda.amp import autocast, GradScaler

model = nn.Sequential(
    nn.Linear(512, 256),
    nn.ReLU(),
    nn.Linear(256, 10),
).cuda()

optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
loss_fn = nn.CrossEntropyLoss()
scaler = GradScaler()

for batch, (inputs, targets) in enumerate(dataloader):
    inputs = inputs.cuda()
    targets = targets.cuda()

    optimizer.zero_grad()

    with autocast():
        outputs = model(inputs)
        loss = loss_fn(outputs, targets)

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

The pattern is consistent: the forward pass runs inside autocast(), the loss is scaled before .backward(), and scaler.step(optimizer) replaces the direct optimizer.step(). The scaler.update() call adjusts the scale factor for the next iteration.

The reason scaler.step() is separate from scaler.update() is that the optimizer step can be skipped. If the scaled gradients contain inf or NaN values, scaler.step() detects that, skips the optimizer step, and leaves the model parameters unchanged. The scale factor is then reduced by update() so the next iteration is more likely to produce finite gradients.

Why Gradient Underflow Is the Main Risk

The dominant failure mode in mixed precision training is not overflow but underflow. FP16 represents numbers with about 10 bits of mantissa and a limited exponent range; the smallest normal FP16 value is roughly 6.1e-5. Gradients in deep networks frequently fall below this threshold, especially in early layers or with small learning rates.

Loss scaling works because multiplying the loss by a large constant before backpropagation multiplies all gradients by the same constant. A gradient that would be 1e-6 in FP32 becomes 0.0655 in FP16 when the scale is 65536, which is well within the representable range. After the backward pass, the gradients are unscaled back to their original magnitude before the optimizer consumes them.

This is why you should not call optimizer.step() directly when using GradScaler. The unscaling step must happen before the optimizer reads the gradients, and scaler.step() performs that internally.

When Mixed Precision Does Not Help

Mixed precision is not universally faster. The speedup depends on the GPU's FP16 throughput. GPUs with Tensor Cores, such as the V100, T4, A100, and RTX 20-series and newer, can execute FP16 matrix operations at a much higher rate than FP32. On those devices, the speedup is real, particularly for compute-bound layers like convolutions and large linear layers.

On GPUs without Tensor Cores, FP16 operations are often slower than FP32 because they require conversion and do not benefit from specialized hardware. On CPU, the CUDA-specific autocast context manager does not change the computation, and PyTorch's CPU kernels do not provide the same FP16 acceleration. If you are training on CPU, mixed precision will not help.

The model size also matters. If the model is small and the bottleneck is Python overhead or data loading rather than tensor math, the FP16 speedup will be masked. Measure the training step time before and after enabling mixed precision rather than assuming the improvement.

Failure Modes and How to Diagnose Them

The most common symptom is a loss that becomes NaN or inf after enabling mixed precision. The first thing to check is whether the model uses operations that autocast does not cover, such as custom autograd functions that assume FP32 inputs. If you write a custom torch.autograd.Function, you are responsible for handling FP16 inputs explicitly; autocast does not intercept the internals of a custom function.

Another common issue is a scale factor that collapses to its minimum value. GradScaler has a floor (default 1.0) and will keep reducing the scale when inf or NaN gradients keep appearing. If the scale stays at the minimum and training is unstable, the problem is usually not the scaling but the model itself, for example an exploding gradient that would also break FP32 training.

When saving and loading a checkpoint, the scaler state must be saved as well. The scaler's scale factor and growth history are part of the training state; restoring only the model and optimizer will restart the scaler from its initial scale, which can cause a temporary spike in instability.

checkpoint = {
    "model": model.state_dict(),
    "optimizer": optimizer.state_dict(),
    "scaler": scaler.state_dict(),
}
torch.save(checkpoint, "checkpoint.pt")

Compatibility with Distributed Training and Custom Code

In distributed data-parallel training, each process has its own GradScaler, and the scaling logic is per-process. The recommended pattern is to enable mixed precision in each rank's training loop identically. Gradient clipping requires care: you must unscale the gradients before clipping, because clipping operates on the true gradient magnitude, not the scaled one.

scaler.unscale_(optimizer)
torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
scaler.step(optimizer)
scaler.update()

Calling unscale_ before clipping ensures that the clip threshold applies to the real gradient values. If you clip before unscaling, the threshold is effectively multiplied by the scale factor, and the clipping behavior changes with the scale.

Custom autograd functions are the other compatibility boundary. If your model uses a custom torch.autograd.Function, autocast does not automatically convert its inputs. You must either implement the FP16 path inside the function or explicitly cast inputs to FP32 before calling it. The same applies to any operation that autocast does not recognize; the context manager only knows about the operations in its internal policy table.