Mixed Precision Training

Mixed precision training keeps the model accurate while halving (or quartering) memory and doubling (or quadrupling) throughput. The key is knowing which tensors need which precision — get it wrong and training diverges.

---

1. Why Mixed Precision?

Precision	Bytes	Tensor Core TFLOPS (H100)	Memory BW savings
FP32	4	67	Baseline
TF32	4 (stored) / 19-bit (compute)	495	Same memory, faster compute
BF16	2	989	2×
FP16	2	989	2×
FP8 (E4M3)	1	1,979	4×

BF16 gives 2× memory savings AND 15× compute speedup over FP32. FP8 doubles that again on Hopper+.

---

2. Number Format Comparison

BF16 vs FP16

Property	FP32	BF16	FP16
Exponent bits	8	8	5
Mantissa bits	23	7	10
Dynamic range	$\pm 3.4 \times 10^{38}$	$\pm 3.4 \times 10^{38}$	$\pm 6.5 \times 10^{4}$
Precision	~7 decimal digits	~2.4 digits	~3.3 digits
Underflow threshold	$1.2 \times 10^{-38}$	$1.2 \times 10^{-38}$	$6.1 \times 10^{-5}$

BF16 wins for training because it has the same dynamic range as FP32 — gradients and activations rarely overflow or underflow. FP16’s limited range ($\pm 65504$) causes gradient underflow without loss scaling.

FP8 Formats

Hopper introduces two FP8 formats:

Format	Exponent	Mantissa	Range	Used for
E4M3	4	3	$\pm 448$	Forward pass (weights, activations)
E5M2	5	2	$\pm 57344$	Backward pass (gradients — need more range)

E4M3 has higher precision (more mantissa bits) — better for weights and activations where values are bounded. E5M2 has higher range (more exponent bits) — better for gradients which have wider dynamic range.

---

3. The Mixed Precision Recipe

BF16 Mixed Precision (Standard)

Working precision:  BF16 (forward, backward, communication)
Master weights:     FP32 (optimizer updates)
Optimizer states:   FP32 (m, v in Adam)
Gradient accumulation: BF16 or FP32

Memory: 16Ψ (2 BF16 params + 2 BF16 grads + 12 FP32 optimizer)

The key operation is the optimizer step:

# Simplified mixed-precision training loop
for data in dataloader:
    # Forward in BF16
    with torch.autocast(device_type='cuda', dtype=torch.bfloat16):
        loss = model(data)
    
    # Backward in BF16 (gradients are BF16)
    loss.backward()
    
    # Optimizer step: BF16 grads → FP32 update → BF16 params
    optimizer.step()  # internally:
    # 1. Cast BF16 gradients to FP32
    # 2. Update FP32 master weights using FP32 Adam states
    # 3. Cast FP32 master weights back to BF16 for next forward
    
    optimizer.zero_grad()

FP16 Mixed Precision (Legacy)

FP16 requires loss scaling because small gradients underflow to zero:

scaler = torch.cuda.amp.GradScaler()

for data in dataloader:
    with torch.autocast(device_type='cuda', dtype=torch.float16):
        loss = model(data)
    
    # Scale loss to prevent gradient underflow
    scaler.scale(loss).backward()
    
    # Unscale gradients, check for inf/nan, then step
    scaler.step(optimizer)
    scaler.update()  # adjust scale factor

Loss scaling multiplies the loss by a large factor (e.g., 1024) before backward, so small gradients are amplified into FP16’s representable range. After backward, gradients are divided by the same factor.

BF16 doesn’t need loss scaling because its dynamic range matches FP32. This is why BF16 is universally preferred on Ampere+ GPUs.

FP8 with Transformer Engine

NVIDIA’s Transformer Engine (TE) handles FP8 automatically:

import transformer_engine.pytorch as te

# Replace nn.Linear with TE's FP8-aware version
model = te.Linear(hidden_size, ffn_size)

# TE handles:
# 1. Per-tensor scaling (calibrate FP8 range for each tensor)
# 2. Delayed scaling (use previous step's statistics)
# 3. E4M3 for forward, E5M2 for backward
# 4. FP32 master weights for optimizer

Per-tensor scaling: Each weight matrix and activation tensor gets its own scaling factor, computed from the max absolute value. This maximizes the use of FP8’s limited range.

Delayed scaling: Instead of computing scales from the current tensor (which requires an extra pass), use scales from the previous training step. Works because tensor statistics change slowly between steps.

---

4. Which Tensors at Which Precision

Tensor	BF16 Training	FP8 Training	Why
Forward activations	BF16	FP8 (E4M3)	Compute-bound, benefits from Tensor Cores
Weights (working copy)	BF16	FP8 (E4M3)	Stored alongside activations
Backward gradients	BF16	FP8 (E5M2)	Wider range needed
Master weights	FP32	FP32	Accumulated small updates
Adam m, v	FP32	FP32	Statistical accuracy
LayerNorm	FP32	FP32/BF16	Numerically sensitive
Softmax	FP32	FP32	Numerically sensitive
Loss computation	FP32	FP32	Precision matters for convergence
AllReduce communication	BF16	BF16	Bandwidth-sensitive

Rule: Compute-heavy operations (GEMM) use low precision. Numerically-sensitive operations (norms, softmax, loss) use high precision. Optimizer always uses FP32.

---

5. Memory Impact

Configuration	Bytes/param	7B total	Speedup vs FP32
FP32 everything	20	140 GB	1×
BF16 mixed (standard)	16	112 GB	~2×
FP8 mixed (Hopper+)	14	98 GB	~3×
8-bit Adam (bitsandbytes)	10	70 GB	~2×

FP8 saves less than you’d expect because the optimizer states (12Ψ) are still FP32. The savings come from parameters and gradients: 2Ψ + 2Ψ (BF16) → 1Ψ + 1Ψ (FP8), saving 2Ψ bytes.

---

6. Common Pitfalls

Pitfall	Symptom	Fix
Training with FP16 without loss scaling	Loss goes to 0, gradients vanish	Use GradScaler() or switch to BF16
LayerNorm in low precision	Training instability, loss spikes	Force FP32 for LayerNorm
Gradient accumulation precision	Accumulated gradients lose precision	Accumulate in FP32, not BF16
FP8 without proper scaling	NaN loss immediately	Use Transformer Engine’s delayed scaling
Mixing autocast scopes	Some layers compute in wrong precision	Ensure entire forward pass is in autocast

For how mixed precision interacts with distributed training, see FSDP.

---

Self-Check

• Compare BF16 vs FP16: dynamic range, precision, and why BF16 is preferred
• Explain the FP8 E4M3 vs E5M2 format choice for forward vs backward
• Write a mixed-precision training loop with torch.autocast
• Explain loss scaling: why FP16 needs it and BF16 doesn’t
• Identify which tensors must remain in FP32 and why
• Calculate memory savings from BF16 vs FP32 mixed precision

---

References

• Mixed Precision Training (Micikevicius et al., 2017)
• FP8 Formats for Deep Learning (Micikevicius et al., 2022)
• NVIDIA Transformer Engine