FSDP and sharded training

What FSDP Actually Does

Fully Sharded Data Parallel is a memory optimization strategy that shards (partitions) three key components across data-parallel workers:

Model parameters ( $\theta$ )
Gradients ( $\nabla_\theta \mathcal{L}$ )
Optimizer states (momentum, variance for Adam)

Traditional data paralelism keeps full copies of all three on every device. FSDP keeps only $\frac{1}{N}$ of each on each of $N$ devices.

How FSDP Works: The Three-Phase Protocol

Phase 1: Forward Pass

Before each layer's computation:

Device $i$ holds shard $\theta_l^{(i)}$ for layer $l$
All-gather operation: collect $\{\theta_l^{(0)}, \theta_l^{(1)}, \ldots, \theta_l^{(N-1)}\}$ from all devices
Reconstruct full $\theta_l$ temporarily in device memory
Compute layer's forward pass: $h_{l+1} = f_l(h_l, \theta_l)$
Discard borrowed shards, keep only $\theta_l^{(i)}$

Why this step? Each layer needs all its parameters to compute activations correctly. We can't split a matrix multiply $y = Wx$ if we only have half of $W$ .

Phase 2: Backward Pass

For each layer in reverse:

All-gather $\theta_l$ again (needed for gradient computation)
Compute local gradients: $\frac{\partial \mathcal{L}}{\partial \theta_l}$
Reduce-scatter gradients: each device gets only $\nabla_{\theta_l^{(i)}}$ (its shard's gradient)
Discard full $\theta_l$ again

Why reduce-scatter? After computing the full gradient, we sum contributions across devices (reduce) and distribute chunks (scatter) so each device updates only its shard.

Phase 3: Optimizer Step

Each device independently updates its shard: $\theta_l^{(i)} \leftarrow \theta_l^{(i)} - \eta \cdot \text{Adam}(\nabla_{\theta_l^{(i)}}, m_l^{(i)}, v_l^{(i)})$

Optimizer states $(m, v)$ are also sharded, so each device stores only $\frac{1}{N}$ of the momentum/variance bufers.

Formula Memory Savings Calculation

For a model with $P$ parameters:

Without FSDP: Memory per device = $P \cdot (4 + 4 + 8) = 24P$ bytes
(4 bytes params fp32, 4 bytes grads, 8 bytes momentum, 8 bytes variance)
With FSDP: Memory per device = $P \cdot \frac{24}{N} + \text{temp}$ bytes
where $\text{temp}$ is the temporary all-gather buffer for one layer Derivation: Sharding divides the persistent memory (params, grads, optimizer states) by $N$ . The temporary buffer scales with layer size, not model size—typically 1-2% of model memory for transformer layers.

For $P = 175 \times 10^9$ , $N = 64$ :

Without FSDP: $24 \times 175 \text{ GB} = 4.2 \text{ TB}$ per GPU❌
With FSDP: $\frac{4.2 \text{ TB}}{64} \approx 66 \text{ GB}$ per GPU ✅ (fits on A100 80GB)

Communication Primitives

FSDP relies on two collective operations:

Comparison with ZeRO

FSDP is PyTorch's implementation of ZeRO (Zero Redundancy Optimizer) from Microsoft DepSpeed:

ZeRO Stage 1: Shard optimizer states only → $4\times$ memory reduction
ZeRO Stage 2: Shard optimizer states + gradients → $8\times$ memory reduction
ZeRO Stage 3: Shard optimizer states + gradients + parameters → $N\times$ memory reduction (FSDP = ZeRO-3)

Why three stages? Each stage introduces more communication overhead. ZeRO-1 is faster for moderate models; ZeRO-3 is necessary only when models exceed single-GPU memory.

Worked example Training GPT-3 with FSDP

Setup: GPT-3 with 175B parameters, mixed precision (fp16 forward, fp32 optimizer), batch size 1536, 1024 A100 GPUs.

Without FSDP:

Memory per param: 2 (fp16 param) + 2 (fp16 grad) + 4 (fp32 master) + 8 (momentum) + 8 (variance) = 24 bytes
Total: $175 \times 10^9 \times 24 = 4.2$ TB per GPU → impossible on 80 GB A100s

With FSDP (ZeRO-3):

Sharded memory: $\frac{4.2 \text{ TB}}{1024} = 4.1$ GB persistent GPU
Temporary buffer (one transformer layer): ~1.5 GB
Activations (recomputed with activation checkpointing): ~30 GB
Total: ~36 GB per GPU ✅

Why this works: The all-gather overhead (fetching other shards) is offset by NVLink/InfiniBand bandwidth. For transformers, communication time< computation time due to large matrix multiplies.

Hybrid Sharding Strategies

Modern frameworks allow hybrid paralelism:

FSDP + Pipeline Parallelism:

Partition model into stages (layers 1-12 on GPU group A, layers 13-24 on GPU group B)
Within each stage, shard with FSDP
Benefit: Reduce all-gather scope to intra-stage devices

FSDP + Tensor Paralelism:

Split large layers (e.g., 12,288 × 12,288 weight matrix) across GPUs within a node
Shard parameters across nodes with FSDP
Benefit: Tensor parallelism for computation, FSDP for memory

Formula for effective parallelism: $\text{Total GPUs} = N_{\text{data}} \times N_{\text{pipeline}} \times N_{\text{tensor}}$

Common mistake Assuming FSDP is Always Faster

Wrong idea: "More memory optimization = better performance."

Why it feels right: Saving memory sounds like pure upside.

Reality: FSDP trades memory for communication. On slow interconnects (PCIe instead of NVLink, or high-latency cross-datacenter links), the all-gather/reduce-scatter overhead dominates. A 10 Gbps network transfering 700 MB per layer = 560 ms per layer, far slower than computation.

Steel-man: The mistake comes from conflating memory efficiency with computational efficiency. FSDP is a memory-constrained optimization—use it when you have no choice, not as a default.

Fix: Profile communication vs. computation ratio. If $\frac{T_{\text{comm}}}{T_{\text{comp}}} > 0.3$ , FSDP likely hurts performance. Consider mixed precision, gradient checkpointing, or smaller batch sizes before sharding.

Common mistake Forgetting Activation Memory

Wrong idea: "FSDP shards everything, so I can train infinite-size models."

Why it feels right: FSDP shards parameters, gradients, and optimizer states—sounds comprehensive.

Reality: FSDP does not shard activations. Each device still stores $\text{batch\_size} \times \text{seq\_len} \times \text{hidden\_dim}$ activations for backprop. For GPT-3 with batch size 8, sequence length 2048, hidden dim 12,288: $8 \times 2048 \times 12288times 2\text{ bytes} = 402$ GB per layer.

Fix: Combine FSDP with activation checkpointing (recompute activations during backward instead of storing). This trades33% more computation for 10× less activation memory.

Implementation Details

Wrapping granularity: FSDP can wrap individual layers, modules, or the entire model. Finer wrapping reduces temporary memory (smaller all-gather buffers) but increases communication frequency.

Typical strategy:

model = TransformerModel(...)
for layer in model.transformer_layers:
    layer = FSDP(layer)  # Wrap each transformer block
model = FSDP(model)  # Outer wrap for top-level

Mixed precision: FSDP stores shards in fp32 (for optimizer) but communicates in fp16 (for bandwidth). All-gather converts fp32 → fp16 before sending.

CPU offloading: Some FSDP implementations allow offloading inactive shards to CPU RAM, trading PCIe bandwidth for GPU memory. Rarely used due to severe latency penalties.

Recall Explain to a 12-Year-Old

Imagine you and 7 friends are building a huge LEGO castle (the AI model). The castle is so big that one person can't hold all the pieces at once.

Normal way (data parallel): Each of you has a complete copy of all pieces. You each build the same castle 8 times. This wastes a lot of space!

FSDP way: You divide the pieces into 8 boxes. You hold box1, your friend holds box 2, etc. When you need to build the tower (which needs pieces from all boxes), everyone quickly passes their boxes around the circle so you can grab the pieces you need. After building the tower, you give the extra boxes back.

Why does this help? You only need to store1/8 of the pieces permanently. The passing-around takes a bit of time, but it's way better than neding 8× the shelf space. For really massive AI models (like ChatGPT's big brother), this "sharing pieces" trick is the only way to build them without buying impossibly expensive computers.

Connections

Data Paralelism — FSDP extends data parallelism with sharding
Model Parallelism — FSDP is orthogonal; can combine both
Mixed Precision Training — FSDP leverages fp16 communication, fp32 computation
Gradient Accumulation — Accumulate over microbatches to reduce per-step communication
Activation Checkpointing — Essential complement to FSDP for activation memory
ZeRO Optimizer — FSDP implements ZeRO-3 from DepSpeed
Communication Collectives — All-gather, reduce-scatter are FSDP's primitives
Pipeline Parallelism — Hybrid: pipeline + FSDP within stages

#flashcards/ai-ml

What are the three components FSDP shards across devices? :: Model parameters, gradients, and optimizer states (momentum, variance).

What collective operation does FSDP use to reconstruct full parameters forward pass?

All-gather (each device broadcasts its shard to all others).

What is the memory reduction factor for ZeRO-3/FSDP with N devices?

Approximately N× reduction for persistent memory (parameters + gradients + optimizer states).

Why does FSDP discard borrowed parameter shards after each layer?

To minimize peak memory usage—keeping full parameters would negate the sharding benefit.

What is the difference between ZeRO-2 and ZeRO-3?

ZeRO-2 shards optimizer states + gradients. ZeRO-3 additionally shards model parameters.

When should you NOT use FSDP?

When the model fits in GPU memory with standard data parallelism, because FSDP's communication overhead reduces throughput.

What memory component does FSDP not shard?

Activations (they remain per-device and scale with batch size).

What is the communication cost of all-gather for a shard of size S on N devices?

\frac{N-1}{N} \times S

per device (each device receives N-1 shards from others).

What technique complements FSDP to handle activation memory?

Activation checkpointing (recompute activations during backward instead of storing).

In hybrid FSDP + pipeline parallelism, what does FSDP shard?

Parameters/gradients/optimizer states within each pipeline stage (not across stages).

Concept Map

Hinglish (regional understanding)

Intuition Hinglish mein samjho

FSDP matlab Fully Sharded Data Parallel ek aisi technique hai jo bahut bade AI models ko train karne ke liye use hoti hai. Samjho ki tumhare paas ek bohot bada model hai, jaise GPT-3 with 175 billion parameters—isko ek GPU me fit karna impossible hai kyunki itna memory hi nahi hota. Traditional data parallelism me har GPU pe pora model copy hota hai, lekin FSDP me har GPU sirf apna assigned "shard" (piece) rakhta hai.

Jab forward pass chalti hai, toh har GPU temporarily dosre GPUs se unke shards mang leta hai (all-gather operation), puri layer compute karta hai, aur fir extra shards ko discard kar deta hai—sirf apna piece rakhta hai. Backward pass me same chez hoti hai: gradients compute hone ke bad, har GPU sirf apne shard ka gradient rakhta hai (reduce-scatter operation). Isse memory usage N GPUs ke sath N times kam ho jata hai—matlab 64 GPUs use karke tum 64× bada model train kar sakte ho bina memory overflow ke.

Lekin ek tradeoff hai: FSDP me communication overhead badh jata hai kyunki har layer ke liye all-gather aur reduce-scatter operations hote hain. Agar tumhara model already GPU me fit ho raha hai, toh FSDP use karna actually slower ho sakta hai. Isliye FSDP ko tab use karo jab model ka size GPU memory sezyada ho—warna normal data parallelism hi better rehta hai. Activation memory bhi separately handle karni padti hai activation checkpointing se, kyunki FSDP sirf parameters, gradients, aur optimizer states ko shard karta hai, activations ko nahi.

What FSDP Actually Does

How FSDP Works: The Three-Phase Protocol

Phase 1: Forward Pass

Phase 2: Backward Pass

Phase 3: Optimizer Step

Communication Primitives

Comparison with ZeRO

Hybrid Sharding Strategies

Implementation Details

Connections

Concept Map

Hinglish (regional understanding)

Go deeper — visual, from zero

Test yourself — Scaling & Efficient Architectures

Connections