PyTorch FSDP: Fully Sharded Data Parallelism

• PyTorch FSDP is a fully sharded approach that partitions model parameters, gradients, and optimizer states across GPUs to save memory.
• It uses on-demand all-gather and reduce-scatter operations to synchronize data, trading memory savings for increased communication overhead.
• FSDP is applied in large-scale model training and research, with variants optimizing compiler-driven runtimes and addressing sparse computations.

PyTorch Fully Sharded Data Parallel (FSDP) is PyTorch’s implementation of zero-redundancy data parallel training in which model parameters, gradients, and optimizer states are partitioned across participating GPUs, full parameters are materialized with all-gather only when needed for computation, and gradients are returned to shards with reduce-scatter after backward. In the literature, FSDP is treated as the ZeRO-3 style fully sharded configuration, and it is widely used for training large-scale models; recent work also uses the same abstraction in self-supervised vision training, formal neural network verification, adaptive batch-size control, and compiler-driven runtime optimization (Ovi, 19 May 2025, Mehta, 5 Jan 2026, Wang et al., 25 Feb 2026).

1. Core execution model

In the standard formulation, FSDP shards model parameters, gradients, and optimizer states across all GPUs, so that each GPU stores only a fragment of the model rather than a full replica. During the forward pass, parameters are gathered on-demand using an all-gather operation; the backward pass employs reduce-scatter to compute averaged gradients; and each GPU applies optimizer updates to its own shard only (Ovi, 19 May 2025). A closely related description states that fully sharded data parallelism shards each parameter tensor equally across $P$ GPUs, performs an all-gather before executing a layer, and frees the full tensor immediately after that layer finishes so that only the shard is retained (Vorobyov et al., 8 Jun 2026).

This execution pattern distinguishes FSDP from Distributed Data Parallel (DDP). DDP maintains identical model copies on all devices and typically communicates gradients with all-reduce after backward, whereas FSDP replaces model replication with sharded persistent state and repeated parameter materialization (Gokmen et al., 3 Nov 2025). In practical terms, DDP has no communication in the forward pass and one or a few gradient all-reduces in backward, while FSDP introduces many all-gather and reduce-scatter calls at module or parameter-group granularity (Ovi, 19 May 2025).

A compiler-oriented reformulation makes the same semantics explicit in PyTorch primitives. In "SimpleFSDP" (Zhang et al., 2024), parameters are stored as sharded DTensors, a parametrization all-gathers shards into a replicated parameter on demand, selective activation checkpointing frees the gathered parameter after forward and re-gathers it for backward, and backward through DTensor redistribution yields reduce-scatter automatically. This suggests that FSDP is not defined by a single wrapper implementation so much as by a stable execution contract: sharded persistent state, transient materialization for compute, and sharded optimizer ownership.

2. Memory and communication semantics

A systematic placement analysis characterizes ZeRO Stage 3 and FSDP with the state placements

$$(\pi_\Theta,\pi_\Omega,\pi_G,\pi_A) = (S^*, S, S, R),$$

where parameters use sharded-with-gather placement, optimizer state and gradients are sharded, and activations are replicated within each data-parallel worker (Mehta, 5 Jan 2026). In that framework, if $|\Theta| = 2P$ bytes, $|G| = 2P$ bytes, and $|\Omega| = 12P$ bytes, then FSDP’s model-state memory per device is

$M_{\text{FSDP}} = \frac{16P}{N} + s_{\text{unit}},$

with $s_{\text{unit}}$ denoting the transient reconstruction unit for gathered parameters (Mehta, 5 Jan 2026). The same framework derives FSDP’s communication volume as

$$C_{\text{FSDP}} = \frac{N-1}{N}\,|G| + 2\frac{N-1}{N}\,|\Theta|,$$

which implies a $1.5\times$ communication factor relative to ordinary data parallelism, while reducing model-state memory by a factor of $N$ in the idealized case (Mehta, 5 Jan 2026).

For decoder-only transformers, one hardware model writes the parameter count as

$$(\pi_\Theta,\pi_\Omega,\pi_G,\pi_A) = (S^*, S, S, R),$$0

where $$(\pi_\Theta,\pi_\Omega,\pi_G,\pi_A) = (S^*, S, S, R),$$1 is the number of decoder blocks and $$(\pi_\Theta,\pi_\Omega,\pi_G,\pi_A) = (S^*, S, S, R),$$2 is the hidden size (Wang et al., 4 Mar 2025). With parameter and gradient memory $$(\pi_\Theta,\pi_\Omega,\pi_G,\pi_A) = (S^*, S, S, R),$$3 and optimizer-state memory $$(\pi_\Theta,\pi_\Omega,\pi_G,\pi_A) = (S^*, S, S, R),$$4, the free memory under full sharding is expressed as

$$(\pi_\Theta,\pi_\Omega,\pi_G,\pi_A) = (S^*, S, S, R),$$5

and the maximum number of tokens per device is bounded by

$$(\pi_\Theta,\pi_\Omega,\pi_G,\pi_A) = (S^*, S, S, R),$$6

(Wang et al., 4 Mar 2025). The same analysis derives a throughput upper bound

$$(\pi_\Theta,\pi_\Omega,\pi_G,\pi_A) = (S^*, S, S, R),$$7

which makes explicit that FSDP performance depends jointly on available GPU memory and interconnect bandwidth, not only on raw FLOP rate (Wang et al., 4 Mar 2025).

These formulas give a compact interpretation of FSDP’s trade-off. Sharding lowers persistent model-state memory from $$(\pi_\Theta,\pi_\Omega,\pi_G,\pi_A) = (S^*, S, S, R),$$8 to approximately $$(\pi_\Theta,\pi_\Omega,\pi_G,\pi_A) = (S^*, S, S, R),$$9, but that saving is accompanied by repeated parameter movement. A plausible implication is that FSDP’s practical regime is determined by which term dominates the step time: transient parameter traffic, activation memory, or arithmetic throughput.

3. PyTorch implementation lineage and system architectures

PyTorch and closely related systems now expose several distinct realizations of fully sharded training. DeepSpeed ZeRO-3 and PyTorch FSDP1 flatten many tensors into large concatenation buffers and shard them element-wise; PyTorch FSDP2 instead represents each parameter as a DTensor with placement Shard(0); and Megatron-FSDP returns to concatenation-based sharding but adds padding so shards conform to Shard(0) checkpoint semantics (Wang et al., 25 Feb 2026). These design choices preserve the ZeRO-3 contract but differ substantially in memory fragmentation, padding, and copy behavior.

"SimpleFSDP" (Zhang et al., 2024) re-implements FSDP semantics in a torch.compile-friendly form using DTensor, parametrizations, and selective activation checkpointing. Its compiler backend performs IR-node bucketing and reordering for communication–computation overlap, and reported evaluations on Llama 3 models, including the 405B model, show up to $|\Theta| = 2P$0 memory reduction and $|\Theta| = 2P$1 throughput improvement compared to the most widely adopted FSDP2 eager framework (Zhang et al., 2024). A distinct compiler-driven line, "DeepCompile" (Tanaka et al., 14 Apr 2025), compiles a fully sharded approach like ZeRO-3 and FSDP and applies proactive prefetching, selective unsharding, and adaptive offloading; it reports up to $|\Theta| = 2P$2 performance improvement over FSDP baselines and up to a $|\Theta| = 2P$3 throughput increase with limited GPU resources using offloading (Tanaka et al., 14 Apr 2025).

At larger scale and with more demanding sharding formats, "veScale-FSDP" (Wang et al., 25 Feb 2026) redesigns FSDP around a flexible sharding placement called RaggedShard, a structure-aware planning algorithm, and a Distributed Buffer abstraction. The stated motivation is that existing FSDP systems struggle with block-wise quantized training and with non-element-wise optimizers such as Shampoo and Muon because fixed element-wise or row-wise formats conflict with block-structured computations (Wang et al., 25 Feb 2026). VeScale-FSDP reports $|\Theta| = 2P$4–$|\Theta| = 2P$5 higher throughput and $|\Theta| = 2P$6–$|\Theta| = 2P$7 lower memory usage than existing FSDP systems while scaling efficiently to tens of thousands of GPUs (Wang et al., 25 Feb 2026).

This implementation lineage shows that “PyTorch FSDP” is best understood as a family of ZeRO-3-compatible systems rather than a single code path. The invariant is the placement of persistent state; the differences lie in how sharded tensors are represented, planned, gathered, and recycled.

4. Correctness, exactness, and alternative semantics

A formal treatment of distributed training correctness isolates two necessary and sufficient conditions: gradient integrity, meaning that the gradient used for updates is exactly the global batch mean,

$|\Theta| = 2P$8

and state consistency, meaning that any state tensor accessed or communicated is bitwise identical across participating devices up to associative reduction and consistent dtypes (Mehta, 5 Jan 2026). In ordinary FSDP, these conditions imply that parameter all-gather must complete before any forward or backward use of a layer, and gradient reduce-scatter must complete before optimizer updates on shards.

A particularly clear correctness-sensitive case study adapts FSDP to formal neural network verification. In "Scaling Neural Network Verification with Tensor Parallelism and Fully Sharded Data Parallelism" (Vorobyov et al., 8 Jun 2026), only weight matrices are sharded; a per-layer all-gather reconstructs the full weight before bound propagation, and the full tensor is freed immediately afterward. Because the arithmetic path is unchanged, the resulting bounds are reported as bitwise identical to the single-GPU baseline, with baseline memory drops of $|\Theta| = 2P$9–$|G| = 2P$0 and peak-memory reductions of $|G| = 2P$1–$|G| = 2P$2 on wide MLPs (Vorobyov et al., 8 Jun 2026). This suggests that FSDP can act as an exact storage transformation rather than an algorithmic approximation when the gathered tensor is used in precisely the same arithmetic order as the baseline.

Not all recent work accepts FSDP’s collective-synchronization semantics as fixed. "Revisiting Parameter Server in LLM Post-Training" (Wan et al., 27 Jan 2026) introduces On-Demand Communication, which adapts parameter-server ideas into FSDP by replacing collective all-gather and reduce-scatter with direct point-to-point communication. The stated effect is to reduce the synchronization barrier from once per layer to once per minibatch, and the paper reports up to a $|G| = 2P$3 speedup over standard FSDP on imbalanced LLM post-training workloads (Wan et al., 27 Jan 2026). A plausible implication is that FSDP’s semantic core—sharded persistent state plus synchronous correctness—does not uniquely determine a collective-communication backend.

5. Performance trade-offs and hardware bottlenecks

Empirical comparisons against DDP consistently show the central trade-off of FSDP. A study on the GMU Hopper cluster reports that FSDP reduces GPU memory usage by over $|G| = 2P$4, but increases training time by up to $|G| = 2P$5 compared to DDP; for ConvNeXt_Large in particular, memory usage fell by more than $|G| = 2P$6, while end-to-end training time was substantially longer because of repeated all-gather and reduce-scatter operations (Ovi, 19 May 2025). The same study observes that DDP typically shows more stable and higher average SM utilization, whereas FSDP exhibits more fluctuations corresponding to communication-heavy periods (Ovi, 19 May 2025).

A broader hardware analysis argues that the critical interplay is between cluster connection bandwidth and GPU memory size compared to computational performance, and that this interplay limits training efficiency under FSDP (Wang et al., 4 Mar 2025). Real experiments up to 512 A100 GPUs and model sizes up to 310B parameters show that longer sequence lengths can raise Model FLOPS Utilization, but that efficiency still degrades as model size increases or as interconnect bandwidth falls from 200 Gbps to 100 Gbps (Wang et al., 4 Mar 2025). This makes the usual rule of thumb precise: FSDP is most attractive when memory is the binding constraint and the interconnect is strong enough that parameter gathers can be overlapped or amortized.

Several systems target this communication side directly. "Network-Offloaded Bandwidth-Optimal Broadcast and Allgather for Distributed AI" (Khalilov et al., 2024) argues that the FSDP pipeline overlaps Allgather and Reduce-Scatter extensively, so concurrent collectives can compete for injection bandwidth and create pipeline bubbles. It proposes a multicast-based Allgather that achieves $|G| = 2P$7 traffic reduction on a 188-node testbed and, when paired with in-network Reduce-Scatter, yields an idealized speedup of $|G| = 2P$8 for overlapping Allgather and Reduce-Scatter (Khalilov et al., 2024). "TAGC" (Polyakov et al., 8 Apr 2025) instead compresses FSDP gradient communication and reports up to $|G| = 2P$9 training speedup compared to standard FSDP, while "QSDP" (Markov et al., 2023) quantizes both weights and gradients in FSDP collectives and reports end-to-end speedups of up to $|\Omega| = 12P$0 while preserving model accuracy on GPT-family models (Polyakov et al., 8 Apr 2025, Markov et al., 2023).

These results are not contradictory. They indicate that FSDP is neither inherently slow nor inherently fast; its realized efficiency depends on whether the system can turn repeated parameter movement into high-bandwidth, overlap-friendly traffic rather than serialized collective overhead.

6. Extensions, applications, and current research directions

FSDP is now used well beyond conventional large-language-model pretraining. In self-supervised vision, "DINO-MX" (Gokmen et al., 3 Nov 2025) exposes FSDP as a configurable alternative to DDP, launched through Hugging Face Accelerate, with a high-level distribution block, bf16 mixed precision, and compatibility with LoRA, layer freezing, knowledge distillation, and multi-crop training. In that framework, FSDP is treated as an execution detail rather than a change in the DINO training algorithm (Gokmen et al., 3 Nov 2025). In language-model pretraining, adaptive batch-size schedules have been implemented with PyTorch FSDP for TinyLlama 1.1B and OpenLlama 3B, with reported improvements over constant batch sizes and heuristic batch warmup schedules (Lau et al., 2024).

Research on parallel composition has likewise treated FSDP as one dimension in broader training topologies. "Breadth-First Pipeline Parallelism" (Lamy-Poirier, 2022) combines a looping pipeline schedule with fully sharded data parallelism and reports up to $|\Omega| = 12P$1 increase in training throughput for a 52B-parameter model using a small batch size per GPU compared to Megatron-LM (Lamy-Poirier, 2022). The paper’s analysis suggests that FSDP communication becomes markedly more effective when grouped and overlapped stage-wise rather than repeated microbatch by microbatch.

Specialized descendants of FSDP now target sparse and matrix-structured training. "Hecate" (Qing et al., 4 Feb 2025) introduces Fully Sharded Sparse Data Parallelism for Mixture-of-Experts, with SparseAllGather and SparseReduceScatter that only move expert parameters and gradients required on each device; the reported result is up to $|\Omega| = 12P$2 speedup over state-of-the-art MoE training systems (Qing et al., 4 Feb 2025). "MatrixFSDP" (Gao et al., 7 Jul 2026) changes ZeRO-3 placement so that each 2D weight matrix has one owner rank holding the full matrix and other ranks holding empty shards, thereby making matrix optimizers communication-free at optimizer time; it reports optimizer-step latency reductions of $|\Omega| = 12P$3 on one node and $|\Omega| = 12P$4 on eight nodes relative to stock FSDP2-Muon, with up to $|\Omega| = 12P$5 end-to-end speedup (Gao et al., 7 Jul 2026).

Taken together, these developments show that FSDP has become a general systems template for distributed deep learning. Its canonical form is ZeRO-3-style state sharding with all-gather and reduce-scatter. Its active research frontier is not whether to shard, but how to express richer placements, relax synchronization in workload-specific ways, and preserve the memory advantages of sharding while accommodating block-wise quantization, sparse expert routing, matrix-valued optimizers, and compiler-scheduled communication.