FSDP › FSDP Architecture

FSDP works by sharding the three memory-intensive quantities, namely parameters, gradients, and optimizer states, across all participating GPUs. The elegance of the design is that the sharding is largely transparent to the model code: from the perspective of each individual layer, it receives its full parameter tensor when it needs it, computes normally, and returns its result. The orchestration of gathering and scattering tensors happens behind the scenes.

Consider a linear layer with a weight matrix $W$ of shape $(d_{\text{out}}, d_{\text{in}})$, containing $P_W = d_{\text{out}} \times d_{\text{in}}$ total parameters. With $N$ GPUs, FSDP flattens $W$ into a 1D tensor of length $P_W$ and assigns an equal contiguous slice to each GPU:

• GPU 0 holds elements $0$ through $\lfloor P_W/N \rfloor - 1$
• GPU 1 holds elements $\lfloor P_W/N \rfloor$ through $\lfloor 2P_W/N \rfloor - 1$
• GPU $k$ holds elements $\lfloor kP_W/N \rfloor$ through $\lfloor (k+1)P_W/N \rfloor - 1$

Each GPU also holds a corresponding $1/N$ shard of the gradient tensor for those same parameters, and a $1/N$ shard of the Adam optimizer states (first moment $m$ and second moment $v$) for those same parameter positions. This three-way alignment is what makes the local optimizer step possible: when it is time to update parameters, GPU $k$ has all the information it needs to perform the Adam update for its slice without any communication.

The critical constraint is that computing the forward pass through a linear layer requires the complete weight matrix $W$, not a shard of it. An input $x$ of shape $(B, d_{\text{in}})$ needs to be multiplied by the full $W$ to produce output of shape $(B, d_{\text{out}})$. FSDP handles this requirement with explicit communication phases.

FSDP wraps each module (or group of modules) in a unit called an FSDP unit. These units are the atoms of the sharding system: all-gather and reduce-scatter operations occur at FSDP unit boundaries. Understanding the communication pattern for one unit clarifies the entire system.

During the forward pass, the following sequence occurs for each FSDP unit:

1. Every GPU broadcasts its local parameter shard to all other GPUs via an all-gather. After this operation, all GPUs temporarily hold the complete parameter tensor for this unit.
2. Each GPU uses the full parameter tensor to compute its portion of the forward pass (processing its slice of the data batch).
3. Immediately after the computation completes, FSDP discards the reconstructed parameter tensor on each GPU, leaving only the original local shard. This is the key memory reclamation step.

During the backward pass, the cycle repeats with one addition:

1. An all-gather reconstructs the full parameter tensor again (necessary for computing parameter gradients).
2. The backward computation runs, producing full-parameter gradient tensors on each GPU.
3. A reduce-scatter aggregates gradients across all GPUs and distributes the result. Unlike all-reduce (which gives every GPU the full summed gradient), reduce-scatter performs the sum and simultaneously shards it: GPU $k$ receives only the portion of the gradient corresponding to its parameter shard.
4. The reconstructed parameter tensor is discarded again.

The optimizer step then runs locally on each GPU: each GPU updates its own parameter shard using its own gradient shard and its own optimizer state shard. No communication is required at this step, because the data for each parameter's update resides entirely on the GPU responsible for that parameter.

This complete cycle, all-gather then compute then discard (forward), all-gather then compute then reduce-scatter (backward), local optimizer step, is what defines FSDP's operation. The name "fully sharded" reflects the fact that all three memory-intensive components remain sharded at rest; they are only temporarily gathered when computation demands it.

Under FSDP with $N$ GPUs, the memory breakdown per GPU changes dramatically compared to DDP:

• Parameters (steady state): $4P/N$ bytes (only the local shard persists)
• Gradients (steady state): $4P/N$ bytes (only the local gradient shard persists)
• Optimizer states: $8P/N$ bytes (first and second moments for the local shard)
• All-gather buffer during forward: $4P$ bytes temporarily (freed after each layer's computation)

The steady-state memory per GPU is approximately $(4 + 4 + 8)P/N = 16P/N$ bytes for the sharded quantities. The peak memory includes the temporary all-gather buffer of $4P$ bytes during computation for the current FSDP unit, but this is freed right after the unit completes. For a model split into $M$ FSDP units, the peak at any instant is $16P/N + 4P/M$ bytes (steady-state plus one unit's full parameters).

Compare this to DDP, where each GPU holds $16P$ bytes at all times. The reduction factor from $16P$ to $16P/N$ for the steady state is exactly $N$, the number of GPUs. This linear scaling in memory capacity with GPU count is what makes FSDP so powerful for scaling to large models.

To see this concretely: a 13B parameter model requires about 208 GB of training memory ($16 \times 13$ GB). With FSDP on 8 GPUs of 40 GB each (320 GB total), each GPU holds only $208/8 = 26$ GB of steady-state memory, fitting comfortably within 40 GB when including activation buffers.