What is mixed precision training in PyTorch and how do you implement it with torch.cuda.amp?

Mixed precision training runs most operations in FP16 (or BF16) for speed while keeping a master copy of weights in FP32 for numerical stability. Modern GPUs (Volta and later) have dedicated hardware (Tensor Cores) that make FP16 matrix multiplication significantly faster than FP32.

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 to prevent FP16 underflow

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

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

    # autocast: automatically runs eligible ops in FP16/BF16
    with autocast(device_type="cuda", dtype=torch.float16):
        y_hat = model(x)                  # matmul runs in FP16 — faster!
        loss  = nn.MSELoss()(y_hat, y)

    # Loss scaling: inflate loss before backward to prevent small
    # gradients from underflowing to zero in FP16's limited range
    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)            # restore original gradient magnitudes
    nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    scaler.step(optimizer)                # skips the step if grads are inf/NaN
    scaler.update()                       # adjusts scale factor for next iteration

# BFloat16: no GradScaler needed (same exponent range as FP32)
with autocast(device_type="cuda", dtype=torch.bfloat16):
    y_hat = model(x)   # no underflow risk — scaling unnecessary

Invest now in Acorns!!! 🚀
Join Acorns and get your $5 bonus!

Acorns is a micro-investing app that automatically invests your "spare change" from daily purchases into diversified, expert-built portfolios of ETFs. It is designed for beginners, allowing you to start investing with as little as $5. The service automates saving and investing. Disclosure: I may receive a referral bonus.

Invest now!!! Get Free equity stock (US, UK only)!

Use Robinhood app to invest in stocks. It is safe and secure. Use the Referral link to claim your free stock when you sign up!.

The Robinhood app makes it easy to trade stocks, crypto and more.

Webull! Receive free stock by signing up using the link: Webull signup.