What is mixed precision training and how does it speed up deep learning with torch.cuda.amp?

Modern GPUs (Volta and later) have dedicated hardware for 16-bit floating-point operations (FP16 / BFloat16) that can be 2–8× faster than FP32 for matrix multiplications. Mixed precision training runs the forward pass and gradient computations in FP16 (or BF16) for speed, while maintaining a master copy of the weights in FP32 for numerical precision during the optimizer update.

Loss scaling addresses a key challenge: FP16's limited dynamic range (smallest positive ≈ 6×10⁻⁸) can cause small gradient values to underflow to zero. The scaler multiplies the loss by a large scalar before backward (inflating gradients into FP16's representable range), then divides the gradients back before the optimizer step. PyTorch's GradScaler automates this and dynamically adjusts the scale factor.

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

model     = nn.Linear(1024, 512).cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
scaler    = GradScaler()           # manages loss scaling automatically

x = torch.randn(256, 1024).cuda()
y = torch.randn(256, 512).cuda()

for step in range(100):
    optimizer.zero_grad()

    # autocast: runs eligible ops in FP16 automatically
    with autocast(device_type='cuda', dtype=torch.float16):
        y_hat = model(x)           # FP16 matrix multiply
        loss  = nn.MSELoss()(y_hat, y)

    # Scale loss -> backward in FP16 -> unscale gradients -> optimizer step
    scaler.scale(loss).backward()  # inflate loss to prevent underflow
    scaler.unscale_(optimizer)     # restore original gradient magnitudes
    nn.utils.clip_grad_norm_(model.parameters(), 1.0)  # clip after unscale
    scaler.step(optimizer)         # skip step if gradients are inf/NaN
    scaler.update()                # adjust scale factor for next step

# BFloat16 (bfloat16): available on A100+ GPUs
# - Same exponent range as FP32 (no underflow problem -> no scaler needed)
# - Less precision (7-bit mantissa vs 10-bit for FP16)
with autocast(device_type='cuda', dtype=torch.bfloat16):
    y_hat = model(x)  # no scaler needed with BF16