Data Parallelism and Sharded Optimizers
Lesson 2 of 5 in Distributed Training.
The first axis is the obvious one. If one device can hold the model but the run is too slow, make copies: put a full replica of the model on each of N devices, split every batch into N shards, and let each replica run its own forward and backward pass on its shard. Each replica now holds gradients computed from different data — so before anyone updates a weight, the replicas all-reduce their gradients: a collective operation after which every device holds the same averaged gradient vector, exactly as if a single device had processed the whole batch. Then every replica applies the identical optimizer step, and all copies stay in lockstep.
That is Data parallelism in one paragraph: replicate the model, split the data, average the gradients, step together. It is the workhorse of distributed training because it is simple, it maps onto any model, and the synchronization point is a single well-understood collective per step.
One training step under data parallelism
- Global batch
One optimizer step’s worth of training examples.
- Split into per-replica shards
Each replica sees a different slice of the batch; no example is processed twice.
- Replica 1: forward + backward
A full copy of the model computes gradients on its shard only.
- Replica 2 … N: forward + backward
All replicas run at the same time, on different data — this is where the speedup comes from.
- All-reduce: average the gradients
Collective communication. Afterwards every replica holds the same averaged gradients, as if one device had seen the whole batch.
- Identical optimizer step on every replica
Same gradients + same optimizer states everywhere → the copies stay exactly in sync.
- Next batch
Now the fine print. Data parallelism buys throughput and nothing else. Every replica still holds the full 16 bytes per parameter from lesson one — adding devices does not shrink per-device memory by a single byte, so a model that cannot train on one device cannot train on a thousand data-parallel devices either. The gradient all-reduce also moves a model’s worth of gradient data every step, so interconnect bandwidth sets a floor on step time. And the global batch grows with the replica count: past some point, ever-larger batches stop improving the model per token seen, so you cannot scale data parallelism forever and expect the Loss curve to keep pace.
Look back at the replicas, though, and there is something almost offensive about the memory picture: N devices holding N identical copies of the optimizer states — the largest single tenant — and N identical copies of everything else. The ZeRO insight is that this redundancy is pure waste.
| Stage | What gets sharded | What each device still holds | Communication vs plain DP |
|---|---|---|---|
Plain DP | Nothing — full replication | All model states: 16 bytes/param | Baseline: one gradient all-reduce per step |
ZeRO-1 | Optimizer states | Full weights and gradients; a 1/N slice of optimizer states | Same volume as plain DP (per the paper’s analysis) |
ZeRO-2 |
| Full weights; 1/N slices of gradients and optimizer states | Same volume as plain DP (per the paper’s analysis) |
ZeRO-3 |
| A 1/N slice of everything; full layers are gathered just-in-time and freed after use | About 1.5× plain DP — the price of gathering weights on the fly |
The fully-sharded end of this spectrum has become a mainstream framework feature: PyTorch’s FSDP (Fully Sharded Data Parallel) is, conceptually, the ZeRO-3 idea productized. Each device permanently owns a shard of every layer; just before a layer runs, the devices all-gather that layer’s full weights, use them, and immediately free them; gradients flow back the same way, reduce-scattered so each device keeps only the slice it owns. The result is data parallelism that does cut per-device memory — the answer to the question plain DP could not touch — paid for with extra communication woven through every forward and backward pass.
Why sharding the first two stages is (nearly) free
The surprise in the ZeRO paper is not that sharding saves memory — it is that the first two stages cost almost nothing extra to communicate. The reason: a standard all-reduce is already implemented as a reduce-scatter (each device ends with one fully-summed slice of the gradients) followed by an all-gather (everyone shares their slice until all hold the full result). If each device owns the optimizer state for exactly the slice it reduces, the reduce-scatter half already delivers each owner precisely the gradients it needs — sharding falls out of the communication pattern that plain data parallelism was using anyway. The paper’s analysis (§7) concludes stages 1 and 2 match plain DP’s communication volume, while stage 3 pays roughly 1.5× because weights must additionally be all-gathered during the forward and backward passes.
That asymmetry explains the practical default ordering: shard optimizer states first (biggest tenant, zero extra traffic), gradients second, and weights only when the model truly cannot be replicated — at which point you are trading step time for feasibility, and the tensor and pipeline axes of the next two lessons become the competing offer.
Interactive checkpoint quiz (1 questions) — open this page in a browser to take it.