This September 2023 paper introduces PyTorch Fully Sharded Data Parallel (FSDP), an advanced solution designed to scale the training of exceptionally large machine learning models. It addresses limitations of previous methods like Distributed Data Parallel (DDP) by sharding model parameters, gradients, and optimizer states across multiple GPUs, thereby drastically reducing individual GPU memory consumption. FSDP employs various techniques, including deferred initialization, flexible sharding strategies, and optimizations for communication overlap and prefetching, to ensure high efficiency and a user-friendly experience. The research demonstrates FSDP's effectiveness in training models with billions of parameters, achieving near-linear scalability and enabling broader access to large model development. The authors also discuss interoperability with other parallelization paradigms and known limitations.
Source:
https://arxiv.org/pdf/2304.11277
