Visual walkthrough — FSDP and sharded training

---

Step 1 — What is a parameter, and why does it cost bytes?

WHAT. A neural network is, at heart, a huge pile of numbers called parameters. When we say a model "has 175 billion parameters," we mean there are $175\times 10^9$ individual numbers that the training process is trying to tune. We call this count $P$.

WHY count bytes. A GPU has a fixed amount of memory (an A100 has 80 gigabytes). Every number we store eats a fixed number of bytes — the unit computers measure memory in. So the whole "can I even train this?" question reduces to one honest tally: how many bytes must live on one GPU at once? Everything below is that tally.

PICTURE. Look at the strip of little boxes. Each box is one parameter. The whole strip is $P$ boxes long. The number written under the strip — how many bytes one box weighs — is the quantity we are about to fight over.

---

Step 2 — Four things live per parameter, not one

WHAT. Training does not store just the parameter. For every single parameter it also keeps three companions:

$\underset{\text{the number itself}}{\underbrace{\theta }}\underset{\text{its gradient}}{\underbrace{{\nabla }_{\theta }\mathcal{L}}}\underset{\text{momentum}}{\underbrace{m}}\underset{\text{variance}}{\underbrace{v}}$

• $\theta$ — the parameter value ("the weight").
• $\mathcal{L}$ — the training loss: a single number measuring how wrong the model's predictions currently are. Training tries to make $\mathcal{L}$ small.
• ${\nabla }_{\theta }\mathcal{L}$ — the gradient of that loss: how much $\mathcal{L}$ would change if we nudged $\theta$. This tells us which way to update.
• $m$ and $v$ — the Adam optimizer's memory of past gradients (a running average $m$ and a running average-of-squares $v$). See ZeRO Optimizer for why Adam needs these.

WHY. People count "175B parameters" and imagine 175B numbers of memory. But the optimizer secretly quadruples the count. If we forget these companions, our memory estimate is off by 4× — the single most common budgeting mistake.

PICTURE. One parameter box now stacks into a column of four boxes. All four are stored in fp32, so each is $4$ bytes: parameter, gradient, momentum, variance.

---

Step 3 — The baseline: replicate everything (plain Data Parallelism)

WHAT. In ordinary data parallelism, every one of the $N$ GPUs (defined at the top) holds a complete photocopy of all four stacks. Multiply the per-parameter cost by $P$ parameters:

${\text{Mem}}_{\text{DP}}=16\cdot P\text{ bytes on }\text{every}\text{ GPU}$

• $16$ — bytes per parameter (Step 2).
• $P$ — number of parameters.
• The key sin: this number does not shrink when you add GPUs. Ten GPUs, ten identical photocopies.

WHY show this first. It is the villain. FSDP's whole reason to exist is that $16P$ is a constant in $N$ — throwing more GPUs at the problem buys you zero memory relief.

PICTURE. Four GPUs, each holding an identical full-height column. The columns are carbon copies — same colours, same order. Redundant.

---

Step 4 — Sharding: cut each strip into $N$ pieces

WHAT. Sharding means: take a strip of numbers and slice it into $N$ equal pieces, giving piece $i$ to GPU $i$. Nobody keeps the whole strip. If GPU $i$'s piece is called ${\theta }^{(i)}$, then:

$\theta =\left[{\theta }^{(0)}\right|{\theta }^{(1)}|\cdots |{\theta }^{(N-1)}],\text{each }{\theta }^{(i)}\text{ has }\frac{P}{N}\text{ entries}$

• The vertical bars mean "concatenated end to end."
• Superscript $(i)$ means "the piece owned by GPU $i$."
• Each piece is $\frac{1}{N}$ the length, so it costs $\frac{1}{N}$ the bytes.

WHY. This is the one idea that turns a constant into a shrinking quantity. FSDP shards all four stacks — $\theta$, gradient, $m$, and $v$ — the same way. That is the "3" in ZeRO Stage 3 (parent note: FSDP = ZeRO-3).

PICTURE. The full strip from Step 1, now sliced by dashed lines into $N=4$ coloured segments, each segment floating over the GPU that owns it. No GPU holds a full row.

---

Step 5 — The persistent memory after sharding

WHAT. Apply Step 4 to all four stacks. Each GPU now permanently stores only its slice of each:

${\text{Mem}}_{\text{persistent}}=\frac{16\cdot P}{N}\text{ bytes per GPU}$

• $16P$ — the total bytes of all four stacks across the whole cluster (unchanged; the numbers still exist somewhere).
• $÷N$ — but now spread across $N$ GPUs, so each holds only its share.

WHY. This is the payoff line. Where DP gave a flat $16P$ regardless of $N$, sharding gives $16P/N$ — more GPUs literally means less memory each. Doubling the cluster halves the per-GPU load.

PICTURE. Four GPUs, each holding one short column (its slice of all four types), sitting beside the tall greyed-out DP column for contrast. Same total, split four ways.

---

Step 6 — The catch: the temporary all-gather buffer

WHAT. A GPU cannot compute a layer $y=Wx$ with only $\frac{1}{N}$ of $W$ — matrix multiply needs all of $W$. So right before a layer runs, FSDP performs an all-gather: every GPU temporarily borrows the other shards to rebuild the full parameters of that one layer.

${\text{Mem}}_{\text{total}}=\underset{\text{persistent}}{\underbrace{\frac{16P}{N}}}+\underset{\text{one full layer, temporarily}}{\underbrace{L_{\max }}}$

• $\frac{16P}{N}$ — the permanent slices (Step 5).
• $L_{\max }$ — the byte-cost of the parameters (the $\theta$ values only, at $4$ bytes each) of the single largest layer, rebuilt in full and held only for the instant that layer computes, then discarded. Concretely, $L_{\max }=4\times (\text{parameters in the largest layer})$ — it is raw parameter bytes, not the full $16\times$ stack, because only $\theta$ is all-gathered for the forward compute.

WHY not the whole model? Because we rebuild one layer at a time and throw it away immediately. So the temporary cost scales with layer size, not model size — typically 1–2% of model memory for transformers.

PICTURE. A GPU's short persistent columns, plus one bright temporary block that swells up (all-gather arrow) while a layer computes, then vanishes (discard arrow) — a bubble that inflates and pops per layer.

---

Step 7 — The degenerate cases (where the formula bites back)

WHAT. Three edge cases the formula does not protect you from:

1. $N=1$ (one GPU). Then $\frac{16P}{N}=16P$ — sharding buys nothing, and you still pay the all-gather machinery. FSDP is pointless on a single device.
2. The model already fits. If $16P<$ GPU memory even without sharding, FSDP adds communication for savings you don't need — the parent's "When FSDP Hurts" example (15–20% slower).
3. Activations are NOT in the formula. The $16P/N$ term covers parameters, gradients, optimizer states — not activations. Activations cost $\text{batch}\times \text{seq}\times \text{hidden}$ per layer and are not sharded by vanilla FSDP.

WHY. A reader who trusts $16P/N$ blindly will still OOM (out-of-memory) on activations. This is the parent's second big mistake callout.

PICTURE. A bar chart: the sharded $16P/N$ bar looks tiny — then a separate, large "activations" bar towers next to it, unaffected by $N$, with an arrow to Activation Checkpointing as the fix.

---

The one-picture summary

Everything on one canvas: a full DP column ($16P$, flat) transforming — slice, distribute, add a small temporary bubble — into the FSDP layout ($16P/N+L_{\max }$), with the untouched activation block standing to the side as the honest caveat.

Recall Feynman retelling — say it back in plain words

Every weight in the model drags along three friends: its gradient, and Adam's two memories ($m$ and $v$). All four are fp32, so that's $4+4+4+4=16$ bytes per weight — and in plain data parallelism every GPU keeps a full photocopy of all of it, so adding GPUs never helps. FSDP's trick: slice each of those four piles into $N$ equal pieces (where $N$ is how many GPUs you have) and give one piece to each GPU, so the permanent cost drops to $16P/N$ — now more GPUs genuinely means less memory each. The price is that to actually compute a layer you need all its weights, so just before each layer runs, everyone borrows the missing pieces (all-gather), computes, then throws the borrowed pieces away — a small temporary bubble sized to one layer's parameters, not the whole model. Three traps: with one GPU it does nothing, if the model already fits it just adds overhead, and it never touches activation memory — those you kill separately with activation checkpointing.

Recall Quick self-check

Bytes per parameter in fp32 Adam training ::: $4+4+4+4=16$ Per-GPU persistent memory with FSDP over $N$ devices ::: $\frac{16P}{N}$ What does $N$ stand for ::: the number of data-parallel GPUs sharing the work Why does plain data parallelism not shrink with more GPUs ::: every GPU keeps a full photocopy, so per-GPU memory stays $16P$ regardless of $N$ What does the temporary buffer $L_{\max }$ scale with ::: the parameter bytes of one layer ($4\times$ params in the largest layer), not the whole model One thing FSDP does NOT shard ::: activations (fix with activation checkpointing)

---

Related: FSDP and sharded training · ZeRO Optimizer · Model Parallelism · Pipeline Parallelism · Mixed Precision Training · Communication Collectives