Chapter 10 · Section 10.3 · 50 min

ZeRO and FSDP

Look at what data parallelism does with 64 GPUs. Every GPU stores the same 16 bytes per parameter, so the cluster holds 64 identical copies of the weights, 64 identical copies of the averaged gradient, and 64 identical copies of Adam's moments. Then every GPU performs the same optimizer update on the same numbers and gets the same result. Sixty-three of those copies, and sixty-three of those updates, are pure waste.

Rajbhandari et al. (2020) turned that observation into ZeRO, the Zero Redundancy Optimizer. Each GPU keeps and updates only its own slice of the model states, and the slices are combined only when and where they are needed. The memory per GPU for model states falls from $16N$ toward $16N/n_{\text{dp}}$ for $n_{\text{dp}}$ data-parallel GPUs, and remarkably, the first two stages cost no more communication than plain data parallelism. PyTorch's Fully Sharded Data Parallel (FSDP) is the same idea, built into the framework. This section derives the memory and communication of each stage, implements sharding in NumPy, and explains when it stops helping.

The idea, one stage at a time

Flatten all parameters into one long vector of $N$ numbers and cut it into $n_{\text{dp}}$ equal contiguous shards, so rank $r$ owns shard $r$ . Ownership means rank $r$ is the only one that stores the optimizer state for those parameters and the only one that updates them. ZeRO comes in three stages, each sharding one more kind of model state.

Stage 1 shards the optimizer state. Each rank keeps the full bfloat16 weights and computes full gradients as in data parallelism, but stores Adam's $\mathbf{m}$ and $\mathbf{v}$ and the float32 master weights only for its own shard. At the end of the backward pass, instead of an all-reduce, it runs a reduce-scatter, so that rank $r$ receives the averaged gradient for shard $r$ only, which is all its optimizer needs. It updates its shard of the master weights, casts them to bfloat16, and an all-gather gives every rank the full updated weights for the next forward pass.

Here is the first surprise. Recall from §10.2 that a ring all-reduce is a reduce-scatter followed by an all-gather. Stage 1 runs exactly those two halves, just with the optimizer update in between. So it moves exactly as many bytes as data parallelism, and saves $12N(1 - 1/n_{\text{dp}})$ bytes per GPU for free.

Stage 2 also shards the gradients. Rank $r$ only ever needs the averaged gradient of shard $r$ , so there is no reason to keep the full gradient around. During the backward pass, as each bucket of gradients becomes ready, it is reduce-scattered at once, and each rank keeps only its own part and frees the rest. The communication is still one reduce-scatter and one all-gather per step.

There is one catch, and interviewers like it. With gradient accumulation over several micro-batches, stage 1 can add each micro-batch's full local gradient into its full gradient buffer and communicate once per optimizer step. Stage 2 has thrown the full buffer away, so it must reduce-scatter after every micro-batch's backward pass, or else keep the full gradient after all and lose its saving. The same is true of stage 3, which must also gather the weights again for every micro-batch. That is why pipelined runs, which always process many micro-batches per step (§10.5), usually stop at stage 1.

Stage 3 shards the parameters themselves. Now each rank stores only its $1/n_{\text{dp}}$ of the bfloat16 weights. Before a layer's forward pass, the ranks all-gather that layer's weights, compute, and free the gathered copy. In the backward pass they all-gather the layer's weights again, compute the gradients, reduce-scatter them so each rank keeps its gradient shard, and free the gathered weights. Only one or two layers' worth of full weights exists at any moment.

Memory per GPU

With $K = 12$ bytes of optimizer state per parameter (float32 master weights plus two float32 moments, as in §10.1), the model states per GPU are

stage	weights	gradients	optimizer	total per GPU
data parallel	$2N$	$2N$	$KN$	$16N$
ZeRO-1 ( $P_{os}$ )	$2N$	$2N$	$KN/n_{\text{dp}}$	$4N + 12N/n_{\text{dp}}$
ZeRO-2 ( $P_{os+g}$ )	$2N$	$2N/n_{\text{dp}}$	$KN/n_{\text{dp}}$	$2N + 14N/n_{\text{dp}}$
ZeRO-3 ( $P_{os+g+p}$ )	$2N/n_{\text{dp}}$	$2N/n_{\text{dp}}$	$KN/n_{\text{dp}}$	$16N/n_{\text{dp}}$

In plain words, each stage moves one more term under the $1/n_{\text{dp}}$ , and as the number of GPUs grows the totals approach $4N$ , $2N$ , and zero.

The paper's own example is a 7.5B model on 64 GPUs. The cell reproduces it and adds our two models.

Python · cell 1

The paper's numbers, 120, 31.4, 16.6, and 1.9 GB, come out exactly. Notice what each stage buys. Stage 1 gives most of the savings for the least complexity, which is why many large runs stop there (DeepSeek-V3 used ZeRO-1 across its data-parallel ranks). Stage 2 is nearly free on top. Stage 3 is the only one that lets a model whose bfloat16 weights alone exceed a GPU, such as a 70B model at 141 GB, train with data parallelism at all.

Warm-up

ZeRO memory per GPU

Implement zero_memory(N, n_dp, stage, K=12) returning the bytes of model states per GPU for ZeRO stage 0 (plain data parallelism), 1, 2, or 3, with 2-byte weights and gradients and K bytes of optimizer state per parameter. Raise ValueError for any other stage.

def zero_memory(N, n_dp, stage, K=12):
    return 16 * N / n_dp

Hidden tests, hints, and AI help, from $1.99 a month.

Sharding, simulated

Let's see stage 3 work, with simulated ranks in one process. The model below is a stack of $L$ linear layers with a tanh, and each rank owns $1/n_{\text{dp}}$ of every layer's flattened weight matrix. The forward and backward passes gather a layer's weights just before using them and drop them right after, and the backward pass reduce-scatters each layer's gradient. At the end we check that every rank's gradient shard equals the matching piece of the gradient plain data parallelism would compute, and we count how many weights were ever resident at once.

Python · cell 2

Every rank ends with exactly its shard of the averaged gradient, and only one layer's weights were ever materialized. What we paid for it is visible in the code, two all-gathers of every layer instead of zero.

Interview

A sharded Adam step

Implement the optimizer step of ZeRO stages 1 and 2, for n_dp simulated ranks in one process. zero_adam_step(grads, w_shards, m_shards, v_shards, t, lr=1e-3, b1=0.9, b2=0.999, eps=1e-8) receives each rank's full local gradient (grads[r], a 1-D array of length n), and each rank's shard of the float32 master weights and Adam moments (w_shards[r] and so on, the r-th piece of np.array_split of a length-n vector into n_dp pieces). Do three things. (1) Reduce-scatter, so rank r gets the average over ranks of shard r of the gradients. (2) Apply Adam with bias correction at step t (starting from 1) to each rank's own shard only, $\mathbf{m} \leftarrow \beta_1\mathbf{m} + (1-\beta_1)\mathbf{g}$ , $\mathbf{v} \leftarrow \beta_2\mathbf{v} + (1-\beta_2)\mathbf{g}^2$ , $\mathbf{w} \leftarrow \mathbf{w} - \eta\,\hat{\mathbf{m}}/(\sqrt{\hat{\mathbf{v}}} + \epsilon)$ with $\hat{\mathbf{m}} = \mathbf{m}/(1-\beta_1^t)$ and $\hat{\mathbf{v}} = \mathbf{v}/(1-\beta_2^t)$ . (3) All-gather, so every rank receives the full updated weight vector. Return (full_w, w_shards, m_shards, v_shards), where full_w[r] is rank r's copy of the full weights, and the shard lists are new arrays. Do not modify the inputs.

import numpy as np

def zero_adam_step(grads, w_shards, m_shards, v_shards, t, lr=1e-3, b1=0.9, b2=0.999, eps=1e-8):
    raise NotImplementedError

Hidden tests, hints, and AI help, from $1.99 a month.

Communication volume

Now count bytes per GPU per step, in units of the gradient or parameter size, using the ring costs of §10.2. A reduce-scatter or all-gather of an $N$ -element vector sends $(n_{\text{dp}}-1)N/n_{\text{dp}} \approx N$ elements per GPU, and an all-reduce twice that.

Data parallelism all-reduces the gradients, about $2N$ elements.
Stages 1 and 2 reduce-scatter the gradients ( $N$ ) and all-gather the updated weights ( $N$ ), about $2N$ , the same as data parallelism.
Stage 3 all-gathers the weights in the forward pass ( $N$ ), again in the backward pass ( $N$ ), and reduce-scatters the gradients ( $N$ ), about $3N$ , or 1.5 times data parallelism.

This is the accounting of Rajbhandari et al. Multiply by the bytes per element. With bfloat16 weights and gradients, stage 3 moves about $6N$ bytes per GPU per step. Many recipes reduce gradients in float32 for accuracy (Llama 3 reduce-scatters its gradients in float32 across FSDP ranks), which makes the gradient term twice as large.

Can stage 3 hide its extra traffic? Its weight all-gathers do not depend on anything computed in the step, so they can be issued ahead of time, prefetched, and overlapped with the previous layer's computation in both the forward and the backward pass. Repeat the argument of §10.2 with the whole step. Each GPU computes $6N$ FLOPs per token over its $n_{\text{tok}}$ tokens, at $F$ FLOP/s, and moves $6N$ bytes over a link of $\mathcal{B}$ bytes per second, so the communication hides when

\frac{6Nn_{\text{tok}}}{F} \ge \frac{6N}{\mathcal{B}} \quad\Longleftrightarrow\quad n_{\text{tok}} \ge \frac{F}{\mathcal{B}} ,

the same threshold as data parallelism, about 8,000 tokens per GPU per step for H100s at 40% utilization over InfiniBand. Stage 3 moves 1.5 times the bytes, but it also gets the forward pass to hide them in. In practice per-layer latency, imperfect overlap, and the memory taken by prefetched layers make it somewhat more expensive than this ideal.

Python · cell 3

For Llama 3 70B on 128 GPUs, stage 3 moves 420 GB per GPU per step, 8.4 seconds over InfiniBand, while computing 8,192 tokens per GPU takes 8.8 seconds. It just barely hides, and only with perfect overlap. Fewer tokens per GPU, or a slower network, and the GPUs wait.

Core

Communication per step for each stage

Implement zero_comm_bytes(N, n_dp, stage, param_bytes=2, grad_bytes=2) returning the bytes each GPU sends per optimizer step under the ring algorithms, for ZeRO stage 0, 1, 2, or 3. Stage 0 all-reduces the gradients. Stages 1 and 2 reduce-scatter the gradients and all-gather the updated weights. Stage 3 all-gathers the weights twice (forward and backward) and reduce-scatters the gradients. Use the exact factor $(n_{\text{dp}}-1)/n_{\text{dp}}$ , and note that gradients and weights may have different sizes per element.

def zero_comm_bytes(N, n_dp, stage, param_bytes=2, grad_bytes=2):
    return 2 * N * grad_bytes          # data parallelism, ignoring the stage

Hidden tests, hints, and AI help, from $1.99 a month.

FSDP, ZeRO-3 inside PyTorch

PyTorch's Fully Sharded Data Parallel (Zhao et al. 2023) implements stage-3 sharding natively. Its main design choice is the unit. You wrap groups of modules, typically one Transformer block each, and FSDP flattens each unit's parameters into one contiguous buffer, the FlatParameter in the original design, so that gathering a unit is a single all-gather and reducing its gradient a single reduce-scatter. Units are the granularity of everything. A larger unit means fewer, larger collectives (less latency) but more memory held while it is gathered. The newer FSDP2 (fully_shard, used below) keeps the unit but drops the flat buffer from the user's view. Each parameter is stored as a distributed tensor sharded along its first dimension, and the unit's shards are still packed into one buffer for each collective. Because every parameter keeps its own identity, frozen and trainable parameters can share a unit, sharded state dicts need no extra communication, and the same tensors compose with tensor parallelism on another axis of the device mesh (§10.6).

The memory per GPU is then about the sharded model states plus the largest gathered unit,

\text{memory} \approx \frac{16N}{n_{\text{dp}}} + 2 \times 2\max_i N_i + \text{activations},

where $N_i$ is the number of parameters in unit $i$ , and each gathered unit costs 2 bytes per parameter in bfloat16. The factor of 2 in front is there because, with prefetching, the next unit is being gathered while the current one computes, so up to two units are resident at once. In the backward pass the unsharded gradient of the unit being differentiated adds about one more unit's worth before it is reduce-scattered. This is why wrapping each block separately matters. Wrapping a whole 70B model as one unit would gather all 141 GB of bfloat16 weights at once and defeat the purpose.

Prefetching is how FSDP overlaps. In the backward pass it issues the all-gather for the next unit (the previous block in the model) before computing the current one's gradients, so the network and the GPU work at the same time. It can also prefetch in the forward pass. A rate limiter keeps it from racing too far ahead and filling memory with gathered weights.

Sharding strategies trade memory against communication. (In the newer fully_shard interface the same choice is a flag, reshard_after_forward=True for full sharding and False to keep the weights gathered until the backward pass.)

strategy	what is sharded	gathers per step	equivalent
`FULL_SHARD`	weights, gradients, optimizer	weights twice (forward and backward)	ZeRO-3
`SHARD_GRAD_OP`	gradients and optimizer; weights kept gathered from forward to backward	weights once	ZeRO-2
`NO_SHARD`	nothing	none, gradients all-reduced	DDP
`HYBRID_SHARD`	everything, within a group; replicated across groups	weights twice, within the group	ZeRO-3 inside, DDP outside

Hybrid sharding, often called HSDP, deserves a closer look because it is how FSDP is usually run across many nodes. Shard within a node of 8 GPUs, so that all the weight all-gathers and gradient reduce-scatters travel over NVLink. Across nodes, the replicas only need to agree on gradients, and each GPU all-reduces just its own $1/8$ shard of the gradient with its counterparts in other nodes. Model states per GPU are $16N/8$ instead of $16N/n_{\text{dp}}$ , but cross-node traffic per GPU drops from about $3 \times 2N$ bytes to about $2 \times 2N/8$ , a factor of 12.

Core

Hybrid sharding, memory against network

A cluster has n_gpus GPUs in nodes of gpus_per_node. You shard model states over groups of shard consecutive GPUs (ZeRO-3 within a group) and replicate across the R = n_gpus // shard groups. Implement hsdp_cost(N, n_gpus, shard, gpus_per_node=8) returning (state_bytes, cross_node_bytes), the model-state bytes per GPU ( $16N/\text{shard}$ ) and the bytes each GPU sends per step in collectives whose group spans more than one node. Weights and gradients are bfloat16. Within a group the collectives are two all-gathers of the weights and one reduce-scatter of the gradients (each $(s-1)/s \cdot 2N$ bytes for group size $s$ ), and they span nodes only if shard > gpus_per_node. Across groups, each GPU all-reduces its gradient shard of $2N/s$ bytes with the $R$ GPUs holding the same shard ( $2(R-1)/R$ times that), and this always spans nodes when $R > 1$ (assume replicas live on different nodes).

def hsdp_cost(N, n_gpus, shard, gpus_per_node=8):
    return 16 * N / n_gpus, 3 * 2 * N

Hidden tests, hints, and AI help, from $1.99 a month.

CPU offload. ZeRO-Offload (Ren et al. 2021) keeps the float32 master weights and Adam moments, $12N$ bytes, in CPU memory and runs the optimizer step on the CPU, streaming bfloat16 gradients down and updated bfloat16 weights back over PCIe. A server typically has far more CPU memory than GPU memory, so a single GPU can then fine-tune a model whose states are several times its capacity. The cost is the PCIe link, tens of GB/s, and a CPU optimizer that is much slower than a GPU one. ZeRO-Infinity extends the same idea to NVMe drives. Offloading makes sense when memory, not time, is the constraint, as in fine-tuning on a small machine. It rarely pays in large pretraining, where GPU memory across many nodes is plentiful and time is not.

Activations are not sharded. Every ZeRO stage shards model states only. Each rank still runs its own micro-batch and stores its own activations in full. ZeRO-R adds partitioning of activation checkpoints across tensor-parallel ranks, and the activations themselves are the job of recomputation, tensor and sequence parallelism, and context parallelism, in the next sections.

Fewer bytes on the wire. ZeRO++ (Wang et al. 2023) quantizes the weight all-gather to 8 bits, keeps a secondary copy of the weights sharded within each node so that the backward all-gather stays on NVLink, and quantizes gradients during the reduce-scatter, cutting cross-node traffic by about four times.

When sharding stops helping

Sharding has diminishing returns, for three reasons that are worth stating precisely.

The floor is activations. Model states per GPU fall like $16N/n_{\text{dp}}$ , but activations do not fall at all. For Llama 3 70B at 8,192 tokens per sequence, model states on 512 GPUs are 2.2 GB, and one sequence's activations are 182 GB even with FlashAttention. ZeRO alone cannot train it. Recomputation brings the activations to 10.7 GB, but then you pay a third more compute, and the next step is to split each sequence's work across GPUs with tensor or context parallelism.

The working set does not shrink. A gathered unit is the same size however many GPUs share it. Once $16N/n_{\text{dp}}$ is small, the gathered blocks and their prefetch buffers dominate the model-state memory.

Communication grows relative to compute per GPU. The bytes each GPU sends per step stay near $3 \times 2N$ however many GPUs you add, while each GPU's share of the work shrinks if the global batch is fixed. With a 4M-token batch on 512 GPUs, each GPU gets 8,192 tokens, right at the break-even point computed above. On 2,048 GPUs each gets 2,048 tokens, and communication takes four times as long as computation. As with plain data parallelism, the global batch caps how far FSDP scales, and the way past the cap is to add model parallelism, so that several GPUs work on the same tokens and the data-parallel degree stays modest.

Python · cell 4

In PyTorch

The current FSDP interface, fully_shard, is applied bottom-up, first to each block and then to the root module. Building the model on the meta device first means no GPU ever holds the full unsharded weights. This code needs several GPUs and is shown for reference.

A detail of the bytes. With this mixed-precision policy, each rank stores its shard of the parameters in float32, and that shard is the master copy. The bfloat16 weights exist only while a unit is gathered, and the gradient shards come out of the float32 reduce-scatter in float32. Per parameter the shard holds $4 + 4 + 8 = 16$ bytes divided by $n_{\text{dp}}$ , the same total as the table above, arranged differently.

DeepSpeed exposes the same stages through a configuration dictionary.

🎯 Interview questions on this section

Explain ZeRO stages 1, 2, and 3. What does each save, and what does each cost?

Stage 1 shards the optimizer state (master weights and Adam moments), stage 2 also the gradients, and stage 3 also the weights. Per GPU, model states fall from $16N$ to $4N + 12N/n_{\text{dp}}$ , $2N + 14N/n_{\text{dp}}$ , and $16N/n_{\text{dp}}$ . Stages 1 and 2 replace the gradient all-reduce with a reduce-scatter and an all-gather of the updated weights, which is the same volume, about $2N$ elements per GPU. Stage 3 must gather the weights for the forward and again for the backward pass, about $3N$ elements, 1.5 times data parallelism.

Why does ZeRO stage 1 cost no extra communication?

A ring all-reduce is a reduce-scatter followed by an all-gather. Stage 1 runs the reduce-scatter on the gradients, so each rank gets the averaged gradient of the shard it owns, updates that shard, and then runs the all-gather on the updated weights instead of on gradients. The bytes moved are identical, and the redundant optimizer updates disappear.

What is the peak memory of FSDP, and how does wrapping affect it?

About $16N/n_{\text{dp}}$ for the sharded states, plus the full bfloat16 weights of the units currently gathered (the computing unit and, with prefetch, the next one), plus their gradient buffers, plus activations. Each unit is gathered as one flat buffer, so the largest unit sets the extra term. Wrapping each Transformer block separately keeps it small. Wrapping the whole model as one unit gathers everything and gives no saving.

When does FSDP stop scaling?

Model states shrink with $n_{\text{dp}}$ but activations do not, so memory bottoms out at the activations and gathered units. Communication per GPU stays near $6N$ bytes per step, while compute per GPU falls as the fixed global batch is spread thinner, so below about $F/\mathcal{B}$ tokens per GPU (thousands, over InfiniBand) the GPUs wait on the network. Hybrid sharding helps the network side, and tensor, pipeline, or context parallelism keeps the data-parallel degree modest.

What is hybrid sharding and why use it?

Shard model states within a group (usually a node), replicate across groups. All-gathers and reduce-scatters then run over NVLink, and across nodes each GPU only all-reduces its own gradient shard, cutting cross-node traffic per GPU from about $6N$ to about $N/2$ bytes for 8-GPU groups. The price is memory, $16N/8$ per GPU instead of $16N/n_{\text{dp}}$ , which is fine whenever the model fits that way.

When would you use ZeRO-Offload?

When GPU memory, not time, is the binding constraint, for example fine-tuning a large model on one or a few GPUs. It moves the $12N$ bytes of optimizer state and the optimizer step to the CPU and streams gradients and weights over PCIe. It slows each step, so in large pretraining runs with many GPUs it is usually better to shard across GPUs instead.

🧭 What to remember

Data parallelism stores $n_{\text{dp}}$ identical copies of the model states. ZeRO gives each rank ownership of one shard and keeps only that.
Per GPU, stages 1, 2, and 3 need $4N + 12N/n_{\text{dp}}$ , $2N + 14N/n_{\text{dp}}$ , and $16N/n_{\text{dp}}$ bytes. For 7.5B parameters on 64 GPUs that is 31.4, 16.6, and 1.9 GB instead of 120.
Stages 1 and 2 cost the same communication as data parallelism (a reduce-scatter plus an all-gather is an all-reduce). Stage 3 costs 1.5 times as much, but prefetching hides it when each GPU has at least $F/\mathcal{B}$ tokens per step.
FSDP is ZeRO-3 in PyTorch, with per-block units, prefetching, and strategies from full sharding to hybrid sharding, which keeps gathers on NVLink.
Sharding does not touch activations, and its communication per GPU does not shrink, so it stops helping once activations dominate or tokens per GPU fall below the break-even point.
Next, §10.4 splits the layers themselves across GPUs, which shrinks activations as well as weights.

Rajbhandari et al. (2020). ZeRO: Memory Optimizations Toward Training Trillion Parameter Models. The three stages, their memory formulas, and the communication analysis of Section 7.
Zhao et al. (2023). PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel. Units, flat parameters, prefetching, the rate limiter, and hybrid sharding.
Ren et al. (2021). ZeRO-Offload: Democratizing Billion-Scale Model Training. Optimizer state and step on the CPU, with the bandwidth analysis that makes it work.
Rajbhandari et al. (2021). ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning. Offloading to CPU and NVMe.
Wang et al. (2023). ZeRO++: Extremely Efficient Collective Communication for Giant Model Training. Quantized gathers and hierarchical partitioning for slow networks.
PyTorch. Getting Started with Fully Sharded Data Parallel (FSDP2). The fully_shard interface used above.
Tazi et al. (2025). The Ultra-Scale Playbook. Its data-parallelism chapter profiles ZeRO stages and their overlap on real hardware.

Zero to Frontier is made for larger screens

ZeRO and FSDP

The idea, one stage at a time

Memory per GPU

ZeRO memory per GPU

Sharding, simulated

A sharded Adam step

Communication volume

Communication per step for each stage

FSDP, ZeRO-3 inside PyTorch

Hybrid sharding, memory against network

When sharding stops helping

In PyTorch

Read more

Watch next

Zero to Frontier is made for larger screens

The idea, one stage at a time

Memory per GPU

Sharding, simulated

Communication volume

FSDP, ZeRO-3 inside PyTorch

Offloading and other refinements

When sharding stops helping

In PyTorch

Read more

Watch next