AI Infra 101 - Day 2 - Fully Sharded Data Parallel (FSDP2)

FSDP2 vs DDP

• DDP: each GPU stores the complete model state and synchronizes gradients through AllReduce.
• FSDP2: model parameters, gradients, and optimizer states are sharded. The required full parameters are temporarily reconstructed through AllGather only when computing a particular model, and gradients are aggregated and sharded again through ReduceScatter after the backward pass.
• Note the difference in gradient handling: DDP synchronizes gradients through AllReduce, with each GPU storing the complete gradients; FSDP2 synchronizes through ReduceScatter and stores only a shard.

  Prefetching: Overlapping Computation and Communication

• Suppose block5 is currently being computed and block6 is known to be next: start AllGather for block6 in advance.

  Reference Code

• benchmark_fsdp2.py