fsdp

Data parallelism is the simplest and most common way to scale training across many GPUs: replicate the entire model on every device, give each replica a different slice of the batch, and average the gradients so all copies stay identical. ZeRO (Zero Redundancy Optimizer) and its PyTorch implementation FSDP (Fully Sharded Data Parallel) keep the same data-parallel structure but remove its biggest weakness — every GPU storing a full copy of the model state — by sharding those states across the GPUs and gathering them only when needed.\n\n**Plain data parallelism trades memory for simplicity.** Each GPU holds the complete model and processes its own micro-batch, then all replicas all-reduce their gradients each step to converge on one update. It is easy and communication-light, but wasteful: every GPU redundantly stores the full parameters, the full gradients, and — the biggest cost — the full optimizer states (for Adam, momentum and variance, often several times the size of the weights). For large models that redundancy, not compute, is what makes the model not fit.\n\n**ZeRO/FSDP shards the redundant state across GPUs.** Instead of N identical copies, ZeRO partitions the model state into N slices and gives each GPU just one. ZeRO does this in stages: stage 1 shards optimizer states, stage 2 adds gradients, stage 3 adds the parameters themselves (this full-shard mode is what FSDP implements). When a layer needs to run, the GPUs all-gather that layer's parameters just in time, compute, then immediately free the gathered copy — so peak memory holds only one shard plus the layer currently in flight. Per-GPU memory drops roughly N-fold.\n\n| State | Plain data parallel | ZeRO-3 / FSDP |\n|---|---|---|\n| Parameters | full copy per GPU | 1/N per GPU |\n| Gradients | full copy per GPU | 1/N per GPU |\n| Optimizer states | full copy per GPU | 1/N per GPU |\n| Communication | all-reduce grads | all-gather params + reduce-scatter grads |\n| Memory per GPU | ~O(full model) | ~O(model / N) |\n\n```svg ```\n\n**The trade is memory for communication.** Sharding replaces plain data parallelism's single gradient all-reduce with an all-gather of parameters on the way into each layer and a reduce-scatter of gradients on the way out — more bytes on the wire per step. Because that traffic is frequent, FSDP leans on fast fabrics (NVLink within a node, InfiniBand across nodes) and overlaps communication with compute to hide it. The payoff is that a model far too large to replicate now fits, letting pure data parallelism scale to model sizes that would otherwise force tensor or pipeline parallelism.\n\nRead data parallelism and ZeRO/FSDP through a quant lens rather than a 'copy the model' lens: plain DP costs O(full model) memory per GPU for one gradient all-reduce, while ZeRO-3/FSDP costs O(model/N) memory in exchange for gathering and re-scattering state each layer. The design question is the memory-versus-bandwidth balance at your N and fabric speed — shard until the model fits and the extra all-gather traffic still overlaps with compute, since past that point communication, not capacity, becomes the binding constraint.