Data Parallelism › Beyond DDP: Fully Sharded Data Parallel

Pure DDP replicates the complete model on every GPU. This is simple and efficient when the model fits in GPU memory, but as models grow beyond the memory capacity of a single GPU (80 GB for A100/H100), pure DDP is impossible without additional techniques.

Fully Sharded Data Parallel (FSDP), introduced by Facebook AI Research and available in PyTorch as torch.distributed.fsdp.FullyShardedDataParallel, extends DDP by sharding the data, model parameters, gradients, and optimizer states across GPUs.

In FSDP, each GPU holds only $1/N$ of the parameters at any given time. Before the forward pass through each layer, FSDP gathers the full parameter tensor for that layer from all GPUs (an all-gather operation). After the forward pass for that layer, FSDP discards the gathered parameters, retaining only its own shard. During the backward pass, FSDP re-gathers parameters on demand, computes gradients, reduces those gradients across GPUs (a reduce-scatter), and again discards the gathered parameters.

The memory savings are dramatic. For a 7B model with AdamW in mixed precision:

• Pure DDP: 16 bytes per parameter times 7 billion = 112 GB per GPU
• FSDP (with optimizer state sharding): approximately 16 bytes per parameter divided by $N$ GPUs = 112/N GB per GPU

With $N = 8$ GPUs, FSDP reduces per-GPU memory from 112 GB to 14 GB, enabling training that is otherwise impossible.

The cost of FSDP is increased communication volume. FSDP performs additional all-gather operations for parameter reconstruction before each layer's forward and backward pass. This roughly doubles the communication volume compared to pure DDP. However, FSDP can overlap these all-gather operations with the computation of the preceding layer, similar to how DDP overlaps gradient all-reduce with the backward pass.

The choice between DDP and FSDP depends primarily on model size relative to GPU memory:

• If your model (plus gradients and optimizer states) fits on a single GPU, use DDP. It is simpler, has lower communication overhead, and is easier to debug.
• If your model does not fit on a single GPU but you need data parallelism, use FSDP. It achieves the same data-parallel scaling as DDP while reducing per-GPU memory by a factor of $N$.
• For the largest models (hundreds of billions of parameters), even FSDP may be insufficient without also combining tensor parallelism or pipeline parallelism, which we cover in subsequent chapters.

The ZeRO optimizer (covered in the ZeRO chapter) implements the same sharding strategy as FSDP through the DeepSpeed library. ZeRO Stage 1 shards only optimizer states, Stage 2 adds gradient sharding, and Stage 3 adds parameter sharding. Stage 3 is roughly equivalent to FSDP in terms of memory savings and communication volume.