Mixed Precision Training (AMP)

Overview

Automatic Mixed Precision (AMP) uses lower precision (float16/bfloat16) for most computations while keeping critical operations in float32. Result: 2-3x speedup, 50% memory reduction, virtually no accuracy loss.

• Frameworks: PyTorch AMP, TensorFlow mixed precision
• Key Innovation: Selectively use lower precision without breaking training
• Speedup: 1.5-3x faster training (hardware-dependent)
• Memory: 40-60% reduction in GPU memory usage
• Adoption: Standard in modern training (no longer optional!)

-

The Problem: float32 is Wasteful

Precision Types

Data Type: Size| Precision| Usage
──────────────────────────────────────
float32: 4B| 7 decimals| Full precision
float16: 2B| 4 decimals| Half precision
bfloat16: 2B| 3 decimals| Half precision (less underflow)
float8: 1B| 2 decimals| Extreme compression (emerging)

float32 breakdown:
 - Sign: 1 bit
 - Exponent: 8 bits
 - Mantissa: 23 bits
 - Total: 32 bits = 4 bytes

float16 breakdown:
 - Sign: 1 bit
 - Exponent: 5 bits
 - Mantissa: 10 bits
 - Total: 16 bits = 2 bytes (50% smaller!)

Range:
 - float32: 10^-38 to 10^38
 - float16: 10^-4 to 10^4 (limited for deep networks!)
 - bfloat16: 10^-38 to 10^38 (matches float32 range)
 - Problem: float16 range too small for training!

Why float32 is Overkill

Observation: Not all operations need full precision

Operation Precision Needed Current Use
─────────────────────────────────────────────────
Matrix multiply float16 OK float32 (wasteful!)
Dot products float16/32 float32
Loss computation float32 (stable) float32 (correct)
Gradient updates float32 float32 (necessary)
Softmax float32 float32 (could be relaxed)

Memory for 70B model:
 - In float32: 70B × 4 bytes = 280GB
 - In float16: 70B × 2 bytes = 140GB
 - Savings: 50%!

Speed improvement:
 - GPUs optimized for float16 matrix ops
 - Tensor cores: Can do 2× float16 ops per float32 op
 - A100: 312 TFLOPS float32 vs 625 TFLOPS float16!
 - Speedup: 2-3x possible!

Trade-off:
 - Loss precision in matrix multiplies: OK (hardware does this anyway)
 - But: Need to be careful about numerical stability
 - Solution: AMP (automatic mixed precision)

---

How Automatic Mixed Precision Works

The Strategy

Goal: Use float16 where possible, float32 where needed

Strategy:
1. Matrix multiplies → float16 (fast, hardware optimized)
2. Reductions (softmax, reduce_sum) → float32 (numerical stability)
3. Loss computation → float32 (stability for backward)
4. Gradient computation → float32 (accurate gradients)
5. Parameter updates → float32 (small changes matter)

Problem: float16 range too small for gradients!
Solution: Master copy in float32, compute in float16

Automatic Mixed Precision Flow

Forward Pass:
- ┌─────────────────────┐
 - Input (float32) │
 - ┬──────────────┘
 ↓
- ┌──────────────────────┐
 - Cast to float16 │
 - ┤
 - Linear + GELU │ (float16 computation, fast!)
 - ┤
 - Linear + Softmax │ (float16 + stable reduction)
 - ┤
 - Cast back to float32 │
 - ┬───────────────┘
 ↓
- ┌──────────────────────┐
 - Output (float32) │
 - ┬───────────────┘
 ↓
 Loss (float32)

Backward Pass:
- ┌──────────────────────┐
 - Gradients (float32) │ (computed in float32 for accuracy)
 - ┤
 - Loss scaling applied │ (prevent gradient underflow)
 - ┤
 - Backward through │ (float32 computation)
 - float16 ops │
 - ┘

-

Implementation

PyTorch Automatic Mixed Precision

import torch
from torch.cuda.amp import autocast, GradScaler

# Setup
model = model.to(device)
optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)
scaler = GradScaler() # Prevents gradient underflow

# Training loop
for epoch in range(num_epochs):
 for batch in dataloader:
 x, y = batch
 x, y = x.to(device), y.to(device)

 # Forward pass with autocast
 with autocast():
 output = model(x)
 loss = criterion(output, y)

 # Backward with scaling
 scaler.scale(loss).backward()
 scaler.unscale_(optimizer)
 torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

 # Optimizer step with scaler
 scaler.step(optimizer)
 scaler.update()
 optimizer.zero_grad()

# Key operations:
# autocast()
# GradScaler

Loss Scaling Explained

Problem: Gradient underflow in float16

float16 range: [6×10^-5, 6×10^4]

Typical gradient: 1e-5
 - In float32: 1e-5 (small but representable)
 - In float16: UNDERFLOWS to 0! (below 6e-5 minimum)
 - Result: No gradient signal! Training fails!

Solution: Loss scaling (multiply loss by large factor)

Loss: 0.5
Scaled loss: 0.5 × 1024 = 512

Gradients: 1e-5 × 1024 = 1.024e-2
 - In float16: 0.01 is representable!
 - Backward computes: 0.01 (in float16)
 - Unscale: 0.01 / 1024 = 1e-5 (back to original)
 - Optimizer sees: 1e-5 (correct gradient!)

Dynamic scaling:
 - If overflow: Reduce scaling factor (2x smaller)
 - If no overflow: Try increasing scaling factor
 - PyTorch GradScaler does this automatically!

Manual Control (if needed)

class MixedPrecisionTrainer:
 def __init__(self, model, optimizer):
 self.model = model
 self.optimizer = optimizer
 self.scaler = GradScaler()

 # Manual loss scale control
 self.loss_scale = 1024.0
 self.max_loss_scale = 65536.0
 self.min_loss_scale = 1.0

 def train_step(self, batch):
 x, y = batch

 # Forward with autocast (float16 for ops)
 with autocast():
 output = self.model(x)
 loss = criterion(output, y)

 # Scale loss to prevent gradient underflow
 scaled_loss = loss * self.loss_scale

 # Backward
 self.optimizer.zero_grad()
 scaled_loss.backward()

 # Check for overflow (NaN or Inf gradients)
 if torch.isnan(scaled_loss) or torch.isinf(scaled_loss):
 print("Overflow detected! Reducing loss scale")
 self.loss_scale = max(self.loss_scale / 2, self.min_loss_scale)
 self.optimizer.zero_grad()
 return

 # Unscale gradients
 for param in self.model.parameters():
 if param.grad is not None:
 param.grad.data /= self.loss_scale

 # Clip gradients (after unscaling!)
 torch.nn.utils.clip_grad_norm_(self.model.parameters(), max_norm=1.0)

 # Optimizer step
 self.optimizer.step()

 # Increase loss scale if training stable
 self.loss_scale = min(self.loss_scale * 1.01, self.max_loss_scale)

-

Memory and Speed Analysis

Memory Breakdown

Model: Llama 7B, Batch size: 1, Sequence length: 4096

Without AMP (float32):
 - Model weights: 7B × 4 = 28GB
 - Optimizer states (Adam): 2 × 28 = 56GB
 - Activations: ~16GB
 - Gradients: 28GB
 - Total: ~128GB (doesn't fit on A100)

With AMP (float16 compute + float32 master):
 - Model weights (float32 master): 28GB
 - Model weights (float16 copy): 14GB
 - Optimizer states (float32): 56GB
 - Activations (float16): ~8GB
 - Gradients (float32): 28GB
 - Total: ~134GB (still doesn't fit!)

Wait... that's more memory! Why?

Issue: Need to maintain both float32 (master) and float16 (compute)
Solution: Use FSDP + AMP + Gradient checkpointing together

With FSDP (8 GPUs) + AMP + Checkpointing:
 - Per GPU weights: 28GB / 8 = 3.5GB
 - Per GPU optimizer: 56GB / 8 = 7GB
 - Per GPU activations: 2GB (checkpointed)
 - Per GPU gradients: 3.5GB / 8 = 0.4GB
 - Per GPU total: ~12GB (fits on A100!)

Techniques combined:
 - AMP: Float16 activations (50% memory)
 - Checkpointing: Recompute activations (90% reduction)
 - FSDP: Shard parameters (8x reduction)
 - Combined: 8 × 2 × 8 = 128x memory reduction!

Speed Improvement

Hardware: A100 GPU

Float32 Performance:
 - Dense matrix multiply: 312 TFLOPS
 - Actual throughput: 300 TFLOPS (realistic)

Float16 Performance:
 - Dense matrix multiply: 625 TFLOPS (2× float32)
 - With tensor cores: up to 1248 TFLOPS
 - Actual throughput: 600-800 TFLOPS (realistic)

Training throughput (iterations/sec):

Model: Llama 7B, Batch: 1, Seq: 4096
 - float32: 0.5 iter/sec
 - float16 AMP: 1.2 iter/sec
 - Speedup: 2.4x!

Model: Llama 70B, Batch: 1, Seq: 4096
 - float32: 0.05 iter/sec (very slow)
 - float16 AMP: 0.1 iter/sec
 - With FSDP: 0.3 iter/sec (6x total)
 - Speedup: 6x!

Wall-clock training time:
 - Without AMP: 48 hours to train 1 epoch
 - With AMP: 20 hours (2.4x faster)
 - Saves: 28 hours per epoch!

---

Practical Considerations

bfloat16 vs float16

Comparison:

 float16 bfloat16
───────────────────────────────────
Exponent bits 5 8
Mantissa bits 10 7
Memory 2B 2B
Range 10^-4-10^4 10^-38-10^38
Precision High Medium
Stability Underflow Stable
Hardware Wide A100+, newer

Recommendation:
 - GPUs with bfloat16: Use bfloat16 (more stable)
 - Older GPUs: Use float16 with careful loss scaling
 - Default modern: bfloat16 is becoming standard

Convergence Impact

Study: Does AMP hurt convergence?

Result: Negligible impact!

Accuracy difference (AMP vs full float32):
 - BERT: <0.1% accuracy loss
 - GPT-2: <0.2% accuracy loss
 - Vision models: <0.5% accuracy loss
 - Conclusion: AMP converges nearly identically

Convergence speed:
 - AMP may converge slightly faster (due to noise)
 - Or slightly slower (depends on loss scale tuning)
 - Typical: Within noise margins

---

When to Use AMP

Always Use AMP When

Training transformer models (default)
GPU has tensor cores (A100, H100, V100+)
Want 2-3x speedup for free
Need to reduce memory usage

Avoid AMP When

Doing numerical research (need full precision)
Model has custom float32-specific operations
Loss landscape is extremely unstable
 (rare, but possible with certain architectures)

---

Key Takeaways

2-3x speedup with virtually no accuracy loss 40-60% memory reduction through lower precision Automatic: PyTorch handles precision decisions Loss scaling prevents gradient underflow Industry standard: Should be default for training

-

Related Notes

• Distributed Training - Combine with AMP
• Gradient Checkpointing - Use together for maximum efficiency
• Lora - Alternative approach for efficient training
• Llm Inference Optimization - Different precision concerns in inference