Distributed AI Systems

4

Scaling with Fully Sharded Data Parallel (FSDP)

Training models larger than single GPU memory with parameter sharding

 

The data center is the new unit of computing.

Jensen Huang, CEO of NVIDIA

From DDP to FSDP

Fully Sharded Data Parallel (FSDP) is a training strategy that shards model parameters, gradients, and optimizer state across multiple devices so that each device holds only a fraction of the full model. In Chapter 3, we used DDP, which replicates the entire model on every GPU—effective when the model fits in a single GPU's memory. When the model (plus gradients and optimizer state) exceeds that memory, DDP is no longer viable. FSDP addresses this by distributing the model and its training state across GPUs, so you can train models ...