Worked examples — FSDP and sharded training

This page is the practice arena for FSDP and sharded training. We do not re-derive theory here — we throw every kind of situation at FSDP and grind through the numbers. If a symbol shows up you haven't met, it was built in the parent note; we recall it in one line the first time it appears.

Two symbols we reuse constantly:

• $P$ = number of parameters in the model (a plain count, e.g. $175\times 10^9$).
• $N$ = number of data-parallel devices sharing the work (GPUs, e.g. $64$).

---

The scenario matrix

Every FSDP exam/interview question lands in one of these cells. The worked examples below are each labelled with the cell they cover, so together they touch every box.

#	Cell class	The question being asked	Degenerate / edge behaviour to watch
A	Memory, fp32	Persistent bytes per GPU, no mixed precision	$N=1$ → FSDP = plain DP (no saving)
B	Memory, mixed precision	Bytes per param when fp16 + fp32 master weights coexist	People forget the fp32 master copy
C	Activation blind spot	Total memory including activations (which FSDP does not shard)	Activations can dwarf the sharded part
D	Communication volume	Bytes each device sends/receives in all-gather & reduce-scatter	$N\to \infty$ → cost per device $\to$ full tensor size
E	Comm-vs-compute trade-off	Is FSDP faster or slower than plain DP here?	Small model on fast link → FSDP loses
F	Slow-interconnect limit	Wall-clock cost of a layer's all-gather on a real network	Bandwidth, not GPU, is the bottleneck
G	ZeRO stage picker	Which ZeRO stage (1/2/3) gives just-enough memory?	Over-sharding wastes bandwidth
H	3-D parallelism bookkeeping	Split $N_{\text{data}}\times N_{\text{pipeline}}\times N_{\text{tensor}}$ over a fixed GPU count	Product must equal total GPUs exactly

Prerequisite links used along the way: Mixed Precision Training, Activation Checkpointing, Communication Collectives, ZeRO Optimizer, Data Paralelism, Pipeline Parallelism, Model Parallelism, Gradient Accumulation.

Reading order note: the examples are numbered by matrix cell, so Example 5 = Cell E and Example 6 = Cell F. If you jump straight to a cell, match on the letter, not the position.

---

Example 1 — Cell A: pure fp32 memory, and the $N=1$ edge

Forecast: Guess before reading — how many bytes per parameter does fp32 + Adam need? Write your number down.

We name one quantity so the formulas read cleanly:

• $M$ = persistent memory per GPU, in bytes — the buffers a device must hold for the whole training run (parameters, gradients, optimizer states), not counting temporary all-gather scratch.

1. Count the bytes each parameter drags along. Why this step? Memory is dominated by four buffers, and each parameter needs a fixed number of bytes in each. In fp32 every number is $4$ bytes.

   • parameter $\theta$: $4$ B
   • gradient ${\nabla }_{\theta }\mathcal{L}$: $4$ B
   • Adam momentum $m$: $4$ B
   • Adam variance $v$: $4$ B

   Total $=4+4+4+4=16$ bytes per parameter.

   Recall Why not

   $24$ like the parent note? The parent note stored momentum/variance in fp64 ($8$ B each) → $24$ B. Here we keep Adam states in fp32 ($4$ B each) → $16$ B. Byte counts depend on the precision you assume — always state it. We use $16$ B (all-fp32) throughout this example. Example 5 will use a different $24$ B count — watch for its own derivation there.

2. No-FSDP memory (full replica on every device). Why this step? Plain Data Paralelism copies all four buffers onto every GPU, so $M$ here is the whole model's byte count. Write $M_{\text{full}}$ for this "full replica" value. $M_{\text{full}}=16\times 7\times 10^9=1.12\times 10^{11}\text{bytes}=112\text{GB}.$ Does not fit on 80 GB. ❌

3. FSDP with $N=1$ — the degenerate case. Why this step? This is the edge the matrix flagged. FSDP divides that persistent memory by $N$; write $M_{\text{FSDP}}(N)$ for the sharded value. Sharding "across one device" gives that device everything. $M_{\text{FSDP}}(N=1)=\frac{M_{\text{full}}}{1}=\frac{112}{1}=112\text{GB}.$ Identical to no-FSDP. FSDP at $N=1$ buys nothing — sharding needs at least two workers to divide anything.

4. FSDP with $N=8$. Why this step? Now the persistent buffers split $8$ ways. $M_{\text{FSDP}}(N=8)=\frac{M_{\text{full}}}{8}=\frac{112}{8}=14\text{GB (persistent)}.$ Fits comfortably (plus a small temporary all-gather buffer). ✅

Verify: Sanity check the $N=1$ limit: $M_{\text{full}}/N$ with $N=1$ must equal the full-replica number $112$ — it does, confirming the formula degrades gracefully. Units: bytes/param × params = bytes → GB. ✔

---

Example 2 — Cell B: mixed precision, the fp32 master trap

Forecast: Is mixed precision more or less bytes per param than pure fp32? (Trap: it's often more, because you keep two copies of the weights.)

1. List every buffer and its precision. Why this step? Mixed precision doesn't delete buffers, it changes their width — and adds a master copy.

   • fp16 parameter: $2$ B
   • fp16 gradient: $2$ B
   • fp32 master parameter (for stable updates): $4$ B
   • fp32 Adam momentum: $4$ B
   • fp32 Adam variance: $4$ B

2. Sum. $2+2+4+4+4=16\text{bytes per parameter}.$ Same total as pure fp32 here — the fp16 savings on params+grads are exactly eaten by the extra fp32 master. This is the trap: mixed precision saves activation/compute memory, not optimizer memory.

3. Persistent memory per GPU at $N=64$. Why this step? All five buffers are sharded by FSDP (ZeRO-3), so persistent memory $M$ is the per-param bytes times $P$, divided by $N$. $M=\frac{16\times 7\times 10^9}{64}=\frac{1.12\times 10^{11}}{64}=1.75\times 10^9\text{B}=1.75\text{GB}.$

Verify: $16\times 7\times 10^9=112$ GB total; $112/64=1.75$ GB. Cross-check against Example 1's $N=8$ result: $112/8=14$, and $14/8=1.75$ — halving-by-8 twice matches. ✔

---

Example 3 — Cell C: the activation blind spot

Forecast: Which is bigger — the sharded weights you cleverly split, or the activations you forgot?

1. Activation tensor size for one layer. Why this step? Backprop must remember each layer's output, shaped $b\times s\times d$. $b\cdot s\cdot d=8\times 2048\times 12288=2.013\times 10^8\text{elements}.$

2. Convert to bytes (fp16). $2.013\times 10^8\times 2=4.026\times 10^8\text{B}\approx 0.40\text{GB per layer}.$

3. Compare to sharded persistent memory. Why this step? This is the whole point — activations are not divided by $N$.

   • Sharded persistent (all layers, per GPU): $\approx 4.1$ GB.
   • Activations for a $96$-layer model without checkpointing: $96\times 0.40\approx 38.4$ GB — nearly ten times the sharded weights.

4. Fix. Why this step? Activation Checkpointing stores activations only at layer boundaries and recomputes the rest. Let $L$ = the number of layers in the model (here $L=96$). Instead of keeping all $L$ layers' activations, checkpointing keeps roughly $\sqrt{L}$ of them, recomputing the rest during the backward pass — trading ~33% extra compute for a ~10× activation-memory cut.

Verify: $8\cdot 2048\cdot 12288=201,326,592$ elements; $\times 2$ B $=402,653,184$ B $\approx 0.402$ GB. Times $L=96$ layers $\approx 38.6$ GB $>$ $4.1$ GB persistent. The blind spot dominates. ✔

---

Example 4 — Cell D: communication volume, and the $N\to \infty$ limit

Forecast: As you add more GPUs, does each GPU send more or less per collective?

The figure below plots the per-device all-gather cost as $N$ grows; we walk it in step 2. Look at the blue curve rising toward the red dashed ceiling, and the two orange dots marking the numbers we compute.

Figure: per-device all-gather traffic $\frac{N-1}{N}S$ versus number of devices $N$. Blue = the cost curve; red dashed = the limit $S=700$ MB; orange dots = the worked values at $N=4$ ($525$ MB) and $N=64$ ($689.06$ MB). The curve never crosses the ceiling — you can never receive more than the whole tensor.

1. All-gather cost per device. Why this step? Each device already holds its own $1/N$ shard; it must receive the other $N-1$ shards to rebuild the full tensor. From Communication Collectives: $\text{cost}=\frac{N-1}{N}S.$

2. Plug in (this is the curve in the figure).

   • $N=4$: $\frac{3}{4}\times 700=525$ MB received per device — the left orange dot.
   • $N=64$: $\frac{63}{64}\times 700=689.06$ MB — the right orange dot, already almost touching the ceiling.
   • $N\to \infty$: $\frac{N-1}{N}\to 1$, so cost $\to S=700$ MB (the red dashed line). Each device essentially receives the whole tensor — you never pay more than the full size, no matter how many GPUs.

3. Reduce-scatter cost. Why this step? Backward pass sums each shard's gradient across devices and hands device $i$ only its slice. Same volume as all-gather by symmetry: $\text{cost}=\frac{N-1}{N}S.$

4. Per-layer total — one all-gather + one reduce-scatter. Why this step? In standard FSDP the layer's parameters are already resident in fp16 from the forward all-gather when the backward reaches it (or are re-gathered once), so a well-implemented step pays one all-gather (to materialise the full weights) and one reduce-scatter (to shard the summed gradients) per layer: $\text{per-layer volume}=\underset{\text{all-gather}}{\underbrace{\frac{N-1}{N}S}}+\underset{\text{reduce-scatter}}{\underbrace{\frac{N-1}{N}S}}=2\cdot \frac{N-1}{N}S.$

Verify: At $N=64$: $\frac{63}{64}\times 700=689.0625$ MB. Per-layer total $=2\times 689.0625=1378.125$ MB. Limit check: ${\lim }_{N\to \infty }\frac{N-1}{N}=1$, matching $700$ MB. ✔

---

Example 5 — Cell E: is FSDP worth it? (small model, fast link — FSDP loses)

Forecast: Memory is fine either way (we'll confirm). Does FSDP still help — or just slow you down?

1. First check: does plain DP already fit? Why this step? FSDP only pays off when memory overflows, so we must first size the plain-DP footprint. This scenario matches the parent note's mixed-precision GPT-3 accounting: per parameter you keep fp16 param ($2$ B) + fp16 grad ($2$ B) + fp32 master ($4$ B) + fp32 momentum ($8$ B) + fp32 variance ($8$ B) $=24$ B/param (the parent note uses fp64-width $8$ B Adam states here).

   Recall Why

   $24$ here but $16$ in Examples 1–2? The byte count is your choice of precision. Examples 1–2 kept Adam states at fp32 width ($4$ B each) → $16$ B. This example follows the parent note's convention with $8$ B-wide momentum/variance → $24$ B. Always state the width before you multiply.

   Plain-DP memory per GPU (full replica, not sharded): $24\times 1.3\times 10^9=3.12\times 10^{10}\text{B}=31.2\text{GB}<80\text{GB}.✅$ The model already fits without sharding — so FSDP's memory benefit is zero here.

2. Compute the comm/compute ratio. Why this step? The parent note's rule of thumb: if $\frac{T_{\text{comm}}}{T_{\text{comp}}}>0.3$, FSDP likely hurts. $r=\frac{T_{\text{comm}}}{T_{\text{comp}}}=\frac{3}{8}=\mathrm{0.375.}$

3. Apply the rule. $0.375>0.30$ → FSDP overhead is not amortized. Since memory already fits without sharding, you are paying communication for a benefit you don't need.

4. Estimate the slowdown — worst case vs. overlapped case. Why this step? The raw number and the parent note's "$15$–$20\%$" differ because of overlap, so we must quantify both.

   • Worst case (no overlap): communication runs after compute, so the layer takes $T_{\text{comp}}+T_{\text{comm}}=8+3=11$ ms. Slowdown $=\frac{11-8}{8}=37.5\%$.
   • Overlapped case: FSDP prefetches the next layer's all-gather while the current layer's matmul is still running (compute and communication happen at the same time). Only the part of $T_{\text{comm}}$ that cannot hide behind compute is charged. If a fraction $\varphi$ of the $3$ ms overlaps, the exposed cost is $(1-\varphi )\times 3$ ms. To land at the parent note's $\approx 8\times 0.175=1.4$ ms exposed (a $17.5\%$ slowdown), we need $\varphi \approx 1-\frac{1.4}{3}\approx 0.53$ — i.e. about half the communication hidden. That is a realistic overlap fraction on NVLink, which is exactly why the measured slowdown ($15$–$20\%$) is smaller than the naive $37.5\%$.

5. Decision. Use plain Data Paralelism here; reserve FSDP for when memory actually overflows.

Verify: $24\times 1.3\times 10^9=31.2$ GB $<80$; $3/8=0.375>0.3$; worst-case slowdown $(11-8)/8=37.5\%$; overlap fraction to hit $17.5\%$: $1-1.4/3\approx 0.533$. ✔

---

Example 6 — Cell F: slow-interconnect wall-clock, bandwidth-bound

Forecast: $10$ Gbps sounds fast. Guess the milliseconds to move $700$ MB.

1. Unit-align bandwidth. Why this step? Bits vs bytes is the classic slip. $10$ Gbps $=10\times 10^9$ bits per second $=1.25\times 10^9$ bytes per second ($÷8$).

2. Transfer time = data ÷ bandwidth. Why this step? Time on a bandwidth-bound link is just "how many bytes" divided by "bytes per second" — the same idea as "distance ÷ speed = time". Both numbers must be in bytes, which is why step 1 converted Gbps to bytes/s first. $t=\frac{700\times 10^6\text{B}}{1.25\times 10^9\text{B/s}}=0.56\text{s}=560\text{ms}.$

3. Interpret and conclude. Why this step? We now compare that $560$ ms to compute. A transformer layer's matrix multiplies on an A100 finish in a handful of milliseconds. So $560$ ms of communication per layer $\gg$ compute → the GPU sits idle waiting for bytes. Conclusion: on this slow link FSDP is bandwidth-bound, and its memory win is drowned by communication — exactly the parent note's warning about PCIe / cross-datacenter links.

Verify: $700\times 10^6/1.25\times 10^9=0.56$ s $=560$ ms. ✔

---

Example 7 — Cell G: pick the smallest ZeRO stage that fits

Forecast: You don't always need full FSDP. Guess the cheapest stage that clears $10$ GB.

Break the $16$ B/param into what each stage shards. Buffers: fp16 param $2$, fp16 grad $2$, fp32 master $4$, momentum $4$, variance $4$. Optimizer states = master + momentum + variance = $12$ B; gradients = $2$ B; params (fp16) = $2$ B.

1. ZeRO-1 (shard the $12$ B optimizer states only). Why this step? Stage 1 leaves params and grads fully replicated (the $2+2=4$ B stay whole) and divides only the $12$ B of optimizer states by $N$. $M_1=\underset{\text{full params+grads}}{\underbrace{(2+2)}}+\frac{12}{8}=4+1.5=5.5\text{B/param}.$ $M_1\times 3\times 10^9=16.5$ GB. ❌ (>10)

2. ZeRO-2 (shard optimizer states + grads = $12+2=14$ B). Why this step? Stage 2 additionally shards the $2$ B gradient, leaving only the $2$ B fp16 param replicated. $M_2=2+\frac{14}{8}=2+1.75=3.75\text{B/param}.$ $\times 3\times 10^9=11.25$ GB. ❌ (still >10, barely)

3. ZeRO-3 / FSDP (shard all $16$ B). Why this step? Stage 3 shards everything — nothing stays replicated, so the whole $16$ B/param divides by $N$. $M_3=\frac{16}{8}=2\text{B/param}.$ $\times 3\times 10^9=6$ GB. ✅ Fits under $10$ GB.

4. Decision. Only ZeRO-3 (full FSDP) clears the budget here. Choosing the smallest stage that fits minimises communication (parent note: higher stages = more collectives).

Verify: $M_1=5.5\Rightarrow 16.5$ GB; $M_2=3.75\Rightarrow 11.25$ GB; $M_3=2\Rightarrow 6$ GB. Monotonic decrease, only stage 3 $\le 10$. ✔

---

Example 8 — Cell H: 3-D parallelism must multiply to your GPU count

Forecast: Given the constraint $N_{\text{data}}\times N_{\text{pipeline}}\times N_{\text{tensor}}=\text{Total GPUs}$, solve for the missing factor.

1. Write the constraint. Why this step? Every GPU belongs to exactly one $(\text{data},\text{pipeline},\text{tensor})$ coordinate, so the three counts must multiply to the total. $N_{\text{data}}\times N_{\text{pipeline}}\times N_{\text{tensor}}=1024.$

2. Solve for $N_{\text{data}}$. Why this step? Two of the three factors are fixed by the hardware plan, so the data-parallel width is whatever makes the product exact. $N_{\text{data}}=\frac{1024}{N_{\text{pipeline}}\times N_{\text{tensor}}}=\frac{1024}{4\times 8}=\frac{1024}{32}=32.$

3. Consistency check for the "tensor inside a node" rule. Why this step? Tensor parallelism needs the fastest link, so its group must fit within one $8$-GPU node — and indeed $N_{\text{tensor}}=8$ equals the node size exactly. ✅ FSDP's all-gathers then run across the $32$ data-parallel replicas, its heaviest traffic kept off the critical tensor-parallel path.

4. Sanity-multiply back. Why this step? Never trust a division without re-multiplying: $32\times 4\times 8=1024$, matching the GPU budget exactly, so no device is left unassigned or double-counted.

Verify: $32\times 4\times 8=1024$. And $N_{\text{tensor}}=8\le$ node size $8$. ✔

---

Recall Quick self-test

All-gather cost per device for tensor size $S$ over $N$ devices ::: $\frac{N-1}{N}S$ Per-layer FSDP communication volume (all-gather + reduce-scatter) ::: $2\cdot \frac{N-1}{N}S$ Does FSDP shard activations? ::: No — combine with activation checkpointing. At $N=1$, FSDP memory vs plain DP ::: Identical; sharding needs $N\ge 2$. Mixed-precision bytes/param with fp32 master + fp32 Adam ::: $2+2+4+4+4=16$ B. Rule of thumb: FSDP hurts when $T_{\text{comm}}/T_{\text{comp}}$ exceeds ::: $0.3$. 3-D parallelism constraint ::: $N_{\text{data}}\times N_{\text{pipeline}}\times N_{\text{tensor}}=$ total GPUs.