Sharded data parallelism in practice
By the end: You can configure FSDP or ZeRO and reason about what it saves and what it costs.
- The three ZeRO stages
- What each stage shards
- The communication it adds
- Activation checkpointing
New to distributed training? A GPU has a fixed amount of memory, and a large model plus everything training needs beside it does not fit in one. Distributed training spreads the work across many GPUs, often across many machines. There are two basic moves: give every GPU a complete copy of the model and a different slice of the data, then average what each of them learned; or cut the model itself into pieces so that each GPU holds only part of it. Why it matters
A model can fit on a GPU and still be untrainable on it, because training also needs gradients and optimiser state, which are three to six times larger than the weights. Sharded data parallelism fixes that without changing how you write the model: each device owns a slice of the state rather than a copy, and the pieces are gathered back only for the moment they are needed. This section covers what each stage shards, what it costs in communication, and activation checkpointing, which is the other half of making something fit and is almost always worth turning on.
Replication is the waste#
Plain data parallelism puts a full copy of the weights, gradients and optimiser state on every device. Eight devices therefore hold eight identical copies of the optimiser state, which for mixed-precision Adam is twelve of the sixteen bytes per parameter. Nothing about that is necessary: each device only ever updates the whole model because it happens to have the whole model. Give each one a slice to own and the memory divides by the device count.
The three stages#
stage shards per GPU, 7B on 8 devices extra traffic
none nothing (plain DDP) 104.3 GiB all-reduce
1 optimiser state 35.9 GiB same
2 + gradients 24.4 GiB same (reduce-scatter)
3 + parameters 13.0 GiB + all-gather, twiceThe stages are ordered by how cheap the sharding is, not by how much it saves. Optimiser state goes first because it is the largest component and is touched exactly once per step, at the update — so each device can hold its slice, update it, and nobody notices. Gradients go second, and this one is genuinely free: the all-reduce that data parallelism already performs is a reduce-scatter followed by an all-gather, and once gradients are sharded the all-gather half is simply not needed. Parameters go last because they are needed by every device on every forward and backward pass, so sharding them means gathering them back twice per step.
Stage 2 is usually the right default
It takes a 7B model from 104 GiB to 24 GiB per device with no additional communication over plain data parallelism, because the collective it uses is the cheaper half of the one that was already running. Reach for stage 3 when the model genuinely does not fit at stage 2, and expect to pay for it on a slow interconnect: the parameter all-gather happens twice per step, so on anything slower than NVLink it can cost more than it saves.
Activation checkpointing#
Sharding handles everything that scales with parameter count. Activations do not: they scale with batch size, sequence length and depth, and on a long-context model they can exceed everything else combined. The backward pass needs the tensors the forward pass produced, and a transformer block saves roughly sixteen hidden-sized tensors per token to make its backward pass cheap.
GiB = 1024 ** 3
layers, hidden, seq, batch, dtype = 32, 4096, 4096, 1, 2 # 7B-shaped, bf16
per_layer = batch * seq * hidden * dtype
full = layers * per_layer * 16 / GiB # ~16 saved tensors per block
checkpointed = layers * per_layer / GiB # only each block's input
print(f"nothing checkpointed : {full:5.1f} GiB")
print(f"blocks recomputed : {checkpointed:5.1f} GiB ({full/checkpointed:.0f}x less)")
print(f"compute cost : {4/3 - 1:.0%} more FLOPs per step")
# nothing checkpointed : 16.0 GiB
# blocks recomputed : 1.0 GiB (16x less)
# compute cost : 33% more FLOPs per stepPut that beside the sharding table and the decision becomes obvious. At stage 2 the per-device total goes from 40.4 GiB to 25.4 GiB, and at stage 3 from 29.0 to 14.0 — and the batch size is one, so on a realistic batch the activation term is the one that was going to break you. Thirty-three per cent more compute to make the job possible at all is a trade almost everyone takes.
The refinement worth knowing is that it does not have to be all or nothing. Checkpointing every block is the default and the simplest; checkpointing every second block recovers about half the saving for half the recompute. If you are close to fitting, that is the dial, and it is usually a single argument.
Choosing#
In practice the decision is short, and the mistake is starting further down it than necessary.
- If it fits on one device, use one device. Distributed training is slower per device, always.
- If it fits with a full replica per device, use plain data parallelism and nothing else.
- If the optimiser state is what does not fit, stage 2 with activation checkpointing. This is where most models land.
- If the parameters themselves do not fit, stage 3, and measure the throughput cost rather than assuming it.
- If a single layer does not fit on a device, sharding will not save you and you need tensor parallelism.
Measure the throughput, not just whether it starts
Every stage past 1 trades bandwidth for memory, and on a slow interconnect that trade can be terrible while still appearing to work. The failure looks like success: the job runs, the memory fits, and the throughput is a third of what it should be. Record samples per second at each configuration you try, and keep the one with the best throughput that fits rather than the one with the most headroom.
Where people go wrong
Reaching for stage 3 by default. It has the highest communication cost, so on a model that fits with stage 1 it makes training slower for no benefit.
If this comes up in an interview
What do the ZeRO stages shard?
Stage 1 shards optimiser state, stage 2 adds gradients, stage 3 adds the weights. Each saves more memory and adds more communication, with weights being the costliest because they must be gathered before every layer.
What is activation checkpointing?
Storing only some forward-pass activations and recomputing the rest during the backward pass. It trades roughly a third more compute for a large reduction in memory, and is often what makes a model fit without new hardware.
Build this: Train the same model with plain data parallelism, then stage 1, then stage 3 with checkpointing. Record peak memory and throughput for each.
Watch it explained
This is one section of Distributed training, 22 hours across 5 sections. Read it in full to get the whole topic in order.