ZeRO and FSDP: Sharding What Data Parallelism Duplicates

Look again at what data parallelism does with 64 GPUs. Each holds the same parameters, computes different gradients that are then averaged to be identical, and runs an optimizer that produces the same update from the same inputs.

After the all-reduce, all 64 replicas hold byte-identical model state. The system stores 64 copies of something that is, by construction, the same on every device.

The Zero Redundancy Optimizer starts from that observation. If the state is identical everywhere, no device needs all of it. Each can hold one sixty-fourth, and the pieces can be gathered when they are actually needed.

The subtlety is that the pieces are needed at different times and for different durations. Optimizer states are touched once per step, during the update. Gradients are needed at the end of the backward pass. Parameters are needed continuously, in every layer of every forward and backward pass.

That ordering, from rarely-needed to constantly-needed, is exactly the order of the three ZeRO stages.

Stage 1 partitions only the 32-bit optimizer states: the master weights, the first moment and the second moment. Those are 12 of the 16 bytes per parameter, so it is the largest single term and the easiest to move.

Each GPU becomes responsible for updating one slice of the parameters. It keeps optimizer state only for that slice.

The step now runs slightly differently. Gradients are still all-reduced as before, so every GPU sees the full averaged gradient. Each then applies the optimizer update to its own slice only, producing updated 16-bit parameters for that slice. Finally an all-gather distributes the updated slices so everyone has the complete parameter set for the next forward pass.

Memory per GPU falls from 16 bytes per parameter to 4 plus 12 over N.

The communication accounting is the appealing part. The reduce-scatter plus all-gather move the same total volume as the single all-reduce they replace. Stage 1 is close to free: large memory saving, no extra communication.

Stage 2 adds gradients to what is sharded. Since each GPU only updates its own parameter slice, it only ever needs the gradients for that slice. So instead of an all-reduce that leaves everyone with every gradient, a reduce-scatter leaves each GPU with only the reduced gradients it will use.

Memory falls to 2 bytes per parameter plus 14 over N. Communication volume is still unchanged from plain data parallelism, because reduce-scatter is precisely the first half of a ring all-reduce.

Stage 3 shards the parameters themselves, and this is where the character changes. No GPU holds the full model at any moment. Before computing a layer, the GPUs all-gather that layer's parameters, use them, and immediately discard the copies they do not own. The same happens again in the backward pass.

Memory falls to 16 over N bytes per parameter, which is full linear scaling.

The cost is real: parameters are now gathered twice per step rather than once, so communication volume rises by about fifty percent over the previous stages.

Abstract formulas hide how sharply the stages differ. Here is the memory of model state per GPU for a 70-billion-parameter model.

Read the columns rather than the rows. Plain data parallelism is flat: more GPUs never help.

Stages 1 and 2 improve quickly at first and then flatten, because each retains an unshardable residue. Stage 1 keeps 4 bytes per parameter no matter what, which is 280 GB for this model, so it can never fit on an 80 GB device however large the cluster. Stage 2 keeps 2 bytes, or 140 GB, with the same conclusion.

Only stage 3 scales without a floor, and it is the difference between a 70B model being untrainable and comfortable.

The four rows are the terms and the stages progressively move them from replicated to sharded.

Under plain data parallelism everything is replicated. Under stage 1, optimizer states become sharded and the rest stay replicated. Under stage 2, gradients join them. Under stage 3, parameters do too, and nothing that scales with model size is replicated any more.

What never changes across the stages is activations. They are per-GPU by construction, since each device processes its own micro-batch, so sharding has no bearing on them. If activations are your problem, no ZeRO stage will help, and you need recomputation, tensor parallelism or a smaller micro-batch.

The right column tracks the communication consequence. Stages 1 and 2 restructure existing communication without adding volume, which is why they are close to free. Stage 3 introduces gathers that did not previously exist, in both the forward and backward pass, which is why it costs about fifty percent more traffic.

Fully Sharded Data Parallel is PyTorch's native implementation of these ideas, and DeepSpeed's ZeRO is the other widely used one. The concepts transfer; the vocabulary differs.

FSDP's unit of sharding is a wrapped module. You choose which parts of the model become FSDP units, and each unit's parameters are flattened, concatenated and split across ranks. At runtime, entering a unit triggers an all-gather of its parameters, the unit executes, and the non-owned shards are freed.

That choice is the main tuning decision. Wrap too coarsely, for instance the whole model as one unit, and you gather everything at once, which defeats the purpose. Wrap too finely, per individual linear layer, and you pay per-gather latency constantly. The usual answer is one unit per transformer block, which the auto-wrap policies provide directly.

Because each gather is a network round trip, FSDP prefetches: while block k computes, the parameters for block k plus one are already in flight. Correctly overlapped, communication hides behind computation and the sharding is close to free in wall-clock terms.

The theory says communication hides behind computation. Whether it actually does is the difference between a run at good utilisation and one at half that, and the failure modes are recognisable.

The blocks are too small. Overlap works when a block computes for longer than the next block's gather takes. Small hidden dimensions or tiny micro-batches mean compute finishes first and the device waits.

The network is slower than assumed. Sharding across nodes on ordinary Ethernet rather than a high-speed fabric can make gather time exceed compute time for any block size, at which point stage 3 is strictly worse than stage 2.

The wrapping is wrong. One giant unit gathers the whole model; hundreds of tiny units pay latency repeatedly.

Something forces a synchronisation. Logging a loss value, computing a metric that requires a full tensor, or a debugging print that calls an item conversion will serialise the pipeline and destroy the overlap in a way no profiler flag makes obvious.

The symptom in all four cases is the same: adding GPUs stops adding throughput.

Both DeepSpeed and FSDP can push state off the accelerator entirely, into host memory or onto NVMe storage. It is the technique that lets a single machine train a model it has no business training, and it should be understood as a last resort rather than a tuning option.

The arithmetic is unforgiving. Host memory is reached over PCIe, which is roughly an order of magnitude slower than the accelerator's own memory bandwidth, and NVMe is slower again by a wide margin.

Optimizer state offload is the most defensible variant, because optimizer state is touched exactly once per step, during the update. Move 12 of the 16 bytes per parameter to the host, and if the transfer is overlapped with the backward pass, the cost can be modest.

Parameter offload is a different proposition, since parameters are needed continuously. Expect a large slowdown.

The honest framing: offload converts an impossible run into a slow one. When the alternative is not training at all, that is a good trade. When more GPUs are available, they are almost always the better answer.

The memory table makes stage 3 look strictly better. It is not, and choosing it reflexively is a common way to leave throughput on the floor.

Stage 3 costs about fifty percent more communication than stages 1 and 2, and that traffic has to hide behind computation. On a well-provisioned cluster with a high-bandwidth fabric and reasonably sized blocks, it usually does. On a cluster with slower interconnect, it does not, and the run gets slower rather than faster.

The decision rule follows the memory table directly. Compute what each stage needs per GPU. Pick the lowest-numbered stage that fits, with headroom for activations.

For a 7B model on 8 GPUs, stage 1 is comfortable and costs nothing. For a 70B model, stages 1 and 2 have floors of 280 GB and 140 GB that no cluster size removes, so stage 3 is the only option and its communication cost is simply the price of admission.

Stage 3 is what makes large models trainable. It is not a free upgrade for models that did not need it.

Two limits remain after every stage, and they are what the next lesson exists to address.

Activations are untouched. Sharding partitions model state across data-parallel ranks, and activations are not model state. A device processing its own micro-batch holds all of that micro-batch's activations. For long contexts this becomes the dominant term, and no stage of ZeRO reduces it by a byte.

And a single layer must still fit. Stage 3 gathers parameters layer by layer, which means that at the moment a layer executes, its full parameters and its activations must be resident on one device. For most models this is unremarkable. For very wide models, where a single feed-forward matrix is itself enormous, the gather does not fit and sharding has run out of road.

Both limits point the same way. Sharding distributes copies of things; it never splits an individual computation. To split the computation itself you need tensor parallelism, which divides a matrix multiply across devices, or pipeline parallelism, which divides the layer stack.