FSDP lets many GPUs train one model too big for any single GPU by giving each GPU only a slice of the model and having them lend slices to each other just for the moment a computation needs them. Everything on this page is a tool for answering one question: "How do we split a model into pieces, and glue pieces back together only when needed?"
Before you can read the parent note , you need to know what every letter and arrow in it means. We build them one at a time — no symbol is used before it is drawn.
Definition GPU (a worker with a fixed-size desk)A GPU is a processor that does the huge number-crunching of training. The key fact: it has a fixed amount of memory (an A100 has 80 gigabytes — GB). Think of it as a desk of fixed size. If your paperwork doesn't fit on the desk, you cannot work — full stop.
A byte is the unit we measure memory in. One number stored in "fp32" (a common precision) takes 4 bytes . So 175 billion numbers take 175 × 1 0 9 × 4 = 700 hold them, before you compute anything.
That overflowing desk is the whole reason FSDP exists. Keep the picture: desk = one GPU's memory; the stack of paper = what we must store.
A model is a big pile of numbers called parameters . But training stores more than just the parameters. Look at the coloured stack:
θ θ whole collection of a model's tunable numbers — the "knobs" that get adjusted during learning. When you see θ
∇ θ L ∇ "the slope with respect to" . L loss — a single number saying how wrong the model currently is. So ∇ θ L "how the loss changes as we nudge each parameter" . It is the arrow telling every knob which way to turn to reduce error. It has exactly one number per parameter, so it is the same size as θ .
Definition Optimizer states (m v m v the same size as θ , and both are stored in 8-byte precision.
So per parameter we store roughly:
N N how many GPUs share the job . If N = 8 device i where i 0 , 1 , 2 , … , N − 1
Definition Shard, and the notationθ ( i ) shard a pile of numbers means to cut it into N i i . We write θ ( i ) θ i θ d θ ( i ) N d
Because N 1 as long as we never need the whole pile at once. The next section handles that "as long as".
f l h l layers . Layer l f l h l h l + 1 = f l ( h l , θ l ) θ l just this layer's slice of the parameters.
The core problem: a layer's computation is usually a matrix multiply y = W x cannot multiply by half a matrix and get half an answer — a single output number needs an entire row of W θ l ( i ) run layer l whole θ l
Intuition Borrow, use, returnOwn a slice → borrow the missing slices → build the full layer → compute → throw the borrowed slices away. This borrow-and-return dance is the heartbeat of FSDP, and it needs the two "collective" operations below.
A collective is an operation all GPUs perform together, passing data between each other.
Definition All-gather (borrow everyone's slice)Before: device i x ( i ) every device holds the full concatenation [ x ( 0 ) , x ( 1 ) , … , x ( N − 1 ) ] θ l
Definition Reduce-scatter (sum, then keep only your slice)Before: every device holds a full vector g i ∑ devices g ( i ) sum over all devices of just the i reduce = add up matching pieces from every GPU; scatter = hand each GPU back only its own summed slice.
sum the gradients?
Each GPU sees a different slice of the training data, so each computes a different opinion of "which way should the knobs turn". The true update is the sum of these opinions. Reduce-scatter does the summing and the slicing in a single message-efficient step — see Communication Collectives .
Bytes and 24 per parameter
All-gather borrows slices
Reduce-scatter sums gradients
FSDP three-phase protocol
Each box is one thing you now know; together they feed the parent's three-phase protocol. Related prerequisites worth a look: Data Paralelism (the "everyone has a full copy" starting point FSDP improves on), Model Parallelism and Pipeline Parallelism (other ways to split a model), Mixed Precision Training (where the 2-byte and 4-byte counts come from), Gradient Accumulation , Activation Checkpointing (the missing piece FSDP does not handle), and ZeRO Optimizer (the theory FSDP implements).
Worked example Does a 175B model fit on 64 GPUs?P = 175 × 1 0 9 N = 64 24 × 175 × 1 0 9 = 4.2 × 1 0 12 = 4.2 64 4.2 TB ≈ 65.6 Fits on an 80 GB A100 (leaving room for the temporary all-gather buffer). This is the number the parent note reports as "≈66 GB".
What does a GPU's fixed memory act like in our pictures? A desk of fixed size — if the paperwork doesn't fit, you can't work.
How many bytes does one fp32 parameter cost, and why 24 per parameter total? 4 bytes; total 4 (param) + 4 (grad) + 8 (m) + 8 (v) = 24.
What does θ The whole collection of a model's tunable numbers (weights/knobs).
Read ∇ θ L How the loss L
What are m v Adam's momentum and variance; each is the same size as θ
What does θ ( i ) Device i θ d / N
Why can't a layer just use its own shard directly? A matrix multiply needs whole rows, so the full θ l
Before/after state of all-gather? Before: each device has its slice; after: every device has the full concatenation.
Before/after state of reduce-scatter? Before: every device has a full gradient; after: device i i
Why must gradients be summed across devices? Each GPU sees different data; the true update is the sum of all their gradient opinions.
