Mixed Precision

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader

# 1) Setup
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model = MyModel().to(device)
opt   = optim.AdamW(model.parameters(), lr=1e-4)
loss_fn = nn.CrossEntropyLoss()
loader  = DataLoader(train_set, batch_size=64, num_workers=4, pin_memory=True)

# 2) Mixed precision training loop (FP16 / AMP)
scaler = torch.cuda.amp.GradScaler()

for epoch in range(epochs):
    for x, y in loader:
        x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
        opt.zero_grad(set_to_none=True)

        with torch.autocast(device_type='cuda', dtype=torch.float16):
            logits = model(x)
            loss   = loss_fn(logits, y)

        scaler.scale(loss).backward()           # scaled to prevent underflow
        scaler.unscale_(opt)                    # before clipping
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        scaler.step(opt)                         # step optimizer if no inf/nan
        scaler.update()                          # adjust scale factor

# 3) bfloat16 — alternative on Ampere+ GPUs (no scaler needed)
for x, y in loader:
    x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
    opt.zero_grad(set_to_none=True)
    with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
        logits = model(x)
        loss = loss_fn(logits, y)
    loss.backward()
    opt.step()

# bfloat16:
# • Wider range than fp16 (same as fp32 exponent) — no scaler needed
# • Slightly less precision in mantissa (rarely an issue)
# • Ampere/Hopper GPUs natively support it; older cards don't

# 4) When to choose which
# fp16 + GradScaler:   pre-Ampere GPUs, broad compatibility
# bfloat16:           Ampere+ GPUs, simpler loop, stable training
# fp32:               debugging, tiny models, fallback

# 5) Why GradScaler
# Some gradients are tiny — fp16 can't represent them; they underflow to 0.
# GradScaler multiplies the loss by a large factor before backward, then unscales gradients.
# If NaN/inf appears, scaler skips the step and halves the scale; if many successful steps, doubles it.

# 6) Where autocast applies
# • Matmul / conv → fp16
# • Element-wise + reductions → mixed (whatever's safe)
# • Normalisation (LayerNorm, BatchNorm) → fp32
# • Loss → kept fp32
# • Optimizer state → fp32 (no change)
# autocast is a CONTEXT manager; it converts ops, not tensors.

# 7) Inference with mixed precision
model.eval()
with torch.no_grad(), torch.autocast(device_type='cuda', dtype=torch.float16):
    output = model(inputs)

# For deployment, often converted to ONNX or TensorRT for production fp16/int8.

# 8) Common patterns

# Gradient accumulation
acc_steps = 4
for step, (x, y) in enumerate(loader):
    x, y = x.to(device), y.to(device)
    with torch.autocast(device_type='cuda', dtype=torch.float16):
        loss = loss_fn(model(x), y) / acc_steps
    scaler.scale(loss).backward()
    if (step + 1) % acc_steps == 0:
        scaler.unscale_(opt)
        torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
        scaler.step(opt); scaler.update()
        opt.zero_grad(set_to_none=True)

# 9) DDP + AMP — works out of the box
from torch.nn.parallel import DistributedDataParallel as DDP
model = DDP(model.to(device), device_ids=[local_rank])
# Use AMP inside the standard training loop; DDP wraps cleanly.

# 10) Detecting / debugging issues
# • Loss = NaN at random steps → fp16 underflow; check scaler state, lower LR, or switch to bf16
# • Inf gradients → exploding; clip_grad_norm_ before scaler.step
# • Accuracy drops 1-2% with fp16 → enable bfloat16 OR train final epoch in fp32
# • TF32 (Ampere): torch.backends.cuda.matmul.allow_tf32 = True — speeds matmul, slight precision loss

# 11) torch.compile + AMP (PyTorch 2+)
model = torch.compile(model, mode='reduce-overhead')
# Combines well with autocast; further 10-30% speedup on supported GPUs.

# 12) FlashAttention + AMP — Transformers
# Modern transformer libraries auto-use FlashAttention when in fp16/bf16.
# Hugging Face Transformers: pip install flash-attn; model.attn_implementation = 'flash_attention_2'

# 13) Memory profile
torch.cuda.memory_summary()
torch.cuda.max_memory_allocated() / 1e9   # GB
# Compare full-precision vs mixed-precision peak memory; expect ~50% reduction.

# 14) When NOT to use mixed precision
# • Tiny models / mostly-elementwise ops → barely any speedup
# • Numerical sensitivity (some physics simulations) → stay fp32
# • Pretrained models with known fp32 weights — fine-tune in fp32 if eval changes drastically
# • CPU-only — torch.autocast(device_type='cpu') uses bfloat16 if supported; gains less dramatic

# 15) Performance checklist
# • Use AdamW with sensible LR for fp16 (lower than fp32 sometimes)
# • Increase batch size to fill memory after mixed-precision reduces footprint
# • num_workers + pin_memory for the DataLoader
# • channels_last memory format for convnets (model = model.to(memory_format=torch.channels_last))
# • Enable cudnn.benchmark for static input shapes

# 16) Common bugs
# • Forgot scaler.scale() — gradients underflow; loss explodes silently
# • Calling backward() outside autocast — uses last cast type; mostly fine but explicit is better
# • clip_grad_norm before scaler.unscale_ → clips SCALED gradients; meaningless
# • Saving model in fp16 weights when ecosystem expects fp32 — convert before save
# • Manually moving tensors to half() inside the loop — defeats autocast; trust it
# • bf16 on a pre-Ampere card → silent FP32 fallback; no speedup
# • Multi-GPU + GradScaler — call scaler.unscale_ once on the local optimizer; DDP averages
# • Forgotten autocast on inference → no speedup, weird kernel switching

from torch import amp
scaler = amp.GradScaler('cuda')
with amp.autocast('cuda', dtype=torch.bfloat16):
    pred = model(xb)
    loss = loss_fn(pred, yb)
scaler.scale(loss).backward()
scaler.step(opt)
scaler.update()