TopolearnMixed precision typically gives 2-3x faster training and halves activation memory on modern GPUs.
Mixed precision keeps weights in float32 but runs most computation in a 16-bit type. Tensor cores on modern GPUs are much faster in 16-bit, and activations take half the memory.
torch.autocast(device_type, dtype) is a context manager around the forward pass and loss. Inside it, ops that benefit (matmul, convolutions) run in low precision, while numerically sensitive ones (softmax, losses, reductions) are kept in float32 automatically. Run backward() outside the context.
Pick the dtype:
- bfloat16 has float32's exponent range, so it rarely underflows. It is the default choice on Ampere-or-newer GPUs and TPUs, and needs no extra machinery.
- float16 has more mantissa but a narrow range; small gradients underflow to zero. Use
torch.amp.GradScaler, which multiplies the loss by a large factor beforebackward(), unscales before the step, and skips steps that produced inf or NaN.
If you clip gradients with a scaler, call scaler.unscale_(opt) before clipping.
Separately, torch.set_float32_matmul_precision("high") lets float32 matmuls use TF32 tensor cores on NVIDIA GPUs, a near-free speedup for code you do not convert.
Keep the model's weights in float32 (do not call model.half() for training); autocast handles the casts.
Example
import torch
from torch import nn
device = "cuda" if torch.cuda.is_available() else "cpu"
dtype = torch.float16 if device == "cuda" else torch.bfloat16
model = nn.Sequential(nn.Linear(512, 512), nn.ReLU(), nn.Linear(512, 10)).to(device)
opt = torch.optim.AdamW(model.parameters(), lr=1e-3)
scaler = torch.amp.GradScaler(device, enabled=(dtype == torch.float16))
x = torch.randn(64, 512, device=device)
y = torch.randint(0, 10, (64,), device=device)
with torch.autocast(device_type=device, dtype=dtype):
logits = model(x) # matmuls in 16-bit
loss = nn.functional.cross_entropy(logits, y)
print(logits.dtype, loss.dtype) # 16-bit logits, float32 loss
scaler.scale(loss).backward() # backward outside autocast
scaler.step(opt) # unscales; skips inf/NaN steps
scaler.update()
opt.zero_grad(set_to_none=True)Remember
- autocast wraps forward + loss; backward goes outside
- bfloat16 needs no scaler; float16 needs GradScaler
- Keep master weights in float32
Watch out for
- Clipping gradients that are still scaled; call
scaler.unscale_(opt)first.