Reproduce “PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel” on OpenResearch

About this paper

Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri and 10 more

FSDP enables large model training by sharding parameters, gradients and optimizer states across GPUs and gathering each unit for computation. The system uses deferred initialization and PyTorch-native integration to address initialization and memory constraints. On minGPT-175B, it reports more than 186 TFLOPS/GPU at batch size 2, across 192–512 80-GB A100 GPUs; a 128-GPU batch-2 run was reduced by memory defragmentation.

Key insights

• The combination of FSDP units and `FlatParameter` with on-demand materialization effectively minimizes peak GPU memory usage by ensuring only one unit's full parameters are present at any given time.
• Sophisticated communication-computation overlapping techniques, such as backward prefetching and using separate CUDA streams for `AllGather` operations, effectively hide communication latency and boost training throughput.
• A novel rate limiter for `AllGather` operations and careful memory management mitigate issues with PyTorch's CUDA caching allocator in multi-stream environments, preventing performance degradation caused by memory defragmentation.

Summary by alphaXiv.