Mixed precision training

What Problem Does This Solve?

Training large neural networks faces three bottlenecks:

Memory: Modern models (GPT, Vision Transformers) have billions of parameters. Each FP32 parameter takes 4 bytes. A175B parameter model needs 700GB just to store weights!
Compute: Matrix multiplications dominate training time. GPUs have specialized hardware (Tensor Cores) that runs FP16 ops 2-8× faster than FP32.
Communication: In distributed training, gradients must sync across GPUs. Halving precision halves bandwidth requirements.

But naively switching to FP16 causes gradient underflow: tiny gradients round to zero, and training diverges.

The key insight: gradients span a huge dynamic range. Early layers might have gradients around $10^{-6}$ , while later layers have $10^{-2}$ . FP16's limited range can't represent both.

How Mixed Precision Training Works

The algorithm has three core components:

Why: Weight updates are cumulative. If learning rate is $10^{-4}$ and gradient is $10^{-3}$ , the update is $10^{-7}$ —smaller than FP16's precision ( $10^{-3}$ ). The update would round to zero, and weights never change.

How: Maintain two copies of weights:

FP16 copy forward/backward passes (fast compute)
FP32 master copy for accumulating updates (preserve precision)

Derivation from First Principles:

Consider a single parameter $w$ updated by gradient descent: $w_{t+1} = w_t - \eta \cdot g_t$

In FP16, if $\eta g_t < \epsilon_{FP16} \approx 10^{-3}$ , the update becomes: $w_{t+1}^{FP16} = \text{round}_{FP16}(w_t - \eta g_t) = w_t$ (the small update rounds to zero)

After $N$ steps with zero updates, we've lost $N \eta g_t$ of accumulated change—training stalls.

Solution: Store master weights in FP32: $w_{t+1}^{FP32} = w_t^{FP32} - \eta \cdot g_t^{FP32}$ Even if individual updates are tiny, they accumulate in the higher-precision master. Then cast to FP16 for the next forward pass: $w_{t+1}^{FP16} = \text{cast}(w_{t+1}^{FP32})$

2. Loss Scaling

Why: Gradients in early layers often fall below FP16's minimum normal value ( $6 \times 10^{-8}$ ). These underflow to zero, creating "dead neurons" that never learn.

What: Multiply the loss by a large scalar $S$ (typically $512$ to $6536$ ) before backprop. This shifts gradients into FP16's representable range. After computing gradients, divide by $S$ to restore true values before updating master weights.

Derivation:

The chain rule for gradients gives: $\frac{\partial L}{\partial w} = \frac{\partial L}{\partial a} \cdot \frac{\partial a}{\partial w}$

If $\frac{\partial L}{\partial a} \approx 10^{-7}$ (common in deep networks), FP16 cannot represent it—underflows to zero.

Loss scaling trick: Define a scaled loss $L' = S \cdot L$ where $S \gg 1$ . Then: $\frac{\partial L'}{\partial w} = S \cdot \frac{\partial L}{\partial w}$

Now if $\frac{\partial L}{\partial w} = 10^{-7}$ and $S = 1024$ , then $\frac{\partial L'}{\partial w} = 10^{-4}$ , which FP16 can represent!

After backprop, unscale: $\frac{\partial L}{\partial w} = \frac{1}{S} \cdot \frac{\partial L'}{\partial w}$

Why This Step?: Scaling happens at the loss (top of computational graph) so it automatically propagates through the entire backward pass. This is more efficient than scaling each layer's gradients individually.

Formula Loss Scaling Algorithm

# Forward pass (FP16)
loss = model(x, y)  # Computed in FP16
 
# Scale loss BEFORE backward
scaled_loss = loss * scale_factor  # scale_factor ~ 1024
 
# Backward pass (gradients are now scaled, stay in FP16 range)
scaled_gradients = backward(scaled_loss)
 
# Unscale gradients before optimizer step
true_gradients = scaled_gradients / scale_factor
 
# Update FP32 master weights
optimizer.step(true_gradients)

Dynamic Loss Scaling: Start with large $S$ (e.g., $6536$ ). If gradients overflow (produce inf or nan), halve $S$ and retry. If training is stable for $N$ steps (e.g., $2000$ ), double $S$ . This adapts to each model's gradient distribution.

3. FP16 Arithmetic for Most Operations

Why: Matrix multiplications (GEMM) dominate training time—often 90%+ of compute. Modern GPUs have Tensor Cores that execute FP16 ops at 2-16× the throughput of FP32.

What: Store activations, gradients, and weights in FP16 during forward/backward passes. Only the optimizer state and master weights stay in FP32.

How:

Cast weights from FP32 master to FP16 at the start of each forward pass
Compute activations in FP16: $a^{FP16} = \sigma(W^{FP16} x{FP16} + b^{FP16})$
Compute loss and scale it
Backward pass in FP16 (gradients accumulate in FP16)
Cast gradients back to FP32 for optimizer

Worked example Example 1: Single Weight Update

Setup: A weight starts at $w_0^{FP32} = 0.5$ , learning rate $\eta = 0.001$ , gradient $g = 0.003$ . Step 1 (Forward): Cast to FP16 $w_0^{FP16} = \text{cast}(0.5) = 0.5$ (No precision loss since $0.5$ is exactly representable in FP16) Step 2 (Backward): Compute gradient in FP16, $g^{FP16} = 0.003$ Step 3 (Optimizer): Update in FP32

g^{FP32} &= \text{cast}(0.003) = 0.003 \\ w_1^{FP32} &= 0.5 - 0.0001 \times 0.003 = 0.4997 \end{aligned}$$ **Why this step?**: The update magnitude is $3 \times 10^{-7}$. FP16 precision is $\sim 10^{-3}$, so: $$w_1^{FP16} = \text{round}_{FP16}(0.5 - 0.0003) = 0.5$$ The update would be lost! FP32 preserves it: $0.4997 \neq 0.5$. **Step 4**: Next forward pass uses $w_1^{FP16} = \text{cast}(0.4997) \approx 0.4997$.

Worked example Example 2: Loss Scaling Prevents Underflow

Setup: A gradient is $g = 5 \times 10^{-8}$ (below FP16's min normal value of $6 \times 10^{-8}$ ). Loss scale $S = 1024$ .

Without scaling: $g^{FP16} = \text{cast}(5 \times 10^{-8}) = 0$ (Underflows to zero)

With scaling:

\text{Loss} &\rightarrow S \cdot \text{Loss} = 1024 \cdot \text{Loss} \\ g_{\text{scaled}}^{FP16} &= 1024 \times 5 \times 10^{-8} = 5.12 \times 10^{-5} \\ &\text{(representable in FP16!)} \\ g_{\text{true}}^{FP32} &= \frac{g_{\text{scaled}}^{FP16}}{1024} = 5 \times 10^{-8} \end{aligned}$$ **Why this step?**: Scaling shifts the gradient into FP16's representable range during the backward pass. We only need the true value at the optimizer, which runs in FP32.

Worked example Example 3: Memory Savings in a Real Model

Setup: ResNet-50 has ~25M parameters. Training with batch size 256.

FP32 only:

Weights: $25M \times 4 \text{ bytes} = 100 \text{ MB}$
Gradients: $100 \text{ MB}$
Activations (batch 256): ~ $500 \text{ MB}$ (depends on intermediate layer sizes)
Optimizer state (Adam: 2× weights): $200 \text{ MB}$
Total: ~ $900 \text{ MB}$ per GPU

Mixed precision:

FP16 weights (for compute): $50 \text{ MB}$
FP32 master weights: $100 \text{ MB}$
FP16 gradients: $50 \text{ MB}$
FP16 activations: ~ $250 \text{ MB}$
FP32 optimizer state: $200 \text{ MB}$
Total: ~ $650 \text{ MB}$ per GPU

Why this step?: We saved $\frac{900 - 650}{900} \approx 28\%$ memory. For larger models, this difference determines whether a model fits on a GPU at all. The activations savings are the biggest win—they dominate memory in deep networks.

Common Pitfalls and Fixes

Common mistake Mistake 1: Forgetting to Unscale Gradients Before Clipping

Why it feels right: "I'll clip gradients at norm 1.0 to prevent exploding gradients." Why it's wrong: If loss scale is $S = 1024$ , scaled gradients have norms $1024×$ larger. Clipping at 1.0 essentially clips unscaled gradients at $1/1024 \approx 0.001$ —way too aggressive! Training barely progresses.

Fix: Unscale before clipping:

scaled_grad = backward(scaled_loss)
true_grad = scaled_grad / scale_factor
clipped_grad = clip_by_norm(true_grad, max_norm=1.0)
optimizer.step(clipped_grad)

Steel-man: The mistake happens because gradient clipping is often added as a separate line, and it's easy to forget the unscaling order. The fix: make unscaling part of your gradient preparation pipeline, always before clipping.

Common mistake Mistake 2: Using Static Loss Scale for All Models

Why it feels right: "Loss scale of 1024 worked for ResNet, so I'll use it everywhere."

Why it's wrong: Different architectures have different gradient distributions. Vision Transformers often have larger gradients than CNs. A fixed scale might be too small (underflow) or too large (overflow).

Evidence: Training aViT with scale $1024$ produces nan after 100 steps. With dynamic scaling (starting at $6536$ , adapted down to $512$ ), training is stable.

Fix: Use dynamic loss scaling or tune per-model. Monitor gradient histograms—if you see many zeros (underflow) or inf (overflow), adjust the scale.

Steel-man: Static scaling seems appealing for reproducibility, but gradient magnitudes vary 1000× across architectures. Dynamic scaling is the robust default.

Common mistake Mistake 3: Storing Batch Normalization Stats in FP16

Why it feels right: "All parameters should be FP16 for consistency."

Why it's wrong: Batch norm running mean/variance are updated with exponential moving average (EMA): $\mu_{\text{running}} = 0.9 \cdot \mu_{\text{running}} + 0.1 \cdot \mu_{\text{batch}}$ Small updates ( $0.1 \times \text{small batch mean}$ ) underflow in FP16. After many steps, running stats become inaccurate, and inference fails.

Fix: Keep running statistics in FP32, even when everything else is FP16. Most frameworks do this automatically.

Steel-man: This is an edge case that's easy to miss when implementing mixed precision from scratch. The lesson: any state updated with small increments needs FP32.

Implementation Pseudocode

# Initialization
model_fp16 = model.half()  # FP16 copy for compute
master_weights = model.parameters()  # FP32 master copy
optimizer = Adam(master_weights)
scale = 2**16  # Initial loss scale (6536)
 
for batch in dataloader:
    # Forward pass in FP16
    with autocast(dtype=torch.float16):
        output = model_fp16(batch.x)
        loss = criterion(output, batch.y)
    
    # Scale loss to prevent gradient underflow
    scaled_loss = loss * scale
    # Backward pass (gradients computed in FP16, but scaled)
    scaled_loss.backward()
    
    # Unscale gradients
    for param in master_weights:
        if param.grad is not None:
            param.grad = param.grad / scale
    # Clip gradients (on unscaled values!)
    torch.nn.utils.clip_grad_norm_(master_weights, max_norm=1.0)
    
    # Optimizer step on FP32 master weights
    optimizer.step()
    optimizer.zero_grad()
    # Copy updated master weights back to FP16 model
    for param_fp16, param_fp32 in zip(model_fp16.parameters(), master_weights):
        param_fp16.data.copy_(param_fp32.data)
    # Dynamic loss scaling: check for overflow
    if any(torch.isnan(p.grad).any() for p in master_weights if p.grad is not None):
        scale = scale / 2  # Halve scale on overflow
    elif step % 2000 == 0:
        scale = min(scale * 2, 2**16)  # Double scale if stable

Theoretical Guarantees

Convergence: Mixed precision training provably converges to the same optimum as FP32 training under mild conditions:

Assumption: Gradients $g$ satisfy $\|g\|_{\infty} < K$ (bounded) and loss scale $S$ ensures \cdot g_{\min} > \epsilon_{FP16}$ (no underflow).
Result: If learning rate $\eta < \frac{1}{L}$ (standard SGD condition), expected distance to optimum decreases at the same rate as FP32: $\mathbb{E}[\|w_t - w^*\|^2] \leq \frac{C}{t}$ where $C$ depends on $L$ (loss smoothness) but not precision.

Proof sketch: Master weights accumulate updates exactly (no precision loss). FP16 forward/backward introduce rounding errors $\delta$ at each operation. But modern GPUs use round-to-nearest mode, so $\mathbb{E}[\delta] = 0$ (unbiased). The error variance is $O(\epsilon_{FP16}^2)$ , which vanishes as $t \to \infty$ due to averaging. Loss scaling ensures $\delta$ is measured relative to scaled gradients, keeping $\delta/g$ small.

When to Use Mixed Precision

Use when:

Training models with >10M parameters (memory savings matter)
Using GPUs with Tensor Cores (V100, A100, RTX 20/30 series)
Distributed training (gradient communication is a bottleneck)
Inference is memory-bound (can deploy FP16 models directly)

Don't use when:

Debugging numerical issues (hard to tell if bug is yours or precision-related)
Working with datasets where values span >5 orders of magnitude (e.g., scientific computing with huge ranges)
Using custom CUDA kernels not tested with FP16

Real-world impact: GPT-3 training used mixed precision and would have been infeasible without it (175B parameters → ~700GB FP32, but fits in ~350GB with mixed precision and activation checkpointing).

Recall Explain Like I'm 12

Pretend you're playing a video game, and you have a pigy bank to store your gold coins. Each coin can be stored in two ways:

Small pouches (FP16): Hold coins faster, take less space, but can only count coins from 1 to 100.
Big chests (FP32): Slower, heavier, but can count from 0.001 to 1,000,000.

Most of the time you're running around collecting and spending coins, so you use small pouches—it's quick! But every night, you pour all your pouches into a big chest to add up your total wealth. That way, even tiny coins (like 0.01 gold) don't get lost.

Now, sometimes you find really tiny treasure (0.001 gold). If you put it straight in a small pouch, it rounds to zero and disappears! So you use a trick: multiply all treasure by 1000 when you collect it. Now0.0001 becomes 0.1, which the pouch can hold. Later, when you pour into the big chest, you divide by 1000 to get the real value back.

Mixed precision training is exactly this: use fast, small storage (FP16) for most work, but keep a careful ledger (FP32 master weights) so tiny updates don't vanish. Scale up gradients so they don't disappear. This makes training huge AI models 2× faster and use 30% less memory—like beating the video game boss with a speed potion!

Connections

3.3.08-Gradient-accumulation – Mixed precision is often combined with gradient accumulation to fit large batches
3.3.05-Learning-rate-schedules – Loss scaling interacts with LR: both affect gradient magnitude
4.2.03-Tensor-cores – Hardware motivation for mixed precision
3.3.10-Gradient-checkpointing – Another memory optimization, often used together
2.4.06-Batch-normalization – Requires special handling (FP32 running stats)
5.1.02-Model-parallelism – Mixed precision reduces communication in distributed training
1.3.07-Numerical-stability – Fundamental tradeoffs between speed and precision

#flashcards/ai-ml

What are the three core components of mixed precision training? :: (1) FP32 master weights for accurate updates, (2) Loss scaling to prevent gradient underflow, (3) FP16 arithmetic for forward/backward passes to speed up compute.

Why can't we just train entirely in FP16?

FP16's limited range (

10^{-8}

6504

) causes (1) gradient underflow (tiny gradients → zero) and (2) weight update underflow (small learning rates × gradients → no change). Master weights in FP32 prevent update underflow.

Derive why loss scaling prevents gradient underflow :: If gradient $g = 10^{-7} < 6 \times 10^{-8}$ (FP16 min), it underflows to 0. Scale loss by $S = 1024$ : backward gives $g' = Sg = 10^{-4}$ (representable). Unscale: $g = g'/S$ recovers true value in FP32 before optimizer step.

What is the memory formula for mixed precision training?

Memory

\approx

(FP16 weights + FP32 master weights + FP16 gradients + FP16 activations + FP32 optimizer state). Typically 25-35% savings vs pure FP32, mainly from activations being ~50% of memory.

When should you unscale gradients—before or after gradient clipping?

Before clipping. Scaled gradients have norms

S×

larger. If you clip scaled grads at norm 1.0, you're actually clipping true grads at

1/S \approx 0.001

(with

S=1024

), which is way too aggressive.

Why must batch normalization running statistics stay in FP32?

Running mean/var updated with EMA:

\mu_{\text{new}} = 0.9\mu_{\text{old}} + 0.1\mu_{\text{batch}}

. The

0.1×

small value update is often

<10^{-3}

(FP16 precision limit), causing underflow. Accumulated error breaks inference.

What is dynamic loss scaling and why is it needed?

Start with large scale (e.g.,

2^{16}

). If gradients overflow (produce inf/nan), halve the scale. If training is stable for

N

steps, double the scale. Needed because different models have gradient distributions spanning 1000× in magnitude—no single static scale works universally.

Compare FP16 vs BF16 :: FP16: 5exp bits, 10 mantissa bits. Range $\sim 10^{-8}$ to $6504$ , precision $\sim 10^{-3}$ . BF16: 8 exp bits, 7 mantissa bits. Same range as FP32 ( $\sim 10^{-38}$ to $10^{38}$ ), less precision. BF16 has fewer overflow issues but less accuracy; preferred on TPUs and newer GPUs.

What is the speedup from mixed precision training on modern GPUs?

Typically 2-3× end-to-end training speedup on V100/A100 with Tensor Cores. Raw matrix multiply can be 8-16× faster in FP16, but overhead (data loading, optimizer CPU work) limits total speedup. Memory savings (25-35%) allow larger batch sizes, giving further speedup.

Concept Map

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Hinglish (regional understanding)

Intuition Hinglish mein samjho

Chalo simple tarike se samajhte hain. Jab hum bade neural networks train karte hain, tab do badi problems aati hain — memory kam padti hai (billions of parameters store karne hote hain) aur compute slow hoti hai. Iska smart solution hai mixed precision training. Iska core idea bilkul aisa hai jaise checkbook balance karna — zyada tar transactions ke liye tumhe dollars aur cents kaafi hain, par kuch critical calculations ke liye extra precision chahiye. Waise hi, hum jyada computations FP16 (lower precision) me karte hain taaki fast aur memory-efficient ho, par jahan numerical accuracy zaroori hai wahan FP32 (higher precision) use karte hain. Yeh cheating nahi, balki smart engineering tradeoff hai.

Ab problem yeh hai ki agar tum sidha sab kuch FP16 me kar do, toh gradient underflow ho jaata hai — matlab bahut chote gradients (jaise $10^{-7}$ ) zero me round ho jaate hain, aur weights kabhi update hi nahi hote, training stall ya diverge ho jaati hai. Isko solve karne ke liye do main tricks hain. Pehla, FP32 master weights rakho — forward/backward pass toh FP16 me fast chalta hai, par actual weight updates ek FP32 master copy me accumulate hote hain, taaki chote updates zero na ban jaayein. Dusra, loss scaling — loss ko ek bade number $S$ se multiply karke gradients ko FP16 ki representable range me shift kar dete hain, aur baad me $S$ se divide karke original value wapas le lete hain.

Yeh matter kyun karta hai? Aaj ke real-world models — GPT, Vision Transformers — inko bina mixed precision ke train karna practically impossible ya bahut mehnga hota. Tumhe 2-8x speedup milta hai GPU ke Tensor Cores ki wajah se, memory aadhi ho jaati hai, aur distributed training me communication bandwidth bhi kam lagti hai. Toh agar tum future me ML engineer ya researcher banna chahte ho, yeh concept samajhna zaroori hai — kyunki har modern deep learning framework (PyTorch, TensorFlow) me yeh feature by default use hota hai. Yeh sirf theory nahi, ekdum practical skill hai jo tum daily kaam me apply karoge.

What Problem Does This Solve?

How Mixed Precision Training Works

1. FP32 Master Weights

2. Loss Scaling

3. FP16 Arithmetic for Most Operations

Common Pitfalls and Fixes

Implementation Pseudocode

Theoretical Guarantees

When to Use Mixed Precision

Connections

Concept Map

Hinglish (regional understanding)

Hinglish (regional understanding)

Go deeper — visual, from zero

Test yourself — Deep Learning Frameworks

Connections