Source-linked AI summary
PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel
Yanli Zhao, Andrew Gu, Rohan Varma, Liang Luo, Chien-Chin Huang, Min Xu, Less Wright, Hamid Shojanazeri, Myle Ott, Sam Shleifer, Alban Desmaison, Can Balioglu, Pritam Damania, Bernard Nguyen, Geeta Chauhan, Yuchen Hao, Ajit Mathews, Shen Li
TL;DR
Large-model training remains difficult because existing approaches can impose memory, usability, or framework-integration barriers. This paper presents PyTorch FSDP, which combines parameter sharding with PyTorch-aligned initialization, communication, and memory techniques. FSDP matches DDP on small models while supporting significantly larger models with near-linear TFLOPS scalability.
Problem
Large models increasingly power important applications, but existing methods can be architecture-specific or dependent on evolving framework internals, limiting generic large-model training.
Method
FSDP decomposes models into units, flattens and shards their parameters, materializes parameters on demand, and combines configurable sharding with deferred initialization and communication optimizations.
Results
FSDP achieves performance similar to DDP on small models while supporting significantly larger models with near-linear scalability in TFLOPS.
Takeaways & Limitations
FSDP provides a PyTorch-integrated training solution for large language and recommendation models across varied hardware configurations.
Takeaways & Limitations
FSDP cannot always preserve mathematical equivalence to local training for optimizer computations that depend on unsharded parameter values, tensor structure, or global parameter states.
Abstract
from arXiv · showhide
It is widely acknowledged that large models have the potential to deliver superior performance across a broad range of domains. Despite the remarkable progress made in the field of machine learning systems research, which has enabled the development and exploration of large models, such abilities remain confined to a small group of advanced users and industry leaders, resulting in an implicit technical barrier for the wider community to access and leverage these technologies. In this paper, we introduce PyTorch Fully Sharded Data Parallel (FSDP) as an industry-grade solution for large model training. FSDP has been closely co-designed with several key PyTorch core components including Tensor implementation, dispatcher system, and CUDA memory caching allocator, to provide non-intrusive user experiences and high training efficiency. Additionally, FSDP natively incorporates a range of techniques and settings to optimize resource utilization across a variety of hardware configurations. The experimental results demonstrate that FSDP is capable of achieving comparable performance to Distributed Data Parallel while providing support for significantly larger models with near-linear scalability in terms of TFLOPS.
1 INTRODUCTION
Large models are increasingly important, but existing training methods can be architecture- or framework-dependent. FSDP presents a PyTorch-aligned approach that reduces memory use while supporting broad usability and large-scale training.
- Motivation: Large language and recommendation models have grown to hundreds of billions or more than one trillion parameters.These models support applications across language processing and services used by billions of people.
- Challenges: Existing scaling methods can be tightly coupled to model architectures or vulnerable to changes in machine-learning framework internals.The paper motivates a native solution co-designed with framework core components.
- Approach: FSDP flattens and shards parameters within model units, materializing only one unit at a time before discarding recovered parameters.This design significantly reduces peak memory consumption while preserving on-demand computation.
- Approach: Deferred initialization, configurable sharding, and memory-management techniques help FSDP provide a local-training-like experience for models that may not fit on one GPU.The paper describes these techniques as responses to user-experience and resource-utilization challenges.
- Results: FSDP matches DDP performance on small models and supports significantly larger models with near-linear TFLOPS scalability.Evaluations used language and recommendation models on up to 512 80GB A100 GPUs.
2 BACKGROUND
The background contrasts replication, partitioning, and sharding strategies for distributed training. These approaches trade memory requirements, communication, and implementation complexity in different ways.
- Model Replication: Model replication distributes computation across devices while maintaining a model copy on each device.DDP synchronizes gradients with AllReduce and overlaps communication with backward computation.
- Model Replication: DDP requires all parameters, gradients, and optimizer states to fit in one GPU’s memory, limiting its support for large models.This constraint follows from maintaining a complete replica on every device.
- Model Partitioning: Pipeline parallelism partitions layers into stages, while tensor parallelism shards parameters and communicates activations at layer boundaries.Both distribute model computation across multiple devices using different partitioning units.
- Parameter Sharding: Parameter sharding reduces per-device memory by assigning each rank only a parameter shard, but requires communication or altered computation for correctness.The background identifies parameter communication and activation communication as two broad approaches.
3 SYSTEM DESIGN
FSDP decomposes a model into independently managed units and keeps parameters and gradients sharded except when a unit’s computation requires unsharded values.
- Unit-Based Execution: FSDP decomposes the model into smaller units and handles each unit independently during training.Only one unit’s unsharded parameters and gradients are materialized at a time.
- Forward Pass: Before forward computation, FSDP gathers the parameters for the active unit, runs its local computation, and frees the gathered peer shards.Other units remain sharded throughout the forward pass.
- Backward Pass: During backward computation, FSDP recovers the active unit’s parameters, then frees peer shards and ReduceScatters gradients after computation.Each rank retains only a shard of parameters and gradients after backward computation.
- System Optimizations: FSDP provides optimizations and configuration knobs covering model initialization, sharding strategies, communication, and memory management.These options address variation in model structures and hardware capabilities.
3.1 Model Initialization
FSDP addresses large-model initialization by postponing tensor storage and replaying initialization operations, then materializing and sharding one unit at a time.
- Initialization Challenges: Before FSDP, PyTorch required full model materialization on one device unless users modified model source code to place submodules separately.This created a challenge for smoothly transitioning from local to distributed training.
- Deferred Initialization: Deferred initialization allocates parameter tensors on a fake device, records initialization operations, and replays them when tensors move to a GPU.The mechanism postpones storage allocation until a concrete device is available.
- Initialization Challenges: FSDP cannot always shard initialization directly because user-defined initialization logic may require unsharded parameters.The system therefore prepares unsharded parameters while controlling memory use.
- Unit-Wise Initialization: FSDP initializes one unit at a time and shards each unit before moving to the next.Combined with deferred initialization, this limits materialization during model construction.
3.2 Sharding Strategies
FSDP exposes sharding strategies that trade memory footprint against communication and throughput, from full replication through full and hybrid sharding. Its flattened parameter units improve collective efficiency while unit granularity controls the memory-throughput balance.
- Sharding strategies: FSDP’s sharding factor F ranges from 1 for full replication to W for full sharding, with intermediate values defining hybrid sharding.Full replication simplifies to vanilla data parallelism, while full sharding leaves each device holding one model shard.
- Full sharding: Full sharding minimizes memory but incurs 1.5x DDP communication overhead under a bandwidth-optimal ring algorithm.This communication cost motivates careful organization of collectives.
- Collective efficiency: Smaller than 33M-element AllGathers sharply increase total communication time when total communication is fixed at approximately 1B FP32 elements.The experiment varies AllGather size while holding total communication constant.
- Parameter flattening: FSDP flattens parameters within each unit, concatenates them into FlatParameters, pads to divisibility by F, and chunks them evenly across ranks.This supports arbitrary original parameter shapes while limiting padding to at most F−1.
- Parameter flattening: FlatParameter layouts match AllGather and ReduceScatter inputs and outputs, enabling collectives without additional tensor copies.The same flattened representation owns the storage of original parameters and gradients.
- Memory-throughput trade-off: Finer-grained FlatParameters reduce peak memory but may lower throughput because the number of collectives grows as O(N).Users control this trade-off by choosing how submodules are wrapped into FSDP units.
- Hybrid sharding: Hybrid sharding combines sharding and replication, can exploit datacenter locality, and offers a tunable memory-throughput trade-off through F.It can reduce cross-host traffic, while smaller-world-size AllReduce operations empirically perform better than global-scale collectives.
3.3 Communication Optimizations
FSDP provides native communication optimizations for overlapping communication with computation and reducing communication overhead. These include overlap, backward and forward prefetching, and two gradient-accumulation modes.
- Communication optimizations: FSDP incorporates four native communication optimizations: overlapping, backward prefetching, forward prefetching, and accumulation.These techniques target communication efficiency across the training loop.
- Overlap: Unlike DDP, FSDP’s eager forward issues the next AllGather after the computation it overlaps because the next FlatParameter is not known early enough.This difference in kernel-issue order reflects FSDP’s parameter materialization pattern.
- Backward prefetching: FSDP’s single NCCL stream can expose consecutive ReduceScatter and AllGather operations on the backward critical path.The current ReduceScatter blocks the next AllGather, which can delay subsequent gradient computation.
- Backward prefetching: Backward prefetching issues the next AllGather before the current ReduceScatter using recorded reverse forward module order.Because forward order is recorded each iteration, the approach supports dynamism across iterations.
- Forward prefetching: Forward prefetching uses the previous iteration’s module order to issue the next AllGather before current-unit computation for static computational graphs.This targets workloads where slow CPU execution delays forward AllGather issuance.
- Gradient accumulation: FSDP supports gradient accumulation with or without communication, trading increased memory for reduced communication in the latter mode.Without communication, ranks retain unsharded gradients rather than reducing them across ranks.
3.4 Memory Management
FSDP’s memory management addresses CUDA caching allocator behavior across asynchronous producer and consumer streams. A rate limiter constrains in-flight AllGathers to preserve block reuse while retaining overlap.
- Caching allocator: The CUDA caching allocator reuses internally managed memory blocks to avoid frequent cudaMalloc and cudaFree calls, but stream interactions complicate safe reuse.Its decisions occur on the CPU thread before the GPU kernel requiring an allocation runs.
- Caching allocator: Per-stream allocation blocks can become unavailable to another stream, causing allocator failures and blocking cudaFrees even when GPU memory remains available.This over-allocation problem is especially relevant when producer and consumer streams run asynchronously.
- Rate limiter: FSDP uses a rate limiter that blocks the CPU thread when necessary to improve caching-allocator block reuse.It permits at most two in-flight AllGathers, the minimum needed to maintain communication-computation overlap.
4 IMPLEMENTATION
FSDP’s implementation integrates parameter sharding, communication scheduling, initialization strategies, wrapping choices, and mixed precision with PyTorch’s module and autograd systems. These mechanisms target memory efficiency, execution-order alignment, and non-intrusive use while preserving model structure where possible.
- APIs: FSDP exposes both a model wrapper and a module annotator, with the latter preserving model structures and fully qualified parameter names.The wrapper replaces submodules with FSDP units, while fully_shard installs forward and backward hooks.
- Initialization: FSDP initialization can use an unsharded model on GPU or CPU, while deferred initialization remains preferable for some large-model settings.GPU initialization requires the entire model to fit on one device; CPU initialization supports larger models but may be substantially slower because of limited bandwidth and parallelization.
- Parameter flattening: FSDP unit boundaries determine when AllGather and ReduceScatter occur, so users should align units with model execution order when possible.Nested module annotation forms FlatParameters from annotated modules and assigns residual parameters to their parents; execution-order-based reconstruction was also explored.
- Runtime communication: FSDP inserts forward and backward communication through module and autograd hooks to coordinate parameter gathering and gradient reduction.Forward hooks place pre-forward and post-forward logic, while Tensor and autograd hooks anchor communication to backward execution.
- Runtime communication: The implementation integrates FSDP with PyTorch’s nn.Module and autograd engine in a non-intrusive and efficient manner.This integration supports the FSDP algorithm without altering its core sharding design.
- Native mixed precision: FSDP’s native mixed precision independently configures parameter, gradient-reduction, and buffer precisions while dynamically materializing unsharded parameters.Its design keeps local sharded FlatParameters on GPU memory, casts per FlatParameter rather than per operator, and can run collectives in low precision.
5 EVALUATION
The evaluation examines FSDP across model sizes, communication strategies, and large-model workloads. Results show that FSDP supports models beyond DDP’s memory limit, benefits variably from communication controls, and scales efficiently on large clusters.
- Model Scale: FSDP and DDP perform similarly on 611M and 2.28B T5 models, but DDP runs out of memory above 2.28B while FSDP accommodates T5-11B.Turning on BF16 gives FSDP significantly higher TFLOPS for the 11B model.
- Model Scale: Backward prefetching produces an approximately 18% speedup on GPT-175B, with the TFLOPS gain persisting across GPU cluster sizes.The evaluation enables backward prefetching in subsequent experiments.
- Throttle Communications: Rate limiting has inconsistent effects: it provides no RegNet speedup, impedes DeepViT, and helps only when aggressive allocation causes defragmentation.CUDA malloc retry can indicate whether defragmentation occurred when latency or traces are inconclusive.
- Throttle Communications: T5 experiments achieve up to 5X speedups with rate limiting, whereas DeepViT incurs 5% overhead when delayed AllGather blocks dependent computations.Practitioners should verify defragmentation before enabling rate limiting.
- Efficient Training for Large Models: Full Sharding with RAF minimizes DHEN memory footprint but reduces QPS, while Hybrid Sharding with NRAF trades higher memory use for higher QPS.Adding GPUs consistently decreases peak memory because each rank holds a smaller model shard.
- Efficient Training for Large Models: The 175B model reaches more than 173 and 186 TFLOPS per GPU at batch sizes 1 and 2, with linear TFLOPS scaling from 128 to 512 GPUs.For T5-11B, per-GPU TFLOPS regresses 7% from 8 to 512 GPUs as communication begins to outweigh computation.
6 RELATED WORK
FSDP is positioned as a drop-in data-parallel alternative that reduces parameter redundancy, unlike approaches requiring model-specific changes or extensive configuration. Related methods offer complementary memory or scaling benefits but may introduce communication, tuning, accuracy, or generalization trade-offs.
- DDP replicates the model on every device and cannot accommodate increasingly large model sizes.
- Pipeline parallelism requires model changes, microbatch tuning, stage partitioning, and intricate scheduling to optimize performance.
- Compiler-based methods search across data, tensor, and pipeline parallelism configurations using profiling, performance modeling, annotations, or search.
- FSDP provides a drop-in replacement for data parallelism that reduces redundancy along the data-parallel axis.
- Gradient compression, mixed precision, rematerialization, and CPU offloading save memory but can affect accuracy or add compression, recomputation, or transfer overhead.
7 DISCUSSION
FSDP can be combined with pipeline and tensor parallelism, but the combination introduces communication or memory constraints. The discussion also identifies optimizer-equivalence and shared-parameter handling as important adoption caveats.
- FSDP can wrap individual pipeline stages, but default full sharding may unshard parameters for every microbatch and cause significant communication overhead.
- Tensor parallelism keeps parameters sharded during computation when a submodule cannot fit in GPU memory, while FSDP applies sharded data parallelism across a second mesh dimension.
- FSDP may not preserve mathematical equivalence to local training for optimizer computations depending on unsharded parameter values, tensor structure, or global parameter states.
- Restoring such optimizer equivalence may require uneven sharding, padding, or extra communication, which hurts performance.
- Shared parameters should belong to the lowest-common-ancestor FSDP unit, though this may keep them unsharded for a large interval and requires model-structure inspection.
8 CONCLUSION
The paper presents FSDP as a PyTorch 2.0 design combining usability and efficiency techniques through close co-design with PyTorch components. Evaluations show near-linear scalability for large language and recommendation models.
- FSDP combines deferred initialization, flexible sharding, communication overlap and prefetching, and rate-limited collectives to improve usability and efficiency.
- FSDP’s techniques are co-designed with key PyTorch components to ensure a sound and robust solution.
- FSDP enables large language and recommendation models with near-linear scalability.