Mixed precision with autocast and GradScaler

Setup#

pip install torch

Runs on CPU and GPU. On CUDA the half format is float16; on CPU, autocast uses bfloat16. Outputs captured with torch 2.5.1 on an NVIDIA RTX A6000.

See what autocast actually changes#

import torch
import torch.nn as nn

device = "cuda" if torch.cuda.is_available() else "cpu"
amp_dtype = torch.float16 if device == "cuda" else torch.bfloat16

model = nn.Linear(64, 64).to(device)
x = torch.randn(8, 64, device=device)

with torch.autocast(device_type=device, dtype=amp_dtype):
    matmul_out = model(x)                # matmuls run in the small dtype
    summed = matmul_out.sum()            # reductions stay in float32

print("inside autocast, linear output:", matmul_out.dtype)
print("inside autocast, the sum:      ", summed.dtype)
print("the weights themselves:        ", model.weight.dtype)

inside autocast, linear output: torch.float16
inside autocast, the sum:       torch.float32
the weights themselves:         torch.float32

Three facts in three lines. autocast is a context manager that reroutes each operation to a chosen precision. Matrix multiplies — safe and fast in half precision — get float16. Sums and other error-accumulating reductions stay float32. And the weights are never touched.

The full training pattern#

import torch
import torch.nn as nn

device = "cuda" if torch.cuda.is_available() else "cpu"
amp_dtype = torch.float16 if device == "cuda" else torch.bfloat16
torch.manual_seed(0)

model = nn.Sequential(nn.Linear(32, 64), nn.ReLU(), nn.Linear(64, 1)).to(device)
opt = torch.optim.SGD(model.parameters(), lr=0.05)
scaler = torch.amp.GradScaler(device, enabled=(amp_dtype == torch.float16))

x = torch.randn(64, 32, device=device)
y = x.mean(dim=1, keepdim=True)

for step in range(5):
    opt.zero_grad(set_to_none=True)
    with torch.autocast(device_type=device, dtype=amp_dtype):
        loss = nn.functional.mse_loss(model(x), y)
    scaler.scale(loss).backward()     # multiply the loss up before backward
    before = scaler.get_scale()
    scaler.step(opt)                  # unscales gradients, skips step if inf
    scaler.update()                   # adjusts the scale for next time
    skipped = "  <- overflow, step skipped" if scaler.get_scale() < before else ""
    print(f"step {step}  loss {loss.item():.4f}  scale {scaler.get_scale():.0f}{skipped}")

step 0  loss 0.0701  scale 65536
step 1  loss 0.0618  scale 65536
step 2  loss 0.0562  scale 65536
step 3  loss 0.0517  scale 65536
step 4  loss 0.0478  scale 65536

scaler.scale(loss).backward() — the loss is multiplied by the scale (65536 to start). Every gradient inherits that factor, lifting tiny values out of the zero zone. scaler.step(opt) divides the gradients back down, checks them for infinities, and skips the update entirely if any appear. scaler.update() halves the scale after an overflow and slowly grows it during calm stretches. A skipped step is normal early in training — the marker line exists to make it visible.

Only the forward pass goes inside autocast. backward() and opt.step() stay outside — autocast handles their precision automatically.

Common mistakes#

Calling .backward() on the unscaled loss. Everything runs, nothing errors, and small gradients quietly die to zero. The model trains worse for no visible reason. The scaler exists for exactly this.

Using float16 without the scaler. Same silent failure. If you want to skip the scaler, use bfloat16 instead — it trades detail for range and cannot underflow this way. On Ampere or newer GPUs, dtype=torch.bfloat16 with no GradScaler is a clean, popular setup.

Clipping gradients while they are still scaled. Clipping at 1.0 when gradients carry a factor of 65536 clips everything. Call scaler.unscale_(opt) first, then clip, then scaler.step(opt).

Putting the whole loop inside autocast. The optimizer update must run in float32. Keep autocast around the forward pass and loss only.

Try it yourself#

Change y to x.sum(dim=1, keepdim=True) — a hundred-fold larger target. Run it and watch the overflow marker appear as the scaler backs off. Then switch amp_dtype to torch.bfloat16 (GPU) and see the scaler sit idle.

What to learn next#

• Gradient checkpointing — the other big activation-memory lever.
• Fast attention with scaled_dot_product_attention — kernels that need half precision to shine.
• Gradient clipping — and the unscale-first rule.