Skip to content

Automatic Mixed Precision

Overview

AMP runs the compute-heavy parts (matmuls, convolutions) in fp16 while keeping sensitive parts (loss, softmax, master weights) in fp32 — halving memory and often 1.5-3x accelerating training on tensor-core GPUs. The catch: fp16 underflows small gradients, so you need a GradScaler that upscales the loss to keep gradient magnitudes representable.

  • torch.amp.autocast(device_type="cuda", dtype=torch.float16) — context for the forward pass.
  • torch.amp.GradScaler() — dynamically scales loss up / gradient down, and detects inf/nan.
  • bf16 mode: same API with dtype=torch.bfloat16, no scaler needed (fp32-like range).
  • Modern path: torch.compile + autocast compose; FP8 hardware paths exist on H100+.

💡 fp16 needs GradScaler; bf16 doesn't. Choose by hardware: NVIDIA tensor cores (fp16) vs Ampere+/H (bf16 also fast).

The Minimal Correct AMP Loop (fp16)

import torch, torch.nn as nn

model = nn.Sequential(nn.Linear(256, 256), nn.ReLU(), nn.Linear(256, 16)).cuda()
opt = torch.optim.SGD(model.parameters(), lr=0.1)
scaler = torch.amp.GradScaler("cuda")          # "cuda" or default

for step in range(10):
    x = torch.randn(64, 256).cuda()
    opt.zero_grad()
    with torch.amp.autocast("cuda", dtype=torch.float16):
        loss = model(x).abs().mean()            # matmuls run in fp16
    scaler.scale(loss).backward()               # scale loss -> safe grads
    scaler.step(opt)                            # unscale, clip inf/nan, step
    scaler.update()                             # adjust scale for next iter

The three lines that save you

scaler.scale(loss).backward()   # multiply loss by scale before backward
scaler.step(opt)                # if grads inf/nan -> skip step (dynamic)
scaler.update()                 # tune scale factor

bf16 — No Scaler, Same API

with torch.amp.autocast("cuda", dtype=torch.bfloat16):
    loss = model(x).abs().mean()
loss.backward()
opt.step()     # NO scaler: bf16's exponent range = fp32's

✅ bf16 is the modern default for LLM training: same range as fp32, no scaling headaches, ~2x memory savings.

What autocast Actually Casts

Op class Behavior under fp16 autocast
Matmul, Conv, Linear fp16 (tensor-core fast)
Loss functions (CE, MSE) fp32 (stability)
Softmax, layernorm fp32 internally (reductions)
BatchNorm fp32 by default (mixed)
Elementwise (relu, add) follows inputs
Reductions (sum, mean) fp32 accumulation 
# You can query the current dtype:
with torch.amp.autocast("cuda"):
    print(torch.get_autocast_gpu_dtype())   # torch.float16

Common Pitfalls

  1. Scaler without autocast — pointless overhead; pair them.
  2. .item()/.numpy() inside autocast — sync + cast surprises; hoist out.
  3. Unscaled gradients for manual clipping — must scaler.unscale_(opt) first: 
    scaler.unscale_(opt)
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
scaler.step(opt); scaler.update()
  4. Forgot scaler.update() — scale never adapts; early overflows → NaN.
  5. Mixed devices: autocast("cuda") vs autocast("cpu") — CPU autocast (bfloat16) separate context.

GradScaler Explained

scale starts at 65536 (or init_scale)
  - overflow detected (inf/nan) -> scale *= backoff_factor (0.5), skip step
  - no overflow for growth_interval steps -> scale *= growth_factor (2)

scaler = torch.amp.GradScaler("cuda", init_scale=2**16, growth_factor=2.0,
                              backoff_factor=0.5, growth_interval=2000)
print("current scale:", scaler.get_scale().item())

AMP + torch.compile

compiled = torch.compile(model)
with torch.amp.autocast("cuda"):
    loss = compiled(x).abs().mean()
scaler.scale(loss).backward(); scaler.step(opt); scaler.update()

✅ Compile + AMP is the standard "free 2x" combo. Ensure the outer call sites use autocast; Inductor respects the dtype inside.

Verifying It Actually Helped

import time
def bench(fn, n=50):
    for _ in range(3): fn()
    torch.cuda.synchronize()
    t0 = time.perf_counter()
    for _ in range(n): fn()
    torch.cuda.synchronize()
    return (time.perf_counter() - t0)/n * 1e3

x = torch.randn(64, 256).cuda()
fp32_time = bench(lambda: model(x))
with torch.amp.autocast("cuda"):
    amp_time = bench(lambda: model(x))
print(f"fp32 {fp32_time:.2f} ms | amp {amp_time:.2f} ms | speedup {fp32_time/amp_time:.2f}x")

Key Takeaways

  • fp16 AMP = autocast + GradScaler; bf16 = autocast only. Know which you're on.
  • scaler.step() skips optimizer steps on overflow; update() adapts scale.
  • Clip grads with scaler.unscale_ first; hoist .item() out of autocast.
  • Compose with torch.compile for the standard fast path.
  • Measure the speedup — and always sanity-check loss vs fp32 baseline.