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.