What is FSDP? Fully sharded data parallel, explained

In this article

1. Before you start
2. Why interviewers ask this
3. FSDP from first principles
4. Worked example
5. Trade-offs and edge cases
6. How to answer it in an interview
7. Common mistakes
8. Follow-up questions
9. Sources

Before you start

This page is one step in the machine learning (ML) coding path.

Why interviewers ask this

FSDP is the first thing a candidate reaches for when a model no longer fits on one GPU, so the question separates people who have trained large models from people who have read about them. It appears in two forms: explain how FSDP works and what it costs, or implement a toy version across simulated ranks (a rank is the process that drives one GPU, numbered 0 to N minus 1). Implementing a toy is the same kind of task as beam search from scratch.

A weak answer says "it splits the model across GPUs", which describes model parallelism instead. A strong answer says what is split (three kinds of state), when full weights exist (only briefly, per layer), and which two collective operations carry the traffic. The spoken order to practise is: the problem, the three shards, the two collectives, the cost. This page builds each part from zero, as one of the questions in an ML coding round.

FSDP from first principles

Shard

A shard is an equal slice of a tensor that one GPU stores, so that N GPUs together hold exactly one copy of the whole tensor.

All-gather

An all-gather is a collective operation in which every GPU contributes its shard and every GPU receives the full, concatenated tensor.

Reduce-scatter

A reduce-scatter is a collective operation that adds the same-sized tensors from all GPUs element by element and gives each GPU only its own slice of the sum.

All-reduce

An all-reduce is a collective operation that leaves every GPU with the full element-by-element sum of everyone's tensors; it can be done as a reduce-scatter followed by an all-gather.

In plain words

Picture four friends who share one huge atlas. Plain data parallelism gives each friend a full copy, which is wasteful. FSDP tears the atlas into four parts and gives each friend one. When someone needs a page from another part, the friends pass the parts around, the page is read, and the borrowed parts go back. The analogy breaks in one place. Real friends read whatever they like, while in FSDP every GPU needs the same layer at the same moment, so the passing is a collective step that every GPU joins together.

In one precise sentence: FSDP keeps one shard of every parameter, gradient and optimizer state on each GPU. It all-gathers a layer's parameters just before using them, and reduce-scatters the gradients so each GPU updates only its own shard. The idea comes from the Zero Redundancy Optimizer (ZeRO) of Rajbhandari et al. (2020), and PyTorch's implementation is described by Zhao et al. (2023).

Why does plain data parallelism run out of memory?

Plain data parallelism stores 16 bytes of model state per parameter on every GPU, so adding GPUs adds speed but never room, because every new GPU repeats the same full copy of the state. The ZeRO paper counts 2 bytes for each 16-bit floating point (fp16) parameter, 2 bytes for its gradient and 12 bytes of optimizer state for mixed-precision Adam, which is 16 bytes in total (Rajbhandari et al., 2020, section 3.1). A 7.5 billion parameter model therefore needs 120 GB of state on each GPU before a single activation (the intermediate output each layer keeps for the backward pass) is stored.¹ See 16-bit formats for why the optimizer keeps higher-precision copies.

Step 1: Shard everything at rest

Each GPU permanently stores one shard of every parameter, gradient and optimizer state, so the memory for state falls from 16 bytes per parameter to 16 divided by N, where N is the number of GPUs. The PyTorch tutorial says that outside forward and backward computation the parameters stay fully sharded, and the optimizer updates the sharded parameters with sharded gradients, so its states are sharded too (PyTorch FSDP tutorial, updated September 2025).

In our four-rank toy, each rank keeps 40 floats of a 160-float model.¹ At this point no GPU can compute anything alone, because it holds only a quarter of every layer. The next step fixes that one layer at a time.

1 · Shard at rest2 · All-gather a layer3 · Forward, then free4 · Backward, gather again5 · Reduce-scatter, updateReplay

Step-by-step description

1. Every GPU stores one slice of every parameter.
2. The slices of one layer are combined on every GPU, and the next layer is fetched early.
3. The layer runs, then the full copy is released.
4. The layer is rebuilt to compute local gradients.
5. Each GPU gets the summed gradient of its slice and updates it.

• Each GPU stores one shard of the parameters, gradients and optimizer states at rest.
• An all-gather rebuilds one layer before it runs, and the forward pass frees the copy afterwards.
• The backward pass gathers the layer again, and a reduce-scatter leaves each GPU the summed gradient for its own shard.

Step 2: All-gather one layer before it runs

Before a layer runs, every GPU receives the missing shards of that layer's parameters, so each GPU briefly holds the whole layer. The PyTorch FSDP tutorial says sharded parameters are all-gathered into unsharded parameters before forward computation, and that implicit prefetching issues the all-gather for the next layer while the current one computes.

The memory cost is bounded, but it is not one layer. With prefetching, the layer that is computing and the next one are whole at the same moment; PyTorch's limit_all_gathers option caps forward at exactly those two, and in backward its default prefetch holds the current and next layers' parameters plus the current layer's gradients (PyTorch 2.14 documentation).

If a 70 billion parameter model had 80 equal layers, one layer in fp16 is about 1.75 GB, so that window is 3.5 GB in forward and about 5.25 GB in backward.¹ The window does not shrink when you add GPUs, while the shard does: across 1,024 GPUs the shard is about 1.09 GB, smaller than the window, so the gathered layers, not the shard, set the floor.¹

Step 3: Run the forward pass, then free the copy

Each GPU runs the layer on its own slice of the batch, then frees the gathered weights so only its shard remains. In FSDP terms this is resharding after forward. The PyTorch documentation (version 2.14) states that under its default FULL_SHARD strategy parameters are all-gathered before forward and backward and resharded after each.

Freeing is what turns a temporary copy into a saving. A variant called SHARD_GRAD_OP skips the free after forward, which saves the second gather and spends more memory. Think of it as a dial: the earlier the copy is released, the lower the peak memory and the more gathers the step must pay for. Each GPU repeats this gather, compute and free cycle for every layer in order, so on top of its shards it holds at most the layer that is computing and the one being prefetched.

Step 4: In the backward pass, gather again and compute local gradients

The backward pass needs the full weights of each layer again, so FSDP all-gathers the layer a second time, in reverse order. The PyTorch FSDP tutorial says parameters are all-gathered again before backward computation and that the gathers follow the reverse of the post-forward order.

Each GPU then computes gradients for the whole layer from its own slice of the batch. These are local gradients: full-sized and different on every GPU, because every GPU saw different data. Keeping them whole would bring back the memory problem, so the next step shrinks them immediately. Notice also that backward visits layers from last to first, which is why the gathers run in reverse and why prefetching has to guess the order correctly.

Step 5: Reduce-scatter the gradients and update your shard

After backward, a reduce-scatter adds the local gradients from all GPUs and hands each GPU only the summed gradient for its own shard (FSDP divides by the number of GPUs, so each shard ends with the average gradient, as in distributed data parallel (DDP)). The PyTorch FSDP tutorial says the local unsharded gradients are reduce-scattered into sharded gradients inside backward. Each GPU then updates only its own parameter shard with its own optimizer-state shard, and the step is complete.

This is where the cost appears. Plain data parallelism uses one all-reduce, which is a reduce-scatter plus an all-gather and so moves about 2 units of data per parameter (Rajbhandari et al., 2020, section 7.1). FSDP moves one unit for the forward gather, one for the backward gather and one for the reduce-scatter, which is 3 units, or 1.5 times as much (Rajbhandari et al., 2020, section 7).¹ The tutorial notes FSDP splits the all-reduce into a reduce-scatter and an all-gather, and the extra all-gather is the price of freeing the weights after forward.

Worked example

Sharding turns a replicated 120 GB bill into a 1.88 GB one, and the table shows the ladder for a 7.5 billion parameter model on 64 GPUs. Each row adds one more thing to shard; numbers follow the ZeRO paper's model, with activations left out.¹

Text description of the diagram

• A replicated 7.5B model needs 120 GB of model state per GPU.
• Sharding the optimizer states across 64 GPUs gives 31.4 GB, and sharding gradients as well gives 16.6 GB.
• Sharding parameters too gives 1.88 GB per GPU.

All four rows match the ZeRO paper's own figures for the same model and GPU count (Rajbhandari et al., 2020, section 5); the ladder factors are computed in the script.¹ A 70 billion parameter model is an illustrative second case: replicated it needs 1,120 GB of state per GPU, which no GPU holds, and fully sharded across 8 GPUs it needs 140 GB, still more than an 80 GB card such as the NVIDIA A100 80GB (vendor specification, as of October 2026). Across 64 GPUs it needs 17.5 GB, which leaves room for the gathered layers and activations on such a card.¹ That is why the number of GPUs, not only the strategy, decides whether a model trains.

The toy below is an excerpt of a numpy script of about sixty lines and shows the logic in numpy. Each rank holds a quarter of every layer; the loop gathers, computes local gradients on its data slice, reduce-scatters, and updates only its shard.

Xs, Ys = np.split(X, N), np.split(Y, N)                         # each rank's data slice
shards = [np.split(W.ravel().copy(), N) for W in layers]       # shards[layer][rank]
for step in range(steps):
    full = [np.concatenate(s).reshape(shape) for s, shape in zip(shards, shapes)]  # all-gather
    local = [grads(full, Xs[r], Ys[r]) for r in range(N)]          # each rank: its own data slice
    for l in range(len(layers)):
        total = sum(g[l].ravel() / N for g in local)             # reduce ...
        for r, piece in enumerate(np.split(total, N)):           # ... scatter
            shards[l][r] -= lr * piece                           # update own shard only

After 5 steps the sharded toy and an ordinary single-process run differ by at most 6e-17, which is rounding error, so sharding changes where the numbers live and not what they are.¹

Trade-offs and edge cases

The central trade is memory against communication: FULL_SHARD saves the most memory and adds the most traffic, so the right strategy depends on the network. The PyTorch documentation lists four public strategies, shown below (a fifth, _HYBRID_SHARD_ZERO2, applies SHARD_GRAD_OP inside a node), and notes that HYBRID_SHARD shards within a node and replicates across nodes, which keeps the expensive gather traffic on the fast links inside one machine and can help medium-sized models.

Text description of the diagram

• FULL_SHARD shards parameters, gradients and optimizer states.
• SHARD_GRAD_OP keeps parameters unsharded after the forward pass.
• HYBRID_SHARD shards inside a node and replicates across nodes.
• NO_SHARD replicates everything and all-reduces gradients.

Communication can hide behind computation. Prefetching the next layer's all-gather while the current layer computes overlaps the two, so the 1.5 times traffic does not always become 1.5 times the step time. Whether it does depends on layer size, network bandwidth and compute speed. The tutorial documents the limits: the first all-gather of a step is always exposed, and when the central processing unit (CPU) is slow to issue the next all-gather, computation waits (PyTorch FSDP tutorial). The PyTorch paper reports performance comparable to plain data parallelism while supporting much larger models (Zhao et al., 2023).

Activations are not sharded. FSDP shards parameters, gradients and optimizer states, while activations still grow with batch size and sequence length. A job that fits its state may still run out of memory on activations, which is a separate technique (activation checkpointing, which recomputes activations instead of storing them).

The 2025-26 state. The PyTorch FSDP tutorial now teaches fully_shard, the second version of FSDP, and its documentation asks users of the older FullyShardedDataParallel, the first version, to consider migrating (PyTorch 2.14 fully_shard documentation). The second version shards each parameter along its first dimension as a distributed tensor instead of flattening groups of tensors, exposes manual prefetch controls, and replaces the strategy names with one setting, reshard_after_forward: True behaves like FULL_SHARD, False like SHARD_GRAD_OP, and a two-dimensional device mesh gives HYBRID_SHARD (PyTorch FSDP tutorial, updated September 2025). The mechanism on this page is the same in both.

How to answer it in an interview

The spoken structure is: the problem in one number, the three shards, the two collectives, then the cost and what stays unsharded.

Common mistakes

"FSDP splits the model across GPUs, like model parallelism." (wrong) Each GPU still computes whole layers on its own slice of the batch; it only stores a shard between layers. (right) "FSDP is data parallelism with sharded state, and it gathers layers on demand."

"Sharding is free." (wrong) FSDP moves 3 units of data per parameter against 2 for plain data parallelism. (right) "It trades about 1.5 times the communication for memory divided by N."

"Each GPU only ever holds its shard." (wrong) A small window of layers is rebuilt in full on every GPU, so peak memory includes that window. (right) "Only a small window is whole at a time: the layer that is computing and the one being prefetched, plus that layer's gradients in backward."

"FSDP shards activations too." (wrong) Activations depend on batch size and sequence length and are handled by other techniques. (right) "It shards parameters, gradients and optimizer states; activations need checkpointing."

Follow-up questions

What does FSDP shard and when does it gather?

It shards parameters, gradients and optimizer states; it all-gathers a layer's parameters before forward and again before backward, and reduce-scatters gradients after backward. The variant SHARD_GRAD_OP skips the second gather by keeping weights after forward, and spends more memory instead.

What is the communication cost of ring all-reduce?

A ring all-reduce is a reduce-scatter plus an all-gather, so each GPU sends and receives about 2 units of data per parameter, which is why plain data parallelism costs 2 and FSDP 3 (Rajbhandari et al., 2020, section 7.1).

How do you implement FSDP from scratch?

Keep 1/N of every flattened layer per rank, concatenate the shards to gather, compute local gradients on each rank's data slice, split the summed gradient back into shards, and update only your own. The toy in the worked example does this in numpy and matches single-process training.

What is the difference between FSDP and DDP?

Distributed data parallel (DDP) replicates the whole model on every GPU and all-reduces gradients, while FSDP shards the state and uses reduce-scatter and all-gather instead. DDP is simpler and moves less data; FSDP fits models that DDP cannot. For very large jobs FSDP is usually combined with other forms of parallelism; the distributed training design page covers the combinations.

Is FSDP the same as ZeRO stage 3?

FSDP and ZeRO stage 3 implement the same idea: shard parameters, gradients and optimizer states. ZeRO's third stage shards all three, which the PyTorch FULL_SHARD strategy matches, while its earlier stages shard fewer things.

What changes when you add more GPUs but the batch stays the same?

Memory per GPU for state falls as 16 bytes divided by N, so a bigger model fits, but each GPU gets fewer examples and the communication share of each step grows. The shard keeps shrinking but the gathered window does not, so past a point it sets the per-GPU floor. When the model and its state already fit on one GPU, plain data parallelism is simpler and sends less data, and high-bandwidth memory explains why memory, not arithmetic, is often the limit.

Sources

1. NVIDIA A100 Tensor Core GPU — NVIDIA, 1 Oct 2026
2. ZeRO: Memory Optimizations Toward Training Trillion Parameter Models — arXiv (Rajbhandari, Rasley, Ruwase, He), v3, 13 May 2020
3. PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel — arXiv (Zhao et al.), v2, 12 Sep 2023
4. Getting Started with Fully Sharded Data Parallel (FSDP) — PyTorch tutorials (last updated), 2 Sep 2025
5. FullyShardedDataParallel and ShardingStrategy — PyTorch 2.14 documentation (date = read date), 24 Oct 2026
6. torch.distributed.fsdp.fully_shard — PyTorch 2.14 documentation (date = read date), 24 Oct 2026