Fully sharded data parallelism — stop copying the whole training state

Yesterday we saw ordinary data parallelism:

\[ \boxed{\text{same model on every GPU, different data on every GPU}} \]

That is wonderfully simple and can scale training throughput well.

But it has a glaring memory problem.

If every GPU stores the full parameters, full gradients, and full AdamW state, then adding GPUs gives us more compute but does not let us train a model whose training state cannot fit on one GPU.

Today we will see how fully sharded data parallelism, often discussed alongside ZeRO, changes that.

The core idea is:

\[ \boxed{\text{replicate the computation logically, but shard the persistent state physically}} \]

Each GPU permanently owns only part of the model state. When a layer needs parameters that are elsewhere, the GPUs temporarily assemble them, do the computation, and then discard the copies.

This is a profound shift: GPU memory becomes a working set rather than a warehouse containing the entire model.

Start with the duplication problem¶

Suppose a model has \(P\) parameters.

Using our simplified AdamW accounting from Day 37:

BF16 parameters       2 bytes/parameter
BF16 gradients        2 bytes/parameter
FP32 first moment     4 bytes/parameter
FP32 second moment    4 bytes/parameter
-----------------------------------------
                     12 bytes/parameter

For a 70B-parameter model:

\[ 70\times10^9\times12 = 840\text{ GB}. \]

Ignore activations for the moment.

With ordinary data parallelism across 16 GPUs, each worker conceptually stores that same training state:

GPU 0 → 840 GB
GPU 1 → 840 GB
GPU 2 → 840 GB
...
GPU 15 → 840 GB

The cluster collectively contains many copies of the same information.

Worse, if a GPU has only 80 GB of HBM, the model does not fit on even one worker.

Sixteen GPUs do not help if each one independently needs 840 GB.

The obvious question¶

Why does every GPU need its own copy of everything?

For optimizer state, the answer is especially weak.

Suppose parameter \(w_7\) has Adam state:

\[ m_7,\qquad v_7. \]

If all workers ultimately perform the same synchronized optimization step, do all 16 GPUs really need permanent copies of:

\[ m_7,\quad v_7? \]

No.

One GPU could own them.

That is the beginning of sharded training.

ZeRO Stage 1: shard optimizer state¶

Imagine four GPUs.

Instead of:

GPU 0: all Adam states
GPU 1: all Adam states
GPU 2: all Adam states
GPU 3: all Adam states

we partition them:

GPU 0: Adam state for parameters 0–24%
GPU 1: Adam state for parameters 25–49%
GPU 2: Adam state for parameters 50–74%
GPU 3: Adam state for parameters 75–100%

Each GPU updates the parameters for the optimizer-state shard it owns.

The updated parameter pieces can then be shared so all workers remain synchronized.

If optimizer state dominated memory, this is a large win.

With \(N\) data-parallel workers, optimizer-state memory per GPU falls roughly toward:

\[ \frac{M_{\text{optimizer}}}{N}. \]

But parameters and gradients are still replicated.

ZeRO Stage 2: shard gradients too¶

Next ask:

Why should every GPU permanently hold a complete gradient tensor?

During backward, every worker computes gradient contributions from its own data.

But after those contributions are combined, we can arrange for each GPU to retain only the gradient shard corresponding to the parameters it owns for optimization.

Conceptually:

GPU 0 → gradient shard 0
GPU 1 → gradient shard 1
GPU 2 → gradient shard 2
GPU 3 → gradient shard 3

Now both:

optimizer state
gradients

are sharded.

This further reduces memory duplication.

But the model parameters themselves are still fully replicated.

ZeRO Stage 3 / FSDP: shard the parameters¶

Now take the final step.

Instead of every GPU permanently storing:

\[ \theta, \]

partition the parameters too:

GPU 0 → θ0
GPU 1 → θ1
GPU 2 → θ2
GPU 3 → θ3

Each GPU stores only about:

\[ \frac{P}{N} \]

parameters persistently.

This is the essential idea behind ZeRO Stage 3 and fully sharded data parallel approaches.

At first this seems impossible.

How can GPU 0 execute a Transformer layer if it owns only one quarter of that layer's weights?

The answer is:

Temporarily gather the weights when the layer needs them.

The key operation: all-gather¶

Suppose a layer's weight matrix is partitioned:

GPU 0: W0
GPU 1: W1
GPU 2: W2
GPU 3: W3

Before computing that layer, the GPUs perform an all-gather.

Afterward, each participating GPU temporarily has:

W = [W0 W1 W2 W3]

and can run the layer on its own local training examples.

So the timeline becomes:

persistent memory:

GPU 0: W0
GPU 1: W1
GPU 2: W2
GPU 3: W3

        ↓ all-gather

temporary working memory:

GPU 0: full W
GPU 1: full W
GPU 2: full W
GPU 3: full W

        ↓ compute layer

discard non-owned pieces

        ↓

persistent memory again:

GPU 0: W0
GPU 1: W1
GPU 2: W2
GPU 3: W3

The full layer exists everywhere only briefly.

The library mental model¶

Imagine four libraries jointly own a giant encyclopedia.

Instead of every library permanently storing all 100 volumes:

Library A → volumes 1–100
Library B → volumes 1–100
Library C → volumes 1–100
Library D → volumes 1–100

they divide ownership:

Library A → volumes 1–25
Library B → volumes 26–50
Library C → volumes 51–75
Library D → volumes 76–100

Suppose every library runs a class that needs volumes 40–43 for ten minutes.

Those volumes are temporarily copied to the libraries that need them.

After class:

temporary copies → discarded

The permanent collection remains distributed.

FSDP treats model parameters similarly.

The cluster collectively owns the whole model.

Each GPU only keeps its assigned shelves permanently.

Backward needs communication too¶

Forward is only half the story.

During backward, each GPU computes gradient contributions from its local examples.

Suppose the full gradient for a layer is:

\[ g. \]

We do not want every GPU to retain all of \(g\).

Instead, a reduce-scatter operation can both combine gradients across workers and distribute the resulting pieces.

Conceptually, before:

GPU 0 → local full gradient g0
GPU 1 → local full gradient g1
GPU 2 → local full gradient g2
GPU 3 → local full gradient g3

We want the global sum:

\[ g_{\text{sum}} = g_0+g_1+g_2+g_3. \]

But rather than giving that entire tensor to every GPU, reduce-scatter leaves:

GPU 0 → summed gradient shard A
GPU 1 → summed gradient shard B
GPU 2 → summed gradient shard C
GPU 3 → summed gradient shard D

So one collective operation performs two conceptual jobs:

reduce
→ combine gradient contributions

scatter
→ leave each GPU only its assigned piece

This fits perfectly with sharded optimizer state.

The whole FSDP rhythm¶

For one layer, a simplified training flow is:

1. each GPU owns only a parameter shard

2. all-gather layer parameters
        ↓
3. run forward on local data
        ↓
4. discard unnecessary full parameter copy

... later during backward ...

5. all-gather parameters if needed again
        ↓
6. compute local gradients
        ↓
7. reduce-scatter gradients
        ↓
8. each GPU keeps only its gradient shard
        ↓
9. update locally owned parameters with locally owned Adam state

Then repeat.

This is why fully sharded training can save enormous amounts of memory while increasing communication.

A simplified memory calculation¶

Return to our 70B model with:

\[ 840\text{ GB} \]

of simplified persistent training state.

Suppose we shard that state evenly across:

\[ N=16 \]

GPUs.

The idealized persistent state per GPU becomes:

\[ \frac{840}{16} = 52.5\text{ GB}. \]

Now a configuration that was impossible on an 80 GB GPU becomes at least plausible from the persistent-state perspective.

We still need room for:

activations
temporary gathered parameters
communication buffers
framework overhead

so 52.5 GB does not mean the full training run automatically fits.

But compare:

\[ 840\text{ GB/GPU} \]

with:

\[ 52.5\text{ GB/GPU}. \]

That is the power of removing replication.

But sharding does not make information disappear¶

This is crucial.

The model still contains 70B parameters.

The optimizer still needs its moments.

We did not compress the mathematical state.

We changed where it lives.

Across the cluster, the information still exists:

\[ M_{\text{cluster}} \approx M_{\text{full training state}}. \]

Per GPU:

\[ M_{\text{GPU}} \approx \frac{M_{\text{sharded state}}}{N} + M_{\text{working set}}. \]

The second term matters.

A GPU temporarily needs full parameter groups while computing them.

So FSDP is not:

\[ \text{memory}=\frac{\text{everything}}{N}. \]

It is closer to:

\[ \boxed{ M_{\text{GPU}} \approx \frac{M_{\text{persistent state}}}{N} + M_{\text{largest temporary materialized state}} + M_{\text{activations}} } \]

This is a much better mental model.

Shard size versus communication frequency¶

Suppose we gather parameters one layer at a time.

That keeps temporary memory relatively small:

gather layer 1
compute
discard

gather layer 2
compute
discard

gather layer 3
compute
discard

But it creates many communication operations.

Alternatively, gather several layers together:

gather layers 1–4
compute them
discard

Now communication may be more efficient, but temporary memory is larger.

So we again face a tradeoff:

\[ \boxed{\text{larger materialization groups} \rightarrow \text{more memory, fewer/larger communications}} \]

versus:

\[ \boxed{\text{smaller groups} \rightarrow \text{less memory, more communication overhead}} \]

Large-scale training is full of these memory-bandwidth tradeoffs.

Prefetching can hide parameter communication¶

Suppose we are computing layer 20.

We know layer 21 will be needed next.

Instead of waiting:

compute layer 20
        ↓
wait for layer 21 weights
        ↓
compute layer 21

the system can begin gathering layer 21 while layer 20 is still computing:

time →

compute:
[layer 20 compute][layer 21 compute][layer 22 compute]

network:
      [gather layer 21]
                         [gather layer 22]

This is parameter prefetching.

The same principle appeared yesterday with gradient all-reduce:

\[ \boxed{\text{hide communication behind useful computation}} \]

If the network finishes before the next layer needs its weights, much of the communication latency disappears from the critical path.

Why topology suddenly matters¶

Suppose eight GPUs inside one server are connected by a very fast interconnect.

But communication between servers is slower.

Then gathering a parameter shard from:

GPU 1 in the same server

may be much cheaper than gathering it from:

GPU 57 in another rack.

At scale, distributed training must understand a hierarchy such as:

GPU
↓
local high-speed GPU fabric
↓
server
↓
network interface
↓
rack
↓
cluster fabric
↓
other servers

Not every byte has the same movement cost.

This is one reason parallelism strategies are often mapped deliberately onto physical topology.

FSDP versus tensor parallelism¶

At this point FSDP may sound like model parallelism.

It is, in a physical-storage sense, but its execution pattern differs from tensor parallelism.

With FSDP:

weights are sharded while idle
↓
full layer weights are temporarily reconstructed
↓
each data-parallel worker performs the full layer computation on its local data

With tensor parallelism:

the layer's mathematical computation itself is split across GPUs

For example, a matrix multiplication:

\[ Y=XW \]

might be partitioned so different GPUs compute different columns of \(Y\).

The full weight matrix may never be materialized on one GPU.

So a useful distinction is:

\[ \boxed{\text{FSDP shards storage; tensor parallelism shards computation}} \]

That distinction is not absolute in every implementation, but it is an excellent mental model.

When FSDP alone stops being enough¶

Suppose a single Transformer layer has an enormous parameter matrix.

Even if its persistent shards fit easily, temporarily gathering the whole layer may exceed GPU memory.

Or perhaps the gathered layer fits, but its matrix multiplication is too large to execute efficiently on one GPU.

Then we need to split the layer computation itself.

That motivates tensor parallelism.

Similarly, if a model has hundreds of layers, we might place different layer ranges on different groups of GPUs.

That motivates pipeline parallelism.

Large training runs therefore combine dimensions:

data parallelism
×
fully sharded state
×
tensor parallelism
×
pipeline parallelism
×
gradient accumulation

The resulting configuration is sometimes described as multi-dimensional parallelism.

Why FSDP can increase network traffic dramatically¶

Ordinary data parallelism mainly synchronizes gradients.

FSDP must additionally move parameters.

During forward:

\[ \text{all-gather parameters}. \]

During backward, often again:

\[ \text{all-gather parameters}. \]

Then:

\[ \text{reduce-scatter gradients}. \]

So we saved HBM by increasing network dependence.

This continues the theme from activation checkpointing.

There we traded:

\[ \text{memory} \leftrightarrow \text{compute}. \]

Here we trade:

\[ \boxed{\text{memory} \leftrightarrow \text{communication}} \]

The model fits because we stopped storing everything locally.

But now the network must deliver missing pieces quickly enough.

A useful three-resource triangle¶

At this point, LLM training can be understood as balancing three scarce resources:

             GPU compute
                /\
               /  \
              /    \
             /      \
            /        \
           /__________\
      HBM memory    network

Different techniques move pressure around this triangle.

Activation checkpointing:

less HBM
more compute

FSDP:

less HBM
more network

Data parallelism:

more aggregate compute
more network synchronization

Gradient accumulation:

smaller instantaneous activation memory
more sequential work per optimizer step

There is rarely a free optimization.

We mostly exchange one resource for another.

The deeper mental model: HBM as a cache for the cluster¶

Yesterday, each data-parallel GPU behaved like a self-contained copy of the model.

With FSDP, that picture changes.

Think of the cluster as owning the model.

Each GPU's HBM contains:

its permanent shard
+
the parameter pieces needed right now
+
the activations needed right now

In that sense, GPU memory behaves partly like a cache.

A layer is approaching:

fetch its missing weights

The layer executes:

use them

The layer is finished:

evict them

Then fetch the next layer.

The model is no longer something that must exist in full on every accelerator.

It is a distributed object whose relevant pieces become locally materialized just in time.

Put the sequence together¶

The last several lessons now form a coherent progression.

AdamW showed why training state is much larger than inference weights.

Gradient accumulation said:

The entire global batch does not need to exist simultaneously.

Activation checkpointing said:

All forward activations do not need to exist simultaneously.

Ordinary data parallelism said:

Different parts of the batch can be computed simultaneously on different GPUs.

FSDP now says:

The entire persistent model state does not need to exist on every GPU.

So we have progressively removed the assumption that large mathematical objects must be physically colocated.

The global batch can be distributed across time.

Activations can be distributed across time.

Optimizer state can be distributed across GPUs.

Parameters can be distributed across GPUs and materialized only when needed.

That gives us a broader principle:

\[ \boxed{ \text{large-scale training is the art of controlling when and where tensors exist} } \]

A naive implementation asks:

What tensors does the mathematics require?

A scalable implementation asks a more powerful question:

Which of those tensors must exist on this GPU, at this exact moment?

That question is one of the keys to understanding modern LLM training systems.