Mixed Precision, Tensor Cores, and Loss Scaling

Core Concept

Mixed precision uses FP16 inputs and Tensor Cores for suitable GEMM/Conv work while retaining FP32 accumulation, reductions, and sensitive operations.

Why It Affects Performance

Lower precision reduces traffic and increases matrix throughput, but FP16 has limited range. Forward values can overflow and small gradients can underflow before reaching an FP32 accumulator.

Example or Pseudocode

scaler = torch.amp.GradScaler('cuda')
with torch.autocast('cuda', dtype=torch.float16):
    loss = loss_fn(model(x), target)
scaler.scale(loss).backward()
scaler.step(optimizer); scaler.update()

Profiling and Performance Judgment

Verify HMMA/Tensor instructions and Tensor Pipe activity. Compare cast kernels, end-to-end time, NaN/inf, gradient norms, and model quality.

Five-Minute Exercise

Does loss scaling mainly address overflow or small-gradient underflow? Can it fix unstable forward formulas?