ZeRO-3 Memory: Budget Every Training Byte

ZeRO-3 memory is not simply model bytes divided by GPU count. This tutorial builds a per-rank ledger for parameters, gradients, optimizer states, collective peaks, activations, runtime overhead, and checkpoint policy before a training run.

Define the ZeRO-3 memory budgeting question

ZeRO-3 memory planning answers a narrower question than “will this model train?” It asks which persistent tensors and temporary communication buffers occupy each rank at a named instant, under a declared precision, optimizer, world size, and activation policy. A useful estimate therefore begins with a byte ledger, not a marketing ratio. Record parameter count, bytes per parameter for every representation, data-parallel degree, largest wrapped unit, bucket settings, and whether the estimate describes steady state or a transient peak.

The original ZeRO paper separates optimizer states, gradients, and parameters because each class can be partitioned independently. Preserve that taxonomy even if a framework exposes one convenient stage number. The ledger should also separate persistent model state from activations, allocator slack, kernels, communication workspaces, and framework bookkeeping. Only the first group follows the simple sharding equations.

A concrete scenario keeps the arithmetic honest. Suppose a model has 7 billion parameters, BF16 parameters and gradients, FP32 master parameters, and two FP32 Adam moments across eight ranks. State the assumed 16 bytes per parameter before sharding: 2 parameter bytes, 2 gradient bytes, 4 master bytes, and 8 moment bytes. Then compute ideal lower bounds and add measured overhead separately.

That distinction is the first ZeRO-3 memory invariant: a theoretical partition is not a device allocation prediction. The bounded lab in this article computes transparent lower bounds and one declared peak-buffer estimate. It does not predict fragmentation, kernel workspaces, throughput, or whether a real training stack fits.

Count the replicated baseline byte by byte

Before applying sharding, reconstruct ordinary data parallelism. Every rank holds the full parameter set, full gradient set, and full optimizer state. For the 7-billion-parameter example, the persistent baseline is 112 GB in decimal units: 14 GB of BF16 parameters, 14 GB of BF16 gradients, 28 GB of FP32 master weights, and 56 GB of Adam moments. This is a semantic inventory, not a single “model size” number.

Binary and decimal labels must not drift. A device marketed in GB is often inspected in GiB, so the receipt should include raw bytes and both human-readable conversions. The difference is large enough to overturn a borderline decision. Keep the calculation in integers where possible, then format only at the display edge.

The optimizer choice belongs in the input. Adam with a master copy has a different ledger from SGD, Adafactor, an eight-bit optimizer, or an optimizer that updates low-precision parameters directly. Do not hide these distinctions inside a generic multiplier. The same warning applies to gradient dtype and mixed-precision policy. “Sixteen bytes per parameter” is a scenario, not a law.

This replicated baseline gives ZeRO-3 memory comparisons a stable denominator. It also complements gradient accumulation token-loss accounting: accumulation changes the number of microbatches contributing to an update, but it does not automatically multiply the persistent parameter and optimizer tensors. Treat schedule, loss denominator, and memory ownership as related but separate ledgers.

Map each ZeRO stage to state ownership

The stages are easiest to remember as an ownership table. Stage 1 partitions optimizer state. Stage 2 also partitions gradients. Stage 3 additionally partitions parameters. Under ideal even partitioning across D data-parallel ranks, each sharded class contributes its full bytes divided by D, while an unsharded class remains fully replicated.

The DeepSpeed ZeRO tutorial describes this progression and the runtime coordination required to make partitioned state usable. Turn that description into explicit equations. For the example, Stage 1 retains 28 GB of parameters plus gradients on each rank and reduces the 84 GB optimizer portion to 10.5 GB at D=8. Stage 2 also divides the 14 GB gradients, while Stage 3 divides all 112 GB of persistent state to an ideal 14 GB per rank.

Those values are lower bounds. Divisibility may be uneven; padding, metadata, redundant copies during initialization, and buckets create additional allocations. World size is also the data-parallel sharding group, not necessarily the total accelerator count when tensor or pipeline parallelism is present. Record the actual group explicitly.

This ownership view makes ZeRO-3 memory debuggable and turns LLM training memory into named categories. If an observed rank holds far more parameter bytes than the ideal share, investigate initialization, wrapping boundaries, persistence thresholds, or a temporarily materialized module. If gradient bytes remain replicated, inspect the stage and synchronization path. The stage label is merely shorthand; the class-by-class ownership receipt is the proof.

Price precision and optimizer states independently

Mixed precision produces multiple representations with different lifetimes. A BF16 forward parameter may coexist with an FP32 master value used by the optimizer. Gradients may be BF16 or FP32. Adam commonly carries two moment tensors. Loss scaling adds small state, but fused implementations can add padding or flattened buffers that matter at scale. ZeRO-3 memory work should name each representation rather than multiply parameter count by an unexplained constant.

Use a table with columns for tensor class, dtype bytes, copies, sharding factor, persistent bytes per rank, and evidence source. The arithmetic is simple enough to review: parameter_count × bytes_per_element × copies ÷ owners. What matters is preventing a missing master copy or doubled moment count from disappearing into a rounded total.

Precision policy also interacts with communication. A framework may communicate gradients in a dtype different from the stored gradient, or cast gathered parameters for compute. Those temporary conversions belong in the peak ledger, not the persistent one. If an implementation changes its mixed-precision defaults, refresh the receipt even when model parameter count is unchanged.

This is where optimizer state sharding provides most of its early-stage savings: optimizer bytes often dominate the persistent baseline. But the correct decision is workload-specific. Report both absolute bytes and the percentage by class. A smaller optimizer state can make parameters or activations the new bottleneck, changing which technique deserves the next engineering hour.

Model all-gather and reduce-scatter peaks

Stage 3 saves persistent parameter memory by materializing parameter shards when a wrapped unit executes, then releasing or resharding them. The important peak is therefore not simply the persistent ZeRO-3 memory total. Add the largest simultaneously gathered parameter unit, communication buckets, prefetched units, and any full-parameter residency required by the selected policy.

The current PyTorch fully_shard documentation frames this behavior in terms of sharded parameters and per-module communication. Its API is not identical to DeepSpeed ZeRO, so use FSDP as a comparative implementation reference rather than claiming interchangeable configuration names. The shared planning principle is precise: wrapping granularity and communication schedule determine the temporary materialization peak.

A conservative lower-bound peak can add the largest unit’s missing remote shards to the persistent total. If one transformer block contains 400 million BF16 parameters across eight ranks, a rank already owns one eighth and may need roughly seven eighths of 800 MB, or 700 MB, for a full gathered view. Prefetching a second unit can nearly double that component. Reduce-scatter and gradient buckets add their own bounded allocations.

The lab makes these assumptions visible with fields for largest unit, prefetched units, and bucket bytes. A mutation that omits the gather peak must fail. This is also why distributed checkpointing needs a separate save/load ledger: checkpoint materialization can create a different peak from the forward and backward schedule.

Keep activation memory in a separate ledger

ZeRO partitions model states; it does not divide every saved activation by the data-parallel degree. Activation memory depends on microbatch size, sequence length, hidden dimensions, attention implementation, checkpoint boundaries, and parallelism along the sequence or model axes. Combining it with state bytes too early creates an attractive but unusable total.

Track activations as a measured or separately modeled input. If the training plan uses activation checkpointing, record the checkpoint policy and the resulting recomputation cost beside the measured peak. The activation checkpointing trade-offs article explains why a memory reduction needs a throughput and reproducibility receipt too. The inbound relationship matters in both directions: checkpointing can make the state budget fit, while aggressive sharding can make communication the dominant cost.

Reserve additional categories for CUDA or accelerator context, compiled graphs, attention workspaces, temporary logits, dataloader staging, and allocator fragmentation. These are not excuses for a vague safety factor. Measure them on a tiny representative run, then record an explicit overhead amount and confidence band.

The complete per-rank feasibility check is therefore persistent sharded state + modeled communication peak + activation peak + measured runtime overhead + contingency. ZeRO-3 memory is one large term in that equation, not the equation itself. When a run fails despite a passing paper budget, compare categories against the receipt instead of increasing a global multiplier until the estimate happens to match.

Use mutation tests to protect the budget

A spreadsheet can contain correct formulas today and quietly regress tomorrow. Treat the estimator as executable policy. Freeze a small scenario with hand-computed values, then mutate one rule at a time: replicate optimizer state at Stage 1, shard parameters at Stage 2, omit the FP32 master copy, divide activations by data parallelism, and erase the all-gather peak. Every mutant should produce a named failure, not merely a different grand total.

The oracle should be independent enough to catch shared mistakes. The article lab computes the production ledger class by class and compares it with a second direct stage table. It also checks monotonic lower bounds: persistent Stage 3 cannot exceed Stage 2 for the same inputs, and Stage 2 cannot exceed Stage 1. Those checks do not prove that a real runtime matches the model, but they protect the declared arithmetic.

Boundary tests matter too. Reject world size below one, negative parameter counts, unsupported byte widths, and prefetched-unit counts beyond the declared cap. Keep the computation bounded so a pasted value cannot freeze the page. Export raw integer bytes and input assumptions as JSON and CSV; formatted GB labels alone are insufficient for review.

The goal is not a single green badge. It is a receipt that tells a reviewer what was counted, who owns it, when it exists, and which claims remain outside the estimator. That is the same discipline needed for tensor parallelism in LLM inference, where device ownership and communication peaks must also be stated rather than implied.

Choose the stage and checkpoint policy together

Start with the least complex policy that clears the measured per-rank budget with explicit contingency. Stage 1 is attractive when optimizer state dominates and replicated parameters plus gradients still fit. Stage 2 adds gradient savings. Stage 3 becomes necessary when parameters themselves prevent the workload from fitting, but it also makes wrapping, gathering, prefetching, and checkpoint behavior first-class performance variables.

Compare candidates in one decision table: persistent bytes, estimated training peak, activation policy, checkpoint peak, expected communication operations, implementation maturity, and recovery procedure. A ZeRO-3 memory win is incomplete if the chosen checkpoint path materializes full state on one rank or if a restart cannot restore the same ownership layout.

Then validate progressively. Run the bounded calculator, inspect a framework initialization without a training step, execute one tiny forward/backward/update, and finally sample a representative sequence and microbatch. At each stage, capture allocated and reserved memory by rank. Differences between the estimate and observation become named overhead entries, not silent corrections to the equation.

The decision should be reversible. Preserve the receipt, configuration, framework versions, wrapper tree, and profiler evidence. Revisit when precision, optimizer, model architecture, world size, activation policy, or sharding implementation changes. ZeRO-3 memory planning succeeds when another engineer can reconstruct why the run was expected to fit—and identify which category moved when it does not.

Runnable local artifact — The lab calculates declared lower-bound state ownership and one conservative temporary-buffer model; it does not predict allocator fragmentation, kernel workspaces, throughput, or framework-specific fit.

Enter a bounded parameter count, world size, dtype bytes, optimizer copies, largest gathered unit, prefetch count, activations, and overhead; compare stages, kill five accounting mutations, and export integer-byte receipts.