Fully Sharded Data Parallel (FSDP)

Technical Overview of FSDP

Fully Sharded Data Parallel (FSDP) is a memory-efficient data-parallel training strategy that shards all model states (parameters, gradients, and optimizer states) across the available GPUs in a cluster.

Standard Data Parallelism (DP) replicates the entire model on every GPU, which creates memory bottlenecks for larger models. FSDP addresses this by sharding the states, dynamically fetching the parameters needed for specific layers during the forward and backward passes. FSDP offers the simplicity of data parallelism while allowing developers to train models that exceed single-GPU memory capacities.

Key Architecture & Implementation

FSDP shards parameters using the Zero Redundancy Optimizer (ZeRO) protocol:

1. Forward Pass:
   • The active layer parameters are gathered from all GPUs using an All-Gather operation.
   • The layer executes, outputting activations.
   • The gathered parameters are immediately discarded, freeing up VRAM.
2. Backward Pass:
   • Layer parameters are gathered again using an All-Gather operation.
   • Gradients are calculated.
   • Parameters are discarded.
   • Gradients are reduced and sharded across GPUs using a Reduce-Scatter operation.

[ GPU 0: Shard A ] ---> All-Gather ---> [ Full Model Layer ] ---> Compute ---> Discard

Core Parameters

• Memory Savings: Reduces weight memory footprint by a factor of $N$ (the number of GPUs).
• Communication Overhead: Increases communication volume by approximately $50\%$ compared to standard data-parallel training.

Real-world Applications

• The primary framework used by Meta to train the Llama series.
• Supported natively in PyTorch FSDP and integrated into Hugging Face Accelerate.

How It Differs From Standard Data Parallelism

Standard data parallelism replicates the full model on every GPU, splitting only the training data across devices — this is simple but means every GPU needs enough memory to hold a complete copy of the model, its gradients, and its optimizer states, which becomes the binding constraint for very large models. FSDP instead shards the model’s parameters, gradients, and optimizer states themselves across all participating GPUs, so no single GPU ever holds a full copy of everything — each GPU only permanently stores its assigned shard, temporarily reconstructing the full parameters for a given layer just before it’s needed during the forward or backward pass, then discarding that reconstructed copy afterward to free the memory again. This trades some additional communication overhead (gathering and re-sharding parameters at each layer boundary) for a dramatic reduction in per-GPU memory requirements, which is what makes it possible to train models far larger than would fit using standard data parallelism on the same hardware. FSDP is built into PyTorch natively and is one of the standard distributed training strategies used for pretraining large open-weight models.