ZeRO and FSDP

Plain DDP wastes memory: every GPU holds a full copy of the parameters, gradients, and optimiser state. With Adam in mixed precision that is 16 bytes per parameter, replicated N times. The Zero Redundancy Optimiser (ZeRO) from DeepSpeed observed that the data-parallel group can shard each of those three buffers across ranks and pay only extra communication to materialise them on demand. Fully Sharded Data Parallel (FSDP) is PyTorch's native implementation of the same idea, now the recommended path for training models that do not fit on a single GPU.

The three ZeRO stages

The ratios assume Adam optimiser and mixed-precision training where optimiser state dominates. ZeRO-1 alone gets most of the memory saving for a fraction of the engineering cost; ZeRO-3 is what you reach for when even the parameters do not fit.

How FSDP implements ZeRO-3

FSDP groups parameters into "FSDP units" (typically each transformer block). For each unit, during the forward pass:

1. AllGather the sharded parameters from all ranks so the full unit's weights exist locally.
2. Compute forward for that unit.
3. Free the gathered parameters; keep only the local shard.

During backward, for the same unit (in reverse order):

1. AllGather the parameters again.
2. Compute backward, producing local gradients for the full unit.
3. ReduceScatter the gradients: each rank ends up with the summed gradient for its shard only.
4. Free the gathered parameters and the full gradients.

The optimiser step runs on the local shard with the local gradient and local optimiser state. No rank ever materialises the full model state.

The overlap that makes it tractable

A naive implementation would gather a unit, compute, free, gather the next unit, and so on - serialising comms and compute. FSDP overlaps them: while unit k is computing, the AllGather for unit k+1 is already in flight, and the ReduceScatter for unit k-1 is also overlapping. With well-tuned bucket sizes and a healthy NVLink / InfiniBand fabric you can hide most of the extra communication behind compute.

forward:    [gather 0][compute 0][free 0]
                      [gather 1][compute 1][free 1]
                                [gather 2][compute 2][free 2]
                                          ...

Communication that is not hidden becomes your bottleneck. Profile with nsys or PyTorch's profiler before adding more parallelism dimensions.

Why FSDP often beats tensor parallelism

Megatron-style tensor parallelism (TP) shards individual matrix multiplications across GPUs and synchronises with AllReduce inside every transformer layer. It works well within a single node where NVLink bandwidth is high, but degrades sharply over InfiniBand. FSDP's communication is at FSDP-unit granularity (whole layers), which packs into larger messages and tolerates slower interconnects much better.

A rough rule:

• TP is competitive only within an NVLink-connected node (typically 8 GPUs).
• FSDP scales across nodes because its messages are large enough to amortise InfiniBand latency.
• For the largest models you compose them: TP within the node, FSDP across nodes.

Communication vs memory trade-off

The communication volume of ZeRO-3 is roughly 1.5x DDP's, not 3x as a naive count suggests, because ReduceScatter + AllGather together carry the same data as one AllReduce. The cost is latency and bookkeeping, not total bytes on the wire.

When it falls down

• Tiny FSDP units. Wrapping at too fine a granularity (every Linear layer) produces many small AllGathers; latency dominates. Wrap at transformer-block boundaries.
• CPU-side bottlenecks. FSDP's per-unit Python dispatch can become the bottleneck on small models with many units. Use use_orig_params=True and the FSDP2 (fully_shard) API where available.
• Mixed-precision sharp edges. If your optimiser state is FP32 but gradients are BF16, the AllGather and ReduceScatter happen in the param dtype. Mismatched dtypes silently break the overlap.
• Slow interconnect. ZeRO-3 over a single 25 Gbps Ethernet link will be slower than just buying a bigger GPU. Real ZeRO-3 needs at least 100 Gbps RDMA between nodes.

Further reading

• ZeRO: Memory Optimizations Toward Training Trillion Parameter Models - Rajbhandari et al., the original DeepSpeed paper that defined the three stages.
• PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel - the engineering paper for FSDP itself.
• Meta Engineering: Fully Sharded Data Parallel - introductory write-up with diagrams of the gather / compute / reduce-scatter dance.
• PyTorch FSDP docs - the current API surface and wrapping policies.