FSDP and sharded training: Exercises — AI-ML deep dive

Before we start, one shared vocabulary reminder, because every problem below leans on it.

We also fix a "bytes per parameter" table used throughout. This is the standard mixed-precision (Mixed Precision Training) count:

Level 1 — Recognition

Exercise 1.1

Which of these does FSDP shard across GPUs, and which does it not? (a) parameters, (b) activations, (c) optimizer momentum, (d) gradients.

Recall Solution

Sharded: (a) parameters, (c) optimizer momentum, (d) gradients. Not sharded: (b) activations — every GPU still stores its own full activation tensors for backprop. This is the single most-missed fact and the seed of the L5 problem.

Exercise 1.2

Match the collective to its job:

All-gather :: ?
Reduce-scatter :: ?

Recall Solution

All-gather ::: each GPU holds one shard $x^{(i)}$ ; afterwards every GPU holds the full concatenation $[x^{(0)}, \dots, x^{(N - 1)}]$ . Used to rebuild a full weight matrix just before a layer's matmul. Reduce-scatter ::: each GPU holds a full vector $g$ ; afterwards GPU $i$ holds only the sum over all GPUs of the $i$ -th shard, $\sum_{j} g_{j}^{(i)}$ . Used to combine and split gradients in one step. See Communication Collectives.

Exercise 1.3

FSDP is PyTorch's version of which ZeRO stage, and what does that stage shard?

Recall Solution

FSDP = ZeRO Stage 3. It shards optimizer states and gradients and parameters — all three. (Stage 1 = optimizer only; Stage 2 = optimizer + gradients.) See ZeRO Optimizer.

Level 2 — Application

Exercise 2.1

A model has $P = 7 \times 1 0^{9}$ parameters. Using 16 bytes/param, compute the per-GPU persistent memory (a) with plain data parallelism (no sharding), and (b) with FSDP on $N = 8$ GPUs. Ignore the temporary buffer.

Recall Solution

(a) Data parallel replicates all 16 bytes on every GPU: $7 \times 1 0^{9} \times 16 = 112 \times 1 0^{9} bytes = 112 GB per GPU.$ (b) FSDP divides the persistent memory by $N = 8$ : $\frac{112 GB}{8} = 14 GB per GPU.$ So sharding across 8 GPUs turns a 112 GB requirement (impossible on an 80 GB card) into 14 GB (comfortable).

Exercise 2.2

For the same 7B model on $N = 8$ , the temporary all-gather buffer for the largest layer is 1.5 GB (fp16). What is the total per-GPU memory, and does it fit an 80 GB A100 (ignoring activations)?

Recall Solution

$sharded persistent 14 + temp all-gather 1.5 = 15.5 GB.$ Yes — $15.5 GB ≪ 80 GB$ , plenty of room (activations still to come, but ignored here).

Exercise 2.3

General formula check. Starting from "16 bytes/param, sharded by $N$ , plus a fixed temp buffer $T$ ", write the per-GPU memory $M (P, N, T)$ and evaluate it for $P = 175 \times 1 0^{9}$ , $N = 64$ , $T = 0$ .

Recall Solution

$M (P, N, T) = \frac{16 P}{N} + T .$ With $T = 0$ : $M = \frac{16 \times 175 \times 1 0 ^{9}}{64} = \frac{2.8 \times 1 0 ^{12}}{64} = 43.75 \times 1 0^{9} bytes \approx 44 GB.$ Comfortably under an A100 80 GB — before activations are added.

Level 3 — Analysis

Exercise 3.1

A layer holds 700 MB of parameters (fp16). You will all-gather it over a network. Compute the per-GPU received volume for $N = 8$ using the all-gather cost $\frac{N - 1}{N} \times size$ .

Recall Solution

Each GPU already owns its own $\frac{1}{N}$ shard, so it must receive the other $N - 1$ shards: $\frac{N - 1}{N} \times 700 = \frac{7}{8} \times 700 = 612.5 MB received per GPU.$ Notice: as $N \to \infty$ , $\frac{N - 1}{N} \to 1$ , so per-GPU received volume approaches the whole layer size — sharding more does not reduce all-gather traffic per GPU, it only reduces stored memory.

Definition Why the "

$\frac{T _{comm}}{T _{comp}} > 0.3$ " rule of thumb? FSDP overlaps (hides) communication behind computation: while a GPU crunches layer $l$ , it prefetches the weights of layer $l + 1$ in the background. If communication is shorter than computation, it hides completely and costs nothing. The ratio $\frac{T _{comm}}{T _{comp}}$ measures "how much of the compute window the comm needs."

Ratio $\leq 1$ in the ideal fully-overlapped case means comm can hide — but overlap is never perfect: kernel launch gaps, buffer allocation, and the last layer (nothing left to overlap with) leak through. Empirically only about the first ~70 % of the compute window gets usefully overlapped, so once comm exceeds roughly 0.3 of compute it starts poking out past the compute shadow and adding wall-clock time.
That is where the practical threshold 0.3 comes from: below it, communication is essentially free (hidden); above it, it becomes visible slowdown. It is a rule of thumb from measured overlap efficiency, not an exact law — treat "0.3" as "roughly a third."

Exercise 3.2

That same 700 MB layer is transferred over a 10 Gbps link. How long does one all-gather take (use full 700 MB for a worst-case single-hop estimate), and if a layer's compute takes 200 ms, is FSDP a good idea here? Also account for the backward pass, which adds a second all-gather plus a reduce-scatter.

Recall Solution

Convert: $10 Gbps = 10 \times 1 0^{9}$ bits/s $= 1.25 \times 1 0^{9}$ bytes/s $= 1250$ MB/s (decimal units). One all-gather: $t_{ag} = \frac{700 MB}{1250 MB/s} = 0.56 s = 560 ms.$ Why backward compute $\approx 2 \times$ forward. The forward pass does one matmul per layer ( $y = W x$ ). The backward pass does two matmuls per layer for the same weight: one for the gradient flowing to the inputs ( $\frac{\partial L}{\partial x} = W^{⊤} \frac{\partial L}{\partial y}$ ) and one for the gradient of the weights ( $\frac{\partial L}{\partial W} = \frac{\partial L}{\partial y} x^{⊤}$ ). Two matmuls of the same size ≈ twice the FLOPs ≈ twice the time. Hence forward compute 200 ms ⇒ backward compute $\approx 2 \times 200 = 400$ ms. Full per-layer communication. FSDP pays this collective three times per layer per step:

forward all-gather ≈ 560 ms,
backward all-gather (weights re-fetched for the gradient) ≈ 560 ms,
backward reduce-scatter of gradients — same 700 MB volume, so ≈ 560 ms as well.

$T_{comm} = 560 + 560 + 560 = 1680 ms.$ Compare to compute (forward + backward, $\approx 200 + 400 = 600$ ms): $\frac{T _{comm}}{T _{comp}} = \frac{1680}{600} = 2.8.$ (If you only counted the single forward all-gather against forward compute — $560/200 = 2.8$ — you get the same ratio here, which is why the earlier single-collective estimate already flagged the problem. But the honest full-step figure sums three collectives.) Either way the ratio $≫ 0.3$ (the overlap threshold defined just above): on a slow 10 Gbps link, do not use FSDP; the network is the bottleneck. (See figure below, which shows the single-collective comparison for clarity.)

Reading the figure. The horizontal axis lists two activities; the vertical axis is time in milliseconds. The blue bar (left) is one layer's forward compute = 200 ms; the pale-yellow bar (right) is one all-gather over the 10 Gbps link = 560 ms. The pink dashed line marks the 200 ms compute level so you can see how far the yellow comm bar overshoots it — the yellow arrow points to that overshoot, the "wasted GPU time" spent waiting on the wire that cannot be hidden behind compute. The bar values are printed above each bar. Remember: a real training step pays this collective three times (two all-gathers + one reduce-scatter), so the true comm burden is even larger than the single yellow bar shown.

Exercise 3.3

The 1.3B-model case from the parent note: it fits in memory under plain DP ( $1.3 \times 1 0^{9} \times 16 = 20.8$ GB on an 80 GB A100). Why does adding FSDP here slow you down, and what is the one-line decision rule?

Recall Solution

FSDP's memory saving is useless here — the model already fits — but its costs stay: every layer now pays an all-gather (forward), a second all-gather (backward), and a reduce-scatter (backward). You bought communication overhead and got no memory benefit in return, hence the parent's "15–20 % slower." Decision rule: use FSDP only when required memory $>$ available GPU memory. If it fits under plain DP, stay on DP.

Level 4 — Synthesis

Exercise 4.1

You have 512 GPUs and want $N_{data} = 8$ (FSDP), $N_{pipeline} = 8$ (Pipeline Parallelism), and some tensor parallelism (Model Parallelism). What is $N_{tensor}$ ?

Recall Solution

The parent's identity: $Total GPUs = N_{data} \times N_{pipeline} \times N_{tensor}$ . $512 = 8 \times 8 \times N_{tensor} \Rightarrow N_{tensor} = \frac{512}{64} = 8.$ A neat $8 \times 8 \times 8$ cube.

Exercise 4.2

In a hybrid FSDP + Pipeline setup, the model is split into 4 pipeline stages, each stage sharded by FSDP across 4 GPUs (16 GPUs total). A stage's parameters are 40 GB (16 bytes/param already folded in). What is the persistent per-GPU parameter memory, and why is the all-gather cheaper than pure FSDP over all 16?

Recall Solution

Each stage's 40 GB is sharded only within its 4 GPUs: $\frac{40 GB}{4} = 10 GB per GPU (persistent).$ Why cheaper communication: all-gather now happens inside a 4-GPU group, not across all 16. Fewer participants and (usually) intra-node NVLink instead of cross-node InfiniBand → smaller, faster collectives. Pipeline splits the model depth-wise so FSDP's expensive all-gather scope shrinks. Trade-off: pipeline introduces bubble/idle time you must hide with micro-batching.

Exercise 4.3

Design task: a 30B model, 8 nodes × 8 GPUs = 64 GPUs, NVLink inside a node (fast), slower InfiniBand across nodes. Where should tensor parallelism go, and where should FSDP go, and why?

Recall Solution

Put tensor parallelism inside a node ( $N_{tensor} = 8$ , over NVLink): tensor parallelism does frequent, latency-sensitive all-reduces every layer, so it needs the fastest link. Put FSDP across nodes ( $N_{data} = 8$ , over InfiniBand): FSDP's all-gather can be prefetched (overlapped with compute) and happens per-FSDP-unit, tolerating higher latency better. Result: $8 \times 8 = 64$ GPUs, fast link for the chatty collective, slower link for the overlappable one. This is exactly why the parent recommends "tensor parallelism for computation within a node, FSDP for memory across nodes."

Level 5 — Mastery

Exercise 5.1

GPT-3-scale layer: batch size $b = 8$ , sequence length $s = 2048$ , hidden dim $d = 12288$ , fp16 (2 bytes). Compute the activation memory for one layer's main tensor, and explain why FSDP does not reduce it. Then estimate the activation total across ~96 transformer layers.

Recall Solution

$b \times s \times d \times 2 = 8 \times 2048 \times 12288 \times 2 bytes .$ Step: $8 \times 2048 = 16384$ ; $16384 \times 12288 = 201, 326, 592$ ; $\times 2 = 402, 653, 184$ bytes $\approx 0.40 GB = 402 MB per layer$ (decimal units). Across ~96 layers (one such tensor each, no safety factor): $402 MB \times 96 \approx 38, 600 MB \approx 38.6 GB.$ (Real transformers store several activation tensors per layer — attention scores, MLP intermediates — so the practical number is a few times higher, which is why one reaches into the tens-of-GB range and why checkpointing matters.) Why FSDP does not reduce it: FSDP shards params, grads, optimizer — not activations. Activations are produced during the forward pass by each GPU for its own micro-batch, and each GPU must keep them until its own backward pass consumes them. There is no shard to fetch from a neighbour — the data is local, per-GPU, and unique. So this 38.6 GB figure is unchanged by any value of $N$ : doubling the GPU count halves the param/grad/optimizer footprint but leaves the activation footprint exactly where it was. That is precisely why FSDP alone cannot make an arbitrarily large model fit, and why you must pair it with Activation Checkpointing.

Exercise 5.2

Activation checkpointing recomputes activations in the backward pass instead of storing them, cutting activation memory ~10× at the cost of ~33 % more compute. If a layer's forward compute is 200 ms and a step is forward + backward, estimate the new backward compute (backward ≈ 2× forward without checkpointing) and the total step time with checkpointing.

Recall Solution

Baseline: forward $= 200$ ms, backward $\approx 2 \times 200 = 400$ ms (see Ex 3.2 for why backward is twice forward), step $= 600$ ms. Checkpointing adds a recompute of the forward inside backward: the extra $\approx 33%$ is applied to the step ⇒ extra $= 0.33 \times 600 = 198$ ms. $new step \approx 600 + 198 = 798 ms (\approx 800 ms) .$ You pay ~33 % more time to buy ~10× less activation memory — usually worth it when activations are what blocks you. Pair FSDP (params) with Activation Checkpointing (activations) to attack both memory sources.

Exercise 5.3

Full budget synthesis. GPT-3, 175B params, $N = 1024$ GPUs, mixed precision (16 bytes/param). Compute the persistent sharded memory. Then, with a temp all-gather buffer of $1.5$ GB and activations with checkpointing of $30$ GB, does the total fit an 80 GB A100? What is the largest single contributor, and what does that tell you about where to optimise next?

Recall Solution

Persistent sharded memory: $\frac{16 \times 175 \times 1 0 ^{9}}{1024} = \frac{2.8 \times 1 0 ^{12}}{1024} \approx 2.73 \times 1 0^{9} bytes \approx 2.73 GB.$ Total budget: $2.73 + 1.5 + 30 = 34.23 GB \approx 34 GB.$ $34 GB < 80 GB$ — fits, with comfortable headroom. Largest contributor: activations (30 GB), dwarfing sharded params (2.73 GB) and the temp buffer (1.5 GB). Lesson: once FSDP has crushed the param/grad/optimizer footprint, activations become the dominant term. The next optimisation lever is activation memory (more aggressive checkpointing, smaller micro-batch, shorter sequences) — not more param sharding.

Recall One-line summary you should now be able to say

FSDP replaces full copies of params/grads/optimizer with $\frac{1}{N}$ shards (memory $\times \frac{1}{N}$ , at 16 bytes/param in mixed precision), paying all-gather + reduce-scatter communication that scales as $\frac{N - 1}{N} \to 1$ per GPU; use it only when the model won't fit under DP, and always budget the un-sharded activations separately with Activation Checkpointing.