Skip to content

Mixed-Precision Training: Choosing Between FP16, BF16 and FP8 ​

Note: This is a review article (no first-hand experimental data), based on public practice and hardware specs.

Experiment Objective ​

Mixed precision is not just "turn on AMP" — FP16's overflow, BF16's precision tradeoff, the new FP8 era (training + inference), and their interplay with loss scaling/stability/checkpoints are all engineering decisions. Goal: dissect a complete mixed-precision training framework — the AMP mechanism and loss scaling (why scaling is needed) → FP16/BF16/FP8 precision profiles and a selection matrix → precision strategy in training (when to drop precision, when not) → linkage with checkpoints and stability. After reading you can answer: which precision for my training, how to configure AMP, and whether FP8 is worth it.

Code Walkthrough ​

Part 1: AMP and loss scaling ​

python
# Standard PyTorch AMP (Automatic Mixed Precision)
scaler = torch.cuda.amp.GradScaler()          # gradient scaler

for x, y in loader:
    with torch.autocast(device_type="cuda", dtype=torch.bfloat16):
        loss = model(x, y)                    # forward: automatic mixed precision
    scaler.scale(loss).backward()             # scale before backward
    scaler.step(optimizer)                    # unscale before update
    scaler.update()                           # dynamically adjust scale factor

Key points:

  1. AMP picks precision per-op in forward/backward: ops auto-select (matmul/conv in low precision, normalization/softmax kept high); weights remain FP32 master copies — this is what "mixed" means.
  2. Loss scaling fixes FP16's gradient underflow: FP16's smallest positive is ~6e-5; tiny gradients vanish. Multiply loss by a scale (e.g. 2¹⁶), backprop with amplified gradients, unscale before update. BF16 needs no loss scaling (same exponent range as FP32, no underflow) — a key training advantage.
  3. scaler.update() adjusts dynamically: on inf/NaN, halve the scale and skip the update step (step handles internally), recover gradually — loss scaling is the core of FP16 training stability.

Part 2: Choosing the three precisions ​

python
# FP16: classic; needs loss scaling; small dynamic range
# BF16: training mainstream; no loss scaling; fewer mantissa bits (slightly lower precision)
# FP8: next-gen; training + inference; finer scaling and special handling

# Selection matrix (empirical):
# general LLM pretraining/finetuning → BF16 (simple, stable)
# need memory/speed savings, accept tiny loss → FP16 or FP8 (half memory)
# small models / low dynamic range → FP16 viable

Key points:

  • Training defaults to BF16: same dynamic range as FP32 (no overflow/underflow), no loss scaling; 7 mantissa bits (fewer than FP16's 10) slightly lower precision, acceptable for most training — "simple + stable" makes BF16 the default for large-model training.
  • FP16's strengths and costs: smaller memory (half FP32), some ops faster, but small dynamic range (large values overflow, small underflow) and mandatory loss scaling; fits "clearly need less memory/faster."
  • FP8 is the frontier: E4M3 (weights/activations) and E5M2 (gradients) with per-tensor scaling (finer than loss scaling) and hardware/framework maturity requirements — evaluate for new projects; production large-scale training stays on BF16 as the safe baseline.

Principles ​

Precision profile comparison ​

PrecisionExponent bitsMantissa bitsDynamic rangeTraining notes
FP32823largemaster copy/baseline
FP16510smallneeds loss scaling
BF1687large (like FP32)no scaling, slightly lower precision
FP8(E4M3)43very smallper-tensor scaling
FP8(E5M2)52smallerfor gradients

Core insight: precision selection is "dynamic range vs mantissa precision" — FP16 loses on range (overflow/underflow), BF16 on mantissa (slightly lower precision), FP8 tight on both (fine-grained scaling). Training (gradient-propagation-sensitive) and inference (forward-only) tolerate precision differently: inference can go lower (ch12's quantization article); training stays conservative (BF16 start).

Linkage with stability and checkpoints ​

  • Stability: on loss spikes under BF16, first investigate data/learning rate (ch10's spike article), not precision — unless reproducible experiments show FP32 helps.
  • Checkpoints: master weights saved in BF16/FP32 (ch11's mixed-precision-checkpoint precision handling); optimizer states (Adam moments) need higher precision — checkpoint precision and training precision are two decisions.

Results ​

(Estimated, not measured; typical LLM training)

PrecisionMemory (vs FP32)Training throughput (rel.)Stability
FP16 + scaling~0.5×1.5-2×needs scaling management
BF16~0.5×1.5-2×good (no scaling)
FP8~0.25×2-3×fine scaling; lower maturity

Engineering Takeaways ​

  1. Default BF16: large-model training starts on BF16, simple and stable; switch to FP16/FP8 only with clear gains (memory/speed).
  2. FP16 must pair with loss scaling: the GradScaler trio (scale/step/update) is standard — never run autocast without scaling.
  3. Checkpoint precision is a separate decision: master weights low precision, optimizer states high — don't conflate "training precision" with "checkpoint precision."
  4. FP8 validates small-scale first: assess hardware/framework support and per-tensor scaling; go to production after clear memory gains are verified.

Next article preview: per-card precision and efficiency done; next, an industry landing form — "Digital Twins & Industrial Simulation: The Loop of Supercomputing, AI Surrogate Models, and Real-Time Data."

AI-HPC Expert · Contact: openaihpc@gmail.com