Mixed Precision: autocast and GradScaler

Half the memory, twice the speed (often)

Mixed-precision training runs most operations in float16 or bfloat16 while keeping a small set of numerically sensitive ops (softmax, normalization, loss reduction) in float32. On Ampere+ GPUs and Apple Silicon this typically gives a 1.5–2x speedup AND roughly halves memory usage for activations.

The modern PyTorch API is torch.amp (not the older torch.cuda.amp, which is deprecated). Two ingredients:

autocast(device_type='cuda', dtype=...) — context manager that picks the right precision per op.
GradScaler(device_type) — only needed for fp16. Scales the loss up before backward to prevent gradient underflow in the narrow fp16 range, then unscales before the optimizer step.

fp16 vs bf16

bfloat16 — same exponent range as float32, less precision in the mantissa. Available on Ampere+ NVIDIA GPUs (A100, RTX 30/40 series, H100) and Apple Silicon. No GradScaler needed. The recommended modern default.
float16 — narrower exponent range. Needs GradScaler. Still relevant for older GPUs (V100, T4, RTX 20-series).

What you actually wrap

Wrap the forward pass and loss computation. Backward stays outside the autocast block. Optimizer step stays outside too.

Code

bfloat16 mixed precision — the modern recipe·python

import torch
import torch.nn as nn
from torch.amp import autocast

model = nn.Linear(1024, 1024).cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()

for x, y in loader:
    x, y = x.cuda(non_blocking=True), y.cuda(non_blocking=True)
    optimizer.zero_grad()

    with autocast(device_type='cuda', dtype=torch.bfloat16):
        out = model(x)
        loss = loss_fn(out, y)

    loss.backward()        # no GradScaler needed for bf16
    optimizer.step()

float16 mixed precision — needs GradScaler·python

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

model = nn.Linear(1024, 1024).cuda()
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3)
loss_fn = nn.CrossEntropyLoss()
scaler = GradScaler('cuda')

for x, y in loader:
    x, y = x.cuda(non_blocking=True), y.cuda(non_blocking=True)
    optimizer.zero_grad()

    with autocast(device_type='cuda', dtype=torch.float16):
        out = model(x)
        loss = loss_fn(out, y)

    scaler.scale(loss).backward()      # scale up to avoid underflow
    scaler.step(optimizer)              # unscale + step (skips on inf/nan)
    scaler.update()                     # adjust scale based on overflow

Inspect what autocast actually picks·python

import torch
from torch.amp import autocast

x = torch.randn(8, 8, device='cuda')

with autocast(device_type='cuda', dtype=torch.bfloat16):
    y_matmul = x @ x        # bfloat16 (fast)
    y_softmax = torch.softmax(x, dim=-1)  # promoted to float32 (numerically sensitive)
    print(y_matmul.dtype, y_softmax.dtype)
# torch.bfloat16 torch.float32

# autocast keeps a built-in op→dtype map. You can usually trust it,
# but for custom ops you may need to specify behavior explicitly.

Exercise

Take any training loop you've written. Time one epoch in float32, then in bfloat16 autocast. On Ampere+ hardware you should see a 1.3–2x speedup with no GradScaler. Watch the loss curve too — if bf16 makes loss noisier or worse, the issue is usually a numerically sensitive op that needs explicit float32 cast.