FSDP › FSDP: Fully Sharded Data Parallel

Training large language models requires distributing both computation and memory across many GPUs. As we explored in earlier chapters on data parallelism, the standard approach replicates model weights on every GPU and splits only the data batches. This works well when the model fits in GPU memory, but modern LLMs with billions of parameters make that assumption impossible. A single GPT-3-scale model requires hundreds of gigabytes of memory just for weights, optimizer states, and gradients combined, while even the most powerful GPUs today offer only 40-80 GB of high-bandwidth memory.

Fully Sharded Data Parallel (FSDP) solves this by sharding everything across GPUs: model parameters, optimizer states, and gradients are each split into equal fragments, so every GPU holds only a $1/N$ slice of each tensor. When a layer needs its full weights for a forward or backward pass, FSDP orchestrates an all-gather communication to reconstruct them temporarily, uses them, and then discards the reconstructed copy to reclaim memory. This approach makes it possible to train models that are $N$ times larger than what any single GPU could accommodate, while still achieving near-linear scaling in compute efficiency.

FSDP was introduced in PyTorch 1.12 as a production-ready implementation of the ideas pioneered by Microsoft's ZeRO optimizer. Understanding FSDP gives you the ability to train frontier-scale models using standard PyTorch, without requiring custom infrastructure or proprietary frameworks.