ZeRO’s Memory Budget, Byte by Byte

A billion parameters stored in two-byte precision occupy two decimal gigabytes. Why might a training job need much more memory before activations even enter the picture? Parameters are only one category of persistent model state.

ZeRODivide training state across workers so each worker does not keep every copy. partitions otherwise duplicated training states across data-parallel workers. Its progression partitions optimizer state first, then gradients as well, and finally parameters. Communication makes the required pieces available when computation needs them. The reduction in persistent redundancy must therefore be considered alongside temporary materialization and communication. ZeRO paper

Write down the bytes before the stages

Choose a specific illustrative mixed-precision Adam layout: two-byte parameters, two-byte gradients, a four-byte master parameter copy, and two four-byte optimizer moments. That totals 16 bytes per parameter. This is an accounting assumption, not a claim that every optimizer or framework uses this exact layout.

For one billion parameters, an unsharded worker stores 2 GB of parameters, 2 GB of gradients, and 12 GB of master parameters plus moments. The total is 16 GB per worker. Here GB means 1,000,000,000 bytes; a GiB uses 1,073,741,824 bytes and would produce a different numeric label.

Now distribute the model state across four logical workers. With optimizer-state partitioning alone, the first two categories remain replicated. Per-worker storage becomes 2 + 2 + 12/4 = 7 GB.

Partitioning gradients too yields 2 + 2/4 + 12/4 = 5.5 GB. Partitioning parameters as well yields (2 + 2 + 12)/4 = 4 GB of persistent model state per worker.

The accompanying code calculates each category separately. That makes a changed dtype visible instead of hiding it in a convenient constant. For example, changing gradients from two to four bytes changes both the unsharded total and the savings attributable to gradient partitioning.

Four gigabytes is not a fit guarantee

The stage-three number is a persistent-state subtotal. A device still needs memory for activations, communication buffers, gathered parameters, temporary operator workspaces, and allocator behavior. Some categories overlap in time; others may be simultaneously live. Peak memory depends on that lifetime schedule, not merely the sum of long-lived tensors.

A useful capacity worksheet therefore has two layers. First, calculate model-state ownership from parameter count, dtypes, optimizer choice, and worker count. Second, account for the peak working set under the intended execution schedule. Mixing these layers can make a configuration appear feasible on paper while failing at the first large layer.

There is also a system-level difference between reducing per-worker storage and reducing total cluster storage. In this example, four unsharded workers hold 64 GB of model state altogether. Fully partitioned persistent state totals 16 GB across them. Additional temporary buffers can still create replication while an operation runs.

The calculator accounts for persistent state from the declared dtypes and ownership rules. A framework-specific capacity estimate must add the activation and gathered-layer working set, then check the resulting peak against the device limit.

parameters = 1_000_000_000
workers = 4
bytes_per_parameter = {
    "parameters": 2,
    "gradients": 2,
    "optimizer": 4 + 4 + 4,  # Master weights and two moments.
}
partitioned = {
    0: set(),
    1: {"optimizer"},
    2: {"optimizer", "gradients"},
    3: {"optimizer", "gradients", "parameters"},
}
totals = []
for stage in range(4):
    parts = {}
    for name, width in bytes_per_parameter.items():
        divisor = workers if name in partitioned[stage] else 1
        parts[name] = parameters * width / divisor / 1_000_000_000
    total = sum(parts.values())
    totals.append(total)
    print(f"stage={stage}: params={parts['parameters']:.1f} GB, "
          f"grads={parts['gradients']:.1f} GB, "
          f"optimizer={parts['optimizer']:.1f} GB, total={total:.1f} GB")
assert totals == [16.0, 7.0, 5.5, 4.0]
assert totals[-1] * workers == totals[0]

stage=0: params=2.0 GB, grads=2.0 GB, optimizer=12.0 GB, total=16.0 GB
stage=1: params=2.0 GB, grads=2.0 GB, optimizer=3.0 GB, total=7.0 GB
stage=2: params=2.0 GB, grads=0.5 GB, optimizer=3.0 GB, total=5.5 GB
stage=3: params=0.5 GB, grads=0.5 GB, optimizer=3.0 GB, total=4.0 GB