Mixed Precision Training (AMP): Speed and Memory with Reduced Precision

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(): Automatically casts operations to float16/float32 as needed
# GradScaler: Scales loss to prevent gradient underflow in float16

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