Mixed Precision Training

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. Floating-Point Formats
5. FP16 Training with Loss Scaling
6. BF16: The Modern Standard
7. The Master Weight Copy
8. FP8 Training (Hopper)
9. Memory and Compute Analysis
10. Common Pitfalls
11. Summary
12. Exercises

Learning Objectives

1. Compare FP32, FP16, BF16, and FP8 in terms of range and precision.
2. Explain why FP16 requires loss scaling and BF16 doesn't.
3. Derive the memory savings from mixed-precision training.
4. Explain the master weight copy and why it's necessary.
5. Describe FP8 training and its per-tensor scaling requirements.

Notation

• FP32: 1 sign + 8 exponent + 23 mantissa bits
• FP16: 1 sign + 5 exponent + 10 mantissa bits
• BF16: 1 sign + 8 exponent + 7 mantissa bits
• FP8 (E4M3): 1 sign + 4 exponent + 3 mantissa bits
• FP8 (E5M2): 1 sign + 5 exponent + 2 mantissa bits

Core Intuition

Neural network training doesn't need full FP32 precision — most gradients and activations use a tiny fraction of the representable range. By using lower precision (FP16/BF16), we halve memory usage and double compute throughput (tensor cores operate at 2x speed in half precision). The challenge: avoiding numerical issues from reduced precision.

Floating-Point Formats

FP32 (standard):

• Range: $\pm 3.4 \times 10^{38}$. Precision: $\sim 7$ decimal digits.
• Used for: master weights, optimizer states.

FP16 (half precision):

• Range: $\pm 65504$. Precision: $\sim 3.3$ decimal digits.
• Problem: Small gradients ($< 2^{-24} \approx 6 \times 10^{-8}$) underflow to zero.

BF16 (brain float):

• Range: $\pm 3.4 \times 10^{38}$ (same as FP32!). Precision: $\sim 2.4$ decimal digits.
• Advantage: Same range as FP32 — no underflow issues. Less precision is acceptable for training.

FP8 (E4M3 for forward, E5M2 for backward):

• E4M3: Range $\pm 448$, precision $\sim 1.5$ digits. For weights and activations.
• E5M2: Range $\pm 57344$, precision $\sim 1$ digit. For gradients (needs more range).

FP16 Training with Loss Scaling

Problem: Gradients for early layers can be $O(10^{-8})$ — below FP16's minimum representable value ($\sim 6 \times 10^{-8}$). These underflow to zero → no learning.

Solution: Loss scaling. Multiply loss by a scale factor $S$:

Gradients are $S$ times larger — within FP16 range. Before the optimizer step, divide by $S$:

Dynamic loss scaling: Start with large $S$; halve it if gradient overflow (inf/nan) detected; double every $N$ steps without overflow.

BF16: The Modern Standard

Why BF16 dominates:

• Same dynamic range as FP32 → no loss scaling needed.
• Simpler training pipeline (no checking for overflows).
• Supported on A100, H100, and newer GPUs at 2x FP32 throughput.

Precision cost: 7 mantissa bits vs 10 for FP16 → slightly more rounding error per operation. Empirically: negligible impact on training quality for models above 1B parameters.

All modern LLM training uses BF16 (LLaMA, GPT-4, Claude, Gemini).

The Master Weight Copy

Problem: Optimizer updates are small: $\Delta\theta \approx \eta \cdot \text{grad} \approx 10^{-5}$. In FP16: $1.0 + 10^{-5} = 1.0$ (lost to rounding). Weights never update.

Solution: Maintain a FP32 master copy of weights:

1. Forward/backward in FP16/BF16 (fast).
2. Accumulate gradients in FP32 (or FP16 reduced to FP32).
3. Update FP32 master weights.
4. Cast updated weights back to FP16/BF16 for next step.

Memory: FP32 master ($4\Psi$) + FP16 working copy ($2\Psi$) = $6\Psi$ for weights (vs $4\Psi$ pure FP32). But activations and gradients are half size → net savings.

FP8 Training (Hopper)

H100 FP8 tensor cores: 2x throughput vs BF16 (1.6 PFLOPS vs 0.8 PFLOPS).

Per-tensor scaling: FP8 has very limited range. Each tensor gets its own scale factor:

Delayed scaling: Use the max from the previous iteration to set the current scale (avoids an extra pass to compute max).

Recipe: Forward activations in E4M3, backward gradients in E5M2, master weights in FP32.

Result: 2x faster training with less than 0.5% quality loss for models larger than 7B.

Memory and Compute Analysis

Memory comparison (model portion, per GPU):

• FP32 training: $4\Psi$ (params) + $4\Psi$ (grads) + $12\Psi$ (Adam) = $20\Psi$ bytes.
• Mixed precision (BF16): $2\Psi$ (BF16 params) + $2\Psi$ (BF16 grads) + $4\Psi$ (FP32 master) + $8\Psi$ (FP32 Adam) = $16\Psi$ bytes.
• Activations: halved (BF16 vs FP32).

Compute speedup: BF16 tensor cores = 2x FP32. FP8 = 4x FP32.

Common Pitfalls

Pitfall 1. Using FP16 without loss scaling for deep networks. Gradient underflow is invisible (gradients silently become zero) — the model appears to train but actually stalls.

Pitfall 2. Accumulating gradients in FP16. Gradient accumulation (summing micro-batch gradients) in FP16 causes severe precision loss. Always accumulate in FP32.

Pitfall 3. Using FP8 without per-tensor scaling. A single outlier value in a tensor wastes the entire FP8 range on that one value, destroying precision for everything else.

Summary

• BF16 is the standard: FP32 range, 2x throughput, no loss scaling needed.
• FP16 requires dynamic loss scaling to prevent gradient underflow.
• Master weight copy (FP32) prevents weight stagnation from rounding.
• FP8 on H100: 2x faster than BF16 with per-tensor scaling.
• Mixed precision halves activation memory and doubles compute throughput.

Exercises

Exercise 1. Compute the smallest positive value representable in FP16 and BF16. Explain why FP16 has underflow problems.

Exercise 2. For a gradient of magnitude $10^{-7}$: what loss scale $S$ is needed to bring it into FP16's representable range?

Exercise 3. Compute the total memory savings from mixed precision (BF16 forward/backward, FP32 optimizer) vs pure FP32 for a 13B model.

Exercise 4. Derive the per-tensor scaling factor for FP8 E4M3 given a tensor with max absolute value 3.7.

Exercise 5. Explain why BF16 (7 mantissa bits) is acceptable for training despite having less precision than FP16 (10 mantissa bits).