ZeRO & FP8 Training Ledger: A Practical Guide to Sharding Memory

Training a 70B model: weights alone are 140 GB; with gradients and optimizer states, one 80G card can't hold it. The industry answer is two "save" tricks: ZeRO shards states across cards, and FP8 lowers precision to save memory. This article settles both ledgers, linking to Training Resource Accounting and Distributed Parallelism.

1. The Four Training States (Recap)

State	Size (70B, FP16 training)	Note
Weights W	140 GB	model params
Gradients G	140 GB	backward gradients
Optimizer state O	280-420 GB	Adam m/v (2-3 copies)
Activations A	batch-dependent	saved forward, recomputed backward

Conclusion: weights are only 20-30% of training memory; optimizer states are the big chunk — that's ZeRO's starting point.

2. ZeRO Three Stages: From "Everyone Full" to "Everyone One Slice"

Stage	What's sharded	Per-card memory (vs full)	Comm overhead
DDP baseline	none	100%	gradient All-Reduce
ZeRO-1	optimizer states	~50%	unchanged
ZeRO-2	optimizer + gradients	~25%	unchanged
ZeRO-3	optimizer + gradients + weights	~1/N (8 cards ≈12%)	weight All-Gather

Key takeaways:

• ZeRO-1/2 add zero communication (still that one gradient All-Reduce) — free memory, enable by default;
• ZeRO-3 trades extra All-Gather for memory; comm grows ~1.5× — use only at the memory limit;
• Recompute per-step comm time with Communication Ledger to size ZeRO-3's scaling cost.

3. Combining ZeRO with Parallelism

DP (data parallel) ── gradient All-Reduce, more cards is better
TP (tensor parallel) ── inside supernode, comm-frequent
PP (pipeline parallel) ── layer-sliced, comm-sparse
ZeRO (state sharding) ── stacks with DP, saves memory without comm (stage 1/2)

Experience: small clusters (8-32 cards) → DP+ZeRO-2 is the most hassle-free; large models (70B+) → "DP×TP×PP×ZeRO-3" four-way combination; fix TP/PP first (topology constraints), then squeeze memory with ZeRO-3.

4. FP8 Mixed Precision: The Other Half of Doubled Compute

FP8 training (Hopper+) gains and costs:

Item	FP16 baseline	FP8	Gain
Tensor Core throughput	1×	~2×	throughput doubles
Weight/gradient memory	2B/param	1B/param	memory halves
Dynamic range	large	small (fewer exponent bits)	needs loss scaling

Engineering points:

• FP8 isn't "FP8 everywhere" — typically FP8 for forward GEMM, backward & optimizer stay FP16/FP32 (mixed precision);
• Needs loss scaling + overflow detection (see Mixed-Precision Checkpoints);
• The gain is capped by the compute ledger: only compute-bound training can eat FP8's 2× (see Accelerator Arithmetic).

5. Reproducible Python: 70B Training Memory Budget

def zero_mem(n_params_b, world, stage, fp8=False):
    """Per-card memory per ZeRO stage (GB, ignoring activations)"""
    w = n_params_b * 1e9 * (1 if fp8 else 2) / 1e9      # weights GB
    g = n_params_b * 1e9 * (1 if fp8 else 2) / 1e9      # gradients
    o = n_params_b * 1e9 * 4 / 1e9 * 2                   # Adam m+v, FP32
    parts = {"0": w+g+o, "1": w+g+o/world, "2": (w+g)/world+o/world,
             "3": (w+g+o)/world}
    return round(parts[str(stage)], 1)

zero_mem(n_params_b=70, world=64, stage=3, fp8=True)
# ≈ 3.3 GB/card (weights/grads/optimizer only; activations extra) → 70B fits in 64 cards

6. Checklist

1. Enable ZeRO-1/2 by default — free memory, zero extra comm;
2. Only add ZeRO-3 when memory is still tight; measure comm ratio (>30% → fix topology);
3. FP8 training on Hopper+: FP8 forward + FP16/32 backward/optimizer + loss scaling;
4. Run the memory budget function above first, then compare with Acceptance & Baselines measurements;
5. Run a small-scale baseline after each parallelism change to confirm scaling efficiency.

---

Organized and written by this site for learning purposes; specific versions and measurements are subject to on-site conditions.