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()
1. Forgot scaler.update()— scale never adapts; early overflows → NaN. 2. 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.

-