Mixed Precision — Module 6.2

How This Lesson Fits the Module

FP32 training is safe but slow and memory-hungry. Mixed precision runs most matmuls in float16 (or bfloat16) while keeping master weights in float32, often doubling throughput on modern GPUs with tensor cores.

PyTorch’s torch.cuda.amp (automatic mixed precision) wraps the training loop with autocast and GradScaler to prevent underflow.

Prior Lesson Gradient Clipping integrates with AMP via scaler.unscale_. Module 6.3 dives deeper into FP16, BF16, and hardware.

Learning Objectives

By the end of this lesson, students should be able to:

Wrap forward pass in torch.cuda.amp.autocast().
Use GradScaler for scaled backward and optimizer steps.
Explain why float16 needs loss scaling to avoid gradient underflow.
Choose float16 vs bfloat16 based on GPU generation.
Combine AMP with gradient clipping correctly.
Recognize layers that should stay in float32 (loss, softmax in some cases).

FP32 vs Mixed Precision

Aspect	FP32	Mixed (AMP)
Speed on Tensor Cores	Baseline	Often 1.5–2× faster
Memory	Higher	Lower activations
Numerical range	Wide	FP16 narrow; BF16 similar to FP32 exponent
Implementation	Default	`autocast` + `GradScaler`

AMP Training Loop

autocast selects float16 for conv/linear ops; GradScaler multiplies loss before backward, then unscales gradients inside the optimizer step.

from torch.cuda.amp import autocast, GradScaler

scaler = GradScaler(enabled=torch.cuda.is_available())

model.train()
for inputs, targets in train_loader:
    inputs, targets = inputs.to(device), targets.to(device)
    optimizer.zero_grad()

    with autocast(dtype=torch.float16):
        outputs = model(inputs)
        loss = criterion(outputs, targets)

    scaler.scale(loss).backward()
    scaler.unscale_(optimizer)
    torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
    scaler.step(optimizer)
    scaler.update()

Validation in Mixed Precision

Inference can use autocast without GradScaler. Many teams validate in FP32 for metric stability on small validation sets.

model.eval()
with torch.no_grad(), autocast(dtype=torch.float16):
    for inputs, targets in val_loader:
        outputs = model(inputs.to(device))
        ...

Critical Mistake — Skipping scaler.update()

After scaler.step(), always call scaler.update(). It adjusts loss scale when inf/NaN gradients appear. Omitting it freezes a bad scale and silently breaks training.

BFloat16 on Ampere and Newer

BF16 shares exponent range with FP32—often no GradScaler needed. Use autocast(dtype=torch.bfloat16) on A100/H100 class GPUs.

with autocast(dtype=torch.bfloat16):
    outputs = model(inputs)
    loss = criterion(outputs, targets)

loss.backward()  # scaler optional for bf16 on many workloads
optimizer.step()

Engineering Habit — Benchmark Before Shipping

AMP speedups vary by model (attention-heavy vs MLP). Profile one epoch FP32 vs AMP on your hardware; confirm accuracy within tolerance.

Knowledge Check

Short Answer: What does autocast do? Answer: Runs eligible ops in lower precision automatically.
True/False: GradScaler is used during validation. Answer: False—only training backward needs scaling (FP16).
Multiple Choice: Gradients underflow in FP16 because: (a) range too small, (b) LR too low only, (c) no GPU. Answer: (a).
Short Answer: Order after scale(loss).backward()? Answer: unscale_, clip (optional), scaler.step, scaler.update.
Short Answer: BF16 advantage over FP16? Answer: Wider exponent range, often more stable without scaling.
True/False: Master weights stay FP32 in standard AMP. Answer: True—optimizer updates FP32 copies.
Multiple Choice: Tensor Cores accelerate: (a) mixed precision matmul, (b) data loading, (c) checkpoint I/O. Answer: (a).
Short Answer: When is AMP disabled safely? Answer: CPU training or when enabled=False in GradScaler/autocast.
Short Answer: Why clip after unscale_? Answer: Clipping must use real gradient magnitudes.
Multiple Choice: Module 6.3 covers hardware details: (a) CPU vs GPU, (b) sklearn pipelines, (c) ETL. Answer: (a).

Key Takeaways

AMP = autocast forward + GradScaler backward (for FP16).
Always scaler.update() after scaler.step().
Consider BF16 on modern NVIDIA/TPU hardware for simpler training.
Next: Distributed Training—scale across GPUs.

Trainer’s Guide

Hands-on idea: Measure images/sec and peak VRAM for ResNet-18 FP32 vs AMP on one GPU.

Discussion prompt: When would you refuse mixed precision (e.g. scientific simulation nets)?

What’s Next Scale to multiple GPUs with Distributed Training, then explore hardware in Module 6.3.