Mixed Precision Training: FP16/FP32, GradScaler, and AMP

🎯 Learning Objectives

• Understand FP16 (half-precision) and its trade-offs
• Implement Automatic Mixed Precision (AMP) training
• Use GradScaler to prevent underflow in FP16 gradients
• Benchmark throughput improvements from mixed precision
• Handle edge cases in mixed precision training

📋 Prerequisites

• PyTorch basics (Week 1): Training loops
• Distributed Training (Week 5): Optional but helpful

1. 📖 Core Content

1.1 Intuition: Why Mixed Precision?

• 2× less GPU memory → train larger models or batches
• Up to 2× faster training on modern GPUs (Tensor Cores)
• 2× faster data transfer (CPU → GPU) The problem: FP16 has smaller range and precision.

1.2 Automatic Mixed Precision (AMP) in PyTorch

python
Run Inline
# runnable
import torch
import torch.nn as nn
# Traditional FP32 training
model = nn.Linear(100, 10).cuda()
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
loss_fn = nn.CrossEntropyLoss()
inputs = torch.randn(32, 100).cuda()
targets = torch.randint(0, 10, (32,)).cuda()
# Mixed precision training
scaler = torch.cuda.amp.GradScaler()
for step in range(10):
    optimizer.zero_grad()
    # AMP context: FP16 for forward pass
    with torch.cuda.amp.autocast():
        outputs = model(inputs)
        loss = loss_fn(outputs, targets)
    # Scale loss, backward, unscale, step
    scaler.scale(loss).backward()
    scaler.step(optimizer)  # Unscales gradients internally
    scaler.update()          # Updates scale for next iteration
    # Optional: gradient clipping
    # scaler.unscale_(optimizer)  # Must call before clip
    # torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
    # scaler.step(optimizer)
    print(f"Step {step}: Loss = {loss.item():.4f}")

1.3 How GradScaler Works

python
Run Inline
# runnable
import torch
# Initialize scaler with default settings
scaler = torch.cuda.amp.GradScaler(
    init_scale=2.**16,      # Initial scale factor: 65536
    growth_factor=2.0,       # Double scale when no inf/nan
    backoff_factor=0.5,      # Halve scale when inf/nan detected
    growth_interval=2000     # Check every N steps
)
# Scale factor dynamics:
# - Normally: increases by 2× every 2000 steps (up to 2^24)
# - On gradient overflow: decreases by 2×, optimizer step is skipped
print(f"Initial scale: {scaler.get_scale():.1f}")

1.4 When to Use Mixed Precision

1.5 Edge Cases

python
Run Inline
# runnable
# 1. Loss functions that aren't AMP-safe
# Softmax, log_softmax, cross_entropy: OK in FP16
# Some custom losses may need FP32
# 2. BatchNorm in FP16
# PyTorch's BatchNorm automatically uses FP32 internally
# No special handling needed
# 3. Gradient accumulation
scaler = torch.cuda.amp.GradScaler()
for micro_batch in range(4):  # Accumulate over 4 micro-batches
    with torch.cuda.amp.autocast():
        loss = model(data) / 4  # Average across micro-batches
    scaler.scale(loss).backward()
scaler.step(optimizer)
scaler.update()
optimizer.zero_grad()
# 4. Model with FP32-only layers
# Set dtype for specific layers to FP32
model.half()  # Convert full model to FP16
# Or keep specific layers in FP32:
for layer in model.fp32_layers:
    layer.float()

1.6 Why This Matters

• GPT-3 (175B params): trained with mixed precision (wouldn't fit in FP32)
• Stable Diffusion: uses AMP for faster training
• BERT Large: 2-3× speedup with AMP on V100 GPUs Every modern training framework (PyTorch Lightning, HF Trainer, JAX) enables mixed precision by default.

2. 📐 Key Formulas / Concepts

3. ⚠️ Common Pitfalls

Pitfall 1: Using AMP Without Tensor Cores

Pitfall 2: Forgetting scaler.update()

Pitfall 3: Gradient Clipping Before Unscaling

4. 📝 Practice Questions

Q1: You're training a ResNet-50 on a V100 GPU. FP32 training takes 4 hours. FP16 with AMP takes 1.5 hours. What causes the 2.67× speedup?

Speedup sources:

1. Tensor Cores: V100 has Tensor Cores that perform FP16 matrix multiply at 8× the rate of FP32. Not all ops use Tensor Cores, so the real speedup is 2-4×.
2. Memory bandwidth: FP16 uses half the memory, reducing data transfer time.
3. Larger batch: If batch size was doubled (same memory), fewer iterations per epoch.

The actual speedup depends on how much of the computation uses Tensor Cores. ResNet-50 has many convolutions (Tensor Core-friendly), so 2.67× is reasonable. Q2: Your mixed precision training produces NaN losses every 50 steps. What's happening and how do you fix it?

Cause: Gradient overflow — the scaled gradients exceed FP16's range (~6e4). The scaler detects these as inf/nan and skips the step.

Fixes:

1. Reduce initial scale: GradScaler(init_scale=2.**10) starts with a smaller scale
2. Increase backoff_factor: backoff_factor=0.25 reduces scale more aggressively on overflow
3. Check model architecture: Some operations (softmax with large logits, exp of large numbers) are prone to overflow in FP16
4. Use BF16 (if available on Ampere+ GPUs): Better range eliminates most overflow issues
5. Keep problematic layers in FP32: with torch.cuda.amp.autocast(enabled=False): for unstable operations

The scaler should handle this automatically by reducing the scale. If it happens too frequently, investigate the root cause.

5. 🔗 Cross-References

• Previous: Distributed Training (Week 5)
• Next: Profiling & Optimization (Week 7)
• Related: Training & Debugging (Week 4)
• External: NVIDIA AMP Documentation Join Discord PreviousDistributed TrainingNextProfiling & Optimization