Mixed Precision Training & Checkpointing: Stable and Recoverable

Abstract: Mixed precision decides how fast training can run; checkpointing decides whether training dares to run at all. The former is a throughput problem, the latter a reliability problem — on a 1000-GPU cluster, a single failure costs hours of recovery. Get both wrong and the biggest cluster is wasted.

1. Mixed Precision: Why Not Just Low Precision Everywhere

1.1 Precision Selection Matrix

Precision	Exponent bits	Mantissa bits	Training use	Key limitation
FP32	8	23	Baseline / stability	Doubles memory, slower
TF32	8	10	Compromise (Ampere+)	Slightly lower precision; fine in most cases
FP16	5	10	Training (needs Loss Scaling)	Small dynamic range; overflow risk
BF16	8	7	Training workhorse	Fewer mantissa bits; relative precision sufficient at scale
FP8	5/4	2/3	Training/inference acceleration	Needs extra stability techniques

Core principle: keep weights and gradients in high precision (FP32), do the compute in low precision (BF16/FP16) — that combination is mixed precision.

1.2 The Two Main Approaches

FP16 (AMP, needs Loss Scaling):
  loss × scale → FP16 forward → FP16 gradients → unscaled → FP32 update
  Problem: gradient underflow/overflow risk; dynamic scale; complex

BF16 (modern default):
  BF16 forward/backward → FP32 master weight update
  Benefit: no scaling; same dynamic range as FP32; naturally stable

• FP16 has a small dynamic range (5 exponent bits); tiny gradients underflow to zero — hence Loss Scaling "amplifies" gradients before updates.
• BF16 shares FP32's 8 exponent bits, so its dynamic range matches FP32 — no under/overflow, no Loss Scaling, stable training. The cost is fewer mantissa bits, proven sufficient for LLM training.
• Engineering conclusion: default to BF16 for new training projects; use FP16 only for legacy-framework compatibility or specific quantization paths; treat FP8 as "fast but not yet stable" — be cautious in production.

2. Checkpoints: The Heart of Training

2.1 What a Checkpoint Must Contain

checkpoint = {
    "model": model.state_dict(),          # weights (FP32 master + BF16 copies)
    "optimizer": optimizer.state_dict(),  # Adam moments/variance (the bulk!)
    "lr_scheduler": scheduler.state_dict(),
    "rng": torch.random.get_rng_state(),  # data order, dropout randomness
    "data_loader": dataloader.state_dict(), # epoch/sample position
    "epoch": epoch, "step": step,
    "config": config,                       # hyperparameters, validated on resume
}

• For 70B BF16 full-parameter training, a checkpoint ≈ weights 140G + optimizer 420G ≈ 560 GB. On a 1000-GPU cluster, save speed directly determines recovery time (RTO).
• Optimizer states are the largest part of a checkpoint, far exceeding the weights themselves.

2.2 Save Strategy and Async Saving

# Async save: push serialization + disk I/O to a background thread
import threading
def save_async(state, path):
    def _save():
        torch.save(state, path + ".tmp")
        os.replace(path + ".tmp", path)   # atomic replace, prevents half-written files
    t = threading.Thread(target=_save, daemon=True)
    t.start()

# Frequency: small intervals + rolling retention (e.g. keep last 3 every 500 steps)
# Milestones permanently saved (pretrain done, each SFT stage)

Three production essentials:

1. Atomic writes: write .tmp first, then os.replace — prevents half-written checkpoints after power loss.
2. Rolling retention: keep the last 2–3 on disk plus permanent milestone copies — a bad checkpoint must never overwrite a good one.
3. Checksums: record weight SHA256 at save time and verify on resume — resuming from a corrupted checkpoint is worse than not resuming (the model silently degrades).

3. Resuming Training: More Than Loading Weights

3.1 Standard Resume Flow

1. Load config; validate hyperparameters match save time (lr, batch, parallelism)
2. Load weights + optimizer state + scheduler
3. Restore RNG state (data sampling order, dropout)
4. Restore dataloader position (epoch/sample)
5. Restore step/epoch counters (affects LR curve, log alignment)
6. Restore communication collective (NCCL world) — run one warmup step to verify

3.2 Consistency Verification: How to Confirm "It Resumed Right"

Check	Method	Failure meaning
Weight checksum	SHA256 vs saved record	Corrupted file
Optimizer state	Loss after resume should sit near the saved point	Optimizer state lost (resume = retrain)
Data order	Loss curve matches the "never interrupted" curve	RNG/dataloader not restored
Communication health	Warmup step: no NCCL errors, no abnormal retransmits	Environment changed (driver/topology)
Throughput	TPS after resume within ±5% of pre-interruption	Node degraded (cooling/clock)

Iron rule: after resume, run a warmup step + observe 3–5 steps of loss before releasing full training — an abnormal resume can silently destroy a training run 100 steps later.

3.3 Common Pitfalls

• LR scheduler not restored: LR restarts from the beginning → training "rewinds".
• RNG state not restored: data order changes → effectively a different training set; loss curve jumps.
• Cross-environment resume: resuming a checkpoint after a GPU model/driver change accumulates numerical drift — validate before full run.
• Distributed resume failure: under torch.distributed, checkpoints saved only on rank 0 and broadcast on resume; network jitter can silently fail — save on every rank + compare checksums is the safest.

4. Production Checklist

• [ ] Precision: BF16 default; FP16 with Loss Scaling configured; FP8 validated at small scale
• [ ] Master weights stay FP32, update path correct
• [ ] Checkpoint contains model/optimizer/scheduler/RNG/dataloader/config
• [ ] Atomic write (.tmp + os.replace) + rolling retention + milestone copies
• [ ] Save checksums, verify on resume
• [ ] Async save does not block training steps
• [ ] Resume flow includes warmup step + 3–5 step loss observation
• [ ] Save frequency matches RTO target: RTO = checkpoint size ÷ resume bandwidth + restart time

Further reading: Loss Spike Debugging (what to do when training gets unstable after resume), NCCL Tuning in Practice (communication-layer stability).