Mixed Precision Training

What is Automatic Mixed Precision (AMP)?

Traditionally, deep neural networks were trained using 32-bit Single Precision Floats (FP32):

• 1 sign bit, 8 exponent bits, 23 mantissa bits (4 bytes per parameter).

Automatic Mixed Precision (AMP - Micikevicius et al., 2017 / NVIDIA) accelerates training by using 16-bit floating point formats for heavy matrix multiplications while keeping FP32 master weights for precision updates.

  FP32 Precision:  [ Sign (1b) | Exponent (8b)  | Mantissa (23b)          ] (4 Bytes)
  FP16 Precision:  [ Sign (1b) | Exponent (5b)  | Mantissa (10b)  ]         (2 Bytes)
  BF16 Precision:  [ Sign (1b) | Exponent (8b)  | Mantissa (7b)   ]         (2 Bytes)

AMP delivers $2\times$ faster training throughput on NVIDIA Tensor Cores while cutting VRAM memory usage by nearly $50%$.

The Mixed Precision Workflow

┌─────────────────────────────────────────────────────────────┐
│ 1. FORWARD PASS IN FP16 / BF16                               │
│ Cast input tensors and layer weights to FP16/BF16.          │
│ Execute heavy GEMM Matrix Multiplications on Tensor Cores.  │
└──────────────────────────┬──────────────────────────────────┘
                           │
                           ▼
┌──────────────────────────┴──────────────────────────────────┐
│ 2. LOSS SCALING & BACKWARD PASS                             │
│ Multiply Loss by Scale Factor (e.g. 2^15) to prevent FP16   │
│ gradient underflow. Compute backward gradients.             │
└──────────────────────────┬──────────────────────────────────┘
                           │
                           ▼
┌──────────────────────────┴──────────────────────────────────┐
│ 3. FP32 MASTER WEIGHT OPTIMIZER UPDATE                      │
│ Unscale gradients back to FP32. Accumulate weight updates   │
│ into FP32 Master Weights to prevent precision loss.         │
└─────────────────────────────────────────────────────────────┘

FP16 vs BFloat16 (BF16)

┌──────────────────────────┬──────────────────────────┐
│ FP16 (IEEE Half Precision)│ BF16 (Brain Float 16)   │
├──────────────────────────┼──────────────────────────┤
│ 5 Exponent / 10 Mantissa │ 8 Exponent / 7 Mantissa  │
│ Small Dynamic Range.     │ HUGE Dynamic Range       │
│ Suffers from Underflow!  │ (Matches FP32 Range!).   │
│ REQUIRES Loss Scaling!   │ NO Loss Scaling Needed!  │
└──────────────────────────┴──────────────────────────┘

Why BFloat16 Has Replaced FP16

BFloat16 (introduced by Google Brain) keeps the exact same 8-bit Exponent width as FP32, sacrificing a few bits of mantissa precision.

Because BF16 covers the same dynamic range as FP32, it never suffers from gradient underflow or overflow, rendering complex Loss Scaling mechanisms completely unnecessary!

BF16 is supported on NVIDIA Ampere (A100, H100, RTX 3000+) GPUs and Google TPUs.

PyTorch Implementation

import torch
scaler = torch.cuda.amp.GradScaler() # For FP16 Loss Scaling

for inputs, targets in dataloader:
    optimizer.zero_grad()
    
    # 1. Forward Pass in Mixed Precision:
    with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
        outputs = model(inputs)
        loss = criterion(outputs, targets)
        
    # 2. Backward Pass & Scaled Optimizer Step:
    scaler.scale(loss).backward()
    scaler.step(optimizer)
    scaler.update()

Say this out loud

Automatic Mixed Precision executes matrix multiplications in 16 bit lower precision (FP16 or BF16) to leverage GPU Tensor Cores, while keeping master weights in FP32 to prevent tiny gradient updates from underflowing to zero. BFloat16 shares the 8 bit exponent range of FP32, eliminating underflow and accelerating training by 2x with zero loss in accuracy.

Followups to expect

1. What is FP8 Training (H100 / Blackwell)? Using 8-bit floating point formats (E4M3 for forward weights and E5M2 for backward gradients) to achieve up to 4x higher FLOP throughput on modern GPU architectures.
2. Why do Loss Values occasionally output NaN during FP16 training? Un-scaled FP16 gradients exceeded maximum float value ($65,504$), causing numerical overflow ($\infty$), which yields NaN during optimizer steps. Fix by switching to BF16.