Memory Sharding: FSDP and ZeRO Are Not a Fifth Parallelism Axis

Recall data parallelism's identity from M7-01: every GPU holds a full copy of the model, processes a different slice of the batch, and an all-reduce averages the resulting gradients so every replica applies an identical update. That description is complete for the weights and the gradients at the instant of the update, but it says nothing yet about what happens to the optimizer's own bookkeeping between steps.

Consider Adam, the optimizer most large models train with. Adam does not just apply a gradient to a weight; it maintains, for every trainable parameter, a running estimate of the mean of past gradients (the first moment) and a running estimate of their variance (the second moment), both updated at every step and both needed again at the next step. Under plain data parallelism, every one of the N data-parallel GPUs is running its own instance of the optimizer, and because every replica ends each step holding the identical updated weights, every replica must also be maintaining its own identical copy of those two moment tensors. Ten data-parallel GPUs are not splitting the optimizer's bookkeeping ten ways; they are each doing the full bookkeeping independently, in parallel, and storing the full result.

The consequence is a term in the training-memory budget that scales with the size of the model but does not shrink at all as more GPUs join a data-parallel group — a genuine waste once memory, not compute, is the binding constraint. A model whose optimizer state alone is, say, 100 GB needs that same 100 GB resident on every data-parallel GPU under the plain scheme, whether there are two GPUs in the group or two hundred.

Identity statement: the distributed optimizer — implemented as ZeRO (Zero Redundancy Optimizer) or FSDP (Fully Sharded Data Parallel) — shards optimizer state, and optionally parameters and gradients as well, across the GPUs of a data-parallel group, so that no single GPU holds a full, redundant copy of state that the group as a whole only needs to hold once.

[GROUND TRUTH] (Sources/ncp-genl/domain-7-gpu-acceleration.md): "The Distributed Optimizer (ZeRO / FSDP-style) shards optimizer states (and optionally params and gradients) across data-parallel GPUs." Read that sentence as a graduated menu rather than a single fixed behavior: the baseline form shards only the optimizer state — the term identified in section 1 as the one that duplicates for free under plain data parallelism and does not need to be fully resident on every rank at every instant. More aggressive configurations extend the same sharding logic to the gradients too, and further still to the parameters (weights) themselves, progressively reducing per-GPU memory at the cost of more frequent communication to reconstruct whatever a given GPU needs for its own slice of computation. Each rank in a sharded configuration is, at rest, responsible for only its own fraction of whichever tensors are being sharded — commonly one-Nth of the optimizer state (and, in the more aggressive configurations, one-Nth of the gradients and parameters) for a group of N data-parallel GPUs.

The critical scope word is data-parallel. This sharding happens specifically across the ranks that would otherwise be holding redundant copies of the same thing — the data-parallel group — not across GPUs performing genuinely different computation the way a tensor-parallel or pipeline-parallel group does. That scoping is what makes the technique a natural fit layered on top of data parallelism rather than a replacement for it, and it is also exactly why the next section's communication pattern looks the way it does.

Identity statement: a sharded optimizer step replaces plain data parallelism's single all-reduce with a reduce-scatter of gradients followed by an all-gather of the updated parameter shards, so that at no point does any one GPU need to hold the full, unsharded tensor for longer than the instant it is actively computing with it.

[GROUND TRUTH] (Sources/ncp-genl/domain-7-gpu-acceleration.md): the distributed optimizer "typically reduce-scatter[s] gradients, then all-gather[s] updated parameter shards." Unpack the two halves in order.

L1 — Intuition: split the averaging step itself, not just the storage

Plain data parallelism's all-reduce combines every rank's gradient tensor and hands the full, combined result back to every rank, so that every rank can apply the full update to its own full copy of the weights. A sharded optimizer cannot afford to do that, because the whole point is that no rank wants to hold the full optimizer state (and, in the more aggressive configurations, the full gradient or the full parameter tensor) at rest. So instead of an all-reduce that ends with everyone holding everything, the reduce-scatter step ends with each rank holding only its own shard of the combined, reduced result — the averaging happens across all ranks exactly as before, but the output is immediately partitioned rather than duplicated everywhere.

L2 — Mechanism: what each rank does with its shard

Once reduce-scatter has left each rank holding one shard of the fully-reduced gradient, that rank runs its own optimizer's update logic on exactly that shard — using its own resident shard of the optimizer state (its slice of the Adam moments, for instance) — and produces an updated shard of the parameters. At this point, every rank has correctly updated only its own fraction of the model's weights; no rank yet holds the complete, updated parameter tensor, because no rank needed to hold the complete tensor to compute its own shard's update. The all-gather step then reassembles the full, updated parameter tensor everywhere it is needed for the next forward pass — each rank contributes its updated shard, and every rank that needs the full tensor receives the concatenation of all the shards.

L3 — Why this produces the identical numerical result, not an approximation

The output of a reduce-scatter followed by an all-gather is exactly the same combined, averaged gradient information that a plain all-reduce would have produced — the operation is just factored into two steps so that the intermediate state (immediately after the reduce-scatter, before the all-gather) never requires any single rank to hold more than its own shard. This is why sharding optimizer state is a memory technique with no inherent accuracy cost of its own: the arithmetic each shard's owner performs is the identical arithmetic a fully-resident, unsharded optimizer would have performed on that same slice of parameters, just computed and stored in a partitioned form rather than a duplicated one. The tradeoff is not correctness for memory; it is communication frequency and volume for memory, since reconstructing the full tensor via all-gather whenever it is needed is real network traffic that a fully-resident, unsharded scheme would not have paid.

⭐ THE EARNED INSIGHT

The reduce-scatter-then-all-gather pattern is not a new communication primitive invented for this purpose — it is the same two-phase decomposition that ring-allreduce itself is built from. A plain all-reduce is mathematically identical to a reduce-scatter immediately followed by an all-gather; the distributed optimizer's trick is simply to stop between those two phases, keep each rank's shard of the result resident rather than immediately reassembling the whole thing, and only pay for the all-gather when something downstream (the forward pass, most obviously) actually needs the full, unsharded tensor back. Sharding optimizer state, in other words, is what you get by refusing to finish an operation your training loop was already implicitly doing.

Take a concrete, illustrative model and quantify what sharding buys. Constructed scenario — every figure below is derived from stated assumptions, not measured from a real run.

Model: 13 billion parameters, trained with Adam, mixed-precision
  bf16 weights and gradients (2 bytes each), fp32 Adam moments
  (4 bytes each, two moments) -- the same per-parameter recipe used
  in this module's other memory arithmetic.

Per-parameter static memory, unsharded (plain data parallelism):
  bf16 weights        2 bytes
  bf16 gradients       2 bytes
  fp32 Adam moment 1   4 bytes
  fp32 Adam moment 2   4 bytes
  ------------------------------
  total               12 bytes/parameter

Total, unsharded, PER GPU (every rank holds all of it):
  13e9 params x 12 bytes = 156e9 bytes = 156 GB PER GPU

That 156 GB is the static memory bill that plain data parallelism forces onto every single rank in the group, regardless of how many ranks there are — the redundancy problem from section 1 stated in real numbers. Now shard.

Sharding optimizer state AND gradients across 8 data-parallel ranks
(weights stay fully resident on every rank, since the forward pass
needs the complete parameter tensor to run at all):

  bf16 weights (unsharded, full copy every rank):  2 bytes/param
  bf16 gradients (sharded across 8 ranks):          2/8 = 0.25 bytes/param
  fp32 Adam moment 1 (sharded across 8 ranks):      4/8 = 0.5 bytes/param
  fp32 Adam moment 2 (sharded across 8 ranks):      4/8 = 0.5 bytes/param
  ------------------------------------------------------------------
  total per parameter, per rank:                    3.25 bytes/parameter

Total, sharded, PER GPU:
  13e9 params x 3.25 bytes = 42.25e9 bytes ~ 42.25 GB PER GPU

Sharding the gradient and optimizer-state terms across 8 ranks took the per-GPU static footprint from 156 GB down to roughly 42 GB — a reduction of nearly 3.7x for this configuration, without changing the model, the precision recipe, or the optimizer. Note precisely what did not shrink: the weight term stayed at 2 bytes per parameter on every rank, because the forward pass on any given rank still needs the complete, assembled parameter tensor to compute with, and that full tensor has to be resident (or reconstructed via all-gather immediately before use) regardless of how the optimizer state itself is partitioned. Sharding the weights too — the more aggressive FSDP configuration — would shrink that term as well, at the cost of an additional all-gather to reassemble the full weight tensor before every forward pass that needs it.

The reason this lesson exists as its own topic, rather than as a footnote to M7-01, is that a memory error at training time does not announce which of several very different underlying causes produced it, and the wrong fix wastes real engineering time. Take two scenarios that produce the same symptom — an out-of-memory error during training — and walk each one to its correct diagnosis.

Scenario A: "Every rank in our 8-GPU data-parallel job runs out of
memory at the exact same point in the optimizer step, right after the
backward pass finishes. The model itself loaded fine; the forward and
backward passes completed. It's the update that fails."

  Diagnosis: this is an optimizer-state memory problem, not a
  model-too-large problem. The model fits; the redundant, fully-resident
  copies of gradients and optimizer state across all 8 ranks do not.
  Fix: shard optimizer state (and optionally gradients) via ZeRO/FSDP
  across the existing 8-GPU data-parallel group. Tensor or pipeline
  parallelism would not address this -- they solve a different memory
  problem (a layer or the model's depth being too large), and this
  model's layers and depth were never the issue.

Scenario B: "A single transformer block's feed-forward matrices alone
need more memory than one GPU has, even before the optimizer runs --
the forward pass itself fails to allocate."

  Diagnosis: this is a within-layer, model-too-large problem -- the
  `M7-01` symptom that points at tensor parallelism, not at the
  distributed optimizer. Sharding optimizer state would not help here,
  because the failure happens before the optimizer's state is even
  relevant; the layer's own weight tensor does not fit regardless of
  what happens to gradients or Adam moments afterward.

Constructed scenario, with both failure signatures authored to isolate the diagnosis cleanly. The generalizable rule: where in the step the failure occurs is diagnostic. A failure during or immediately after the optimizer step, with the forward and backward passes having already succeeded, points at gradient/optimizer-state memory — the distributed optimizer's territory. A failure during the forward pass itself, before any gradient or optimizer computation has even begun, points at a model-too-large problem that the distributed optimizer does not address at all.

This is the section that resolves the trap the exam's own material names directly. [GROUND TRUTH] (Sources/ncp-genl/domain-7-gpu-acceleration.md): "This is the memory-saving lever when optimizer state doesn't fit — it's a memory-sharding technique layered on data parallelism, not a new parallelism axis," and, restated as its own flagged trap, "FSDP/ZeRO is not a replacement for TP/PP; it's a way to shard state across DP ranks to cut memory."

Compare the distributed optimizer's scope against the six families M7-01 established. Tensor parallelism changes how a layer's forward and backward computation is distributed — different GPUs are doing genuinely different arithmetic on the same layer at the same instant. Pipeline parallelism changes which layers a GPU even holds. Sequence and context parallelism change how the sequence dimension is distributed across the computation. Expert parallelism changes which experts a GPU holds and routes to. Every one of those five is a decision about how to split the model's own computation across devices — a genuinely new axis of "what is different about what each GPU is doing."

The distributed optimizer does none of that. Every rank in a data-parallel group sharded by ZeRO or FSDP is still doing exactly the same kind of computation data parallelism always assigns it — a forward and backward pass over its own slice of the batch, using (after reconstruction via all-gather) the same complete model every other rank in the group is using. What changed is not which computation each GPU performs, but how much of the optimizer's bookkeeping each GPU is required to store at rest between the moments it actually needs the full picture. That is a storage decision riding on top of an existing parallelism strategy, not a new strategy for distributing computation. A scenario whose real constraint is "optimizer state does not fit" is answered by sharding that state across the data-parallel ranks already doing the work — not by introducing tensor or pipeline parallelism, both of which solve a different memory problem (a layer, or the model's depth, being too large) that this scenario never described.

Because the distributed optimizer is scoped to the data-parallel group specifically, it composes cleanly with tensor and pipeline parallelism rather than competing with either. A large-scale training job might combine all three: tensor parallelism sized to fit inside each node's fast interconnect, splitting each layer's matrices across a handful of GPUs; pipeline parallelism assigning different contiguous blocks of layers to different groups of nodes; and, layered on top of that whole tensor-plus-pipeline unit, data parallelism replicating it across still more groups of nodes to process more of the batch simultaneously — with the distributed optimizer sharding the resulting optimizer state specifically across those data-parallel replicas, so the state is not fully duplicated across every one of them.

The scoping question worth asking before reaching for any of these: is the constraint about what computation a GPU needs to be able to run at all (a layer too wide, a model too deep, a sequence too long, an MoE structure with too many experts) — in which case tensor, pipeline, sequence, context, or expert parallelism is the answer — or is the constraint about redundant storage across ranks that are already doing the same kind of work — in which case the distributed optimizer is the answer, and it is applied within whichever data-parallel group already exists, not instead of the other techniques.

Domain 7, GPU Acceleration and Optimization, is 14% of the NCP-GENL blueprint, the second-largest domain after Model Optimization's 17%, and its own scope note requires knowing "how memory-sharding differs from a parallelism axis" as one of the domain's explicit study priorities. Objectives 7.1 through 7.4 ask candidates to configure multi-GPU training and diagnose bottlenecks, and a training run that fails specifically at the optimizer step — succeeding through the forward and backward passes first — is exactly the kind of symptom this domain's scenario questions are built to test, because the correct diagnosis requires distinguishing a storage-redundancy problem from a computation-distribution problem.

How the question tends to be phrased

Expect a scenario naming a specific failure point in the training step ("training fails during the optimizer update, after the backward pass completes without error") and asking which technique addresses it, with the distributed optimizer as the keyed answer against tensor or pipeline parallelism as plausible-sounding but wrong-axis distractors. A second recurring shape states the mechanism directly and asks what it shards: "which technique shards optimizer states across data-parallel GPUs to cut memory?" with the distributed optimizer (ZeRO/FSDP-style) as the answer against tensor parallelism, pipeline parallelism, or expert parallelism as distractors that shard something else entirely or shard nothing at all.

What the distractors typically look like

The house style in this domain offers a real, nameable technique attached to the wrong problem. Expect tensor or pipeline parallelism offered as the fix for an optimizer-state memory error (wrong axis: those solve a computation-distribution problem, not a storage-redundancy one); expect the distributed optimizer described as "a fourth parallelism strategy" or "a replacement for tensor and pipeline parallelism" (the trap this lesson's section 6 resolves directly); and expect a claim that sharding optimizer state also shrinks activation memory, which it does not — activation memory is a separate term entirely, tied to batch size, sequence length, and the model's depth, and unaffected by how optimizer state is partitioned across ranks.

ZeRO is commonly described in three progressively more aggressive stages, and the stage number tracks exactly which of the three static-memory terms — optimizer state, gradients, parameters — is currently being sharded rather than fully duplicated. Stage 1 shards only the optimizer state, leaving gradients and parameters fully resident on every data-parallel rank; this is the smallest memory win but also the smallest change to the communication pattern, since gradients still arrive at every rank in full via an ordinary all-reduce before each rank extracts and updates only its own shard of the optimizer state. Stage 2 additionally shards the gradients themselves, replacing that all-reduce with the reduce-scatter this lesson's section 3 describes, so gradients are never fully reconstructed on any single rank except transiently. Stage 3 goes further still and shards the parameters too, which is the closest ZeRO analog to what FSDP calls fully sharded: at rest, no rank holds the complete weight tensor for any layer it is not actively computing with, and an all-gather reconstructs the needed slice of weights just before each layer's forward or backward computation and discards it again afterward.

FSDP's terminology maps onto the same three-way choice using its own sharding-strategy names rather than ZeRO's stage numbers, but the underlying tradeoff is identical: shard less and pay less communication overhead but keep more memory duplicated; shard more and free more memory at the cost of more frequent reconstruction traffic. Neither framework's naming changes the section 6 conclusion — whichever stage or strategy is selected, the sharded dimension is still the data-parallel one, and the technique is still riding on top of data parallelism rather than replacing it or standing beside it as an independent parallelism axis.

No — the number of data-parallel ranks is still set by the same considerations that determine data parallelism's degree in the first place, chiefly how much you want to divide the batch and how much redundant throughput you want across replicas. What sharding changes is how expensive each existing rank is to keep, not how many ranks the training plan calls for. A team that was already planning an 8-way data-parallel group because that split the batch usefully will still run an 8-way group after adopting ZeRO or FSDP; the difference is that each of those 8 GPUs now stores roughly one-eighth of the optimizer state (and, at more aggressive stages, one-eighth of the gradients or parameters) instead of the full amount. If a team's actual goal is only to reduce per-GPU memory rather than to increase data-parallel throughput, adding more data-parallel ranks purely to get a finer-grained shard is a legitimate lever — more ranks means a smaller shard per rank — but it is a different decision than the one this lesson's mechanism is about, and it comes with its own consequence: a larger data-parallel group also means a larger effective batch size at a fixed per-device batch, which M7-05 covers as its own topic.

• ZeRO/FSDP shard optimizer state — and optionally parameters and gradients too — across data-parallel ranks, eliminating the redundant full copies plain data parallelism otherwise duplicates on every GPU.
• This is a memory-sharding technique layered on data parallelism, not a fifth or sixth parallelism axis alongside M7-01's tensor, pipeline, sequence, context, and expert parallelism — the exam's own material names conflating the two as a specific trap.
• The mechanism is reduce-scatter followed by all-gather, which produces the identical numerical result as a plain all-reduce, factored so no single rank ever needs to hold more than its own shard at rest.
• Where a training failure occurs is diagnostic. A failure at the optimizer step, after forward and backward succeed, points at redundant optimizer-state memory; a failure during the forward pass itself points at a model-too-large problem M7-01's taxonomy addresses instead.
• Weight memory stays fully resident unless parameters themselves are also sharded (the more aggressive FSDP configuration); sharding only gradients and optimizer state leaves weights as the dominant remaining term.
• The technique composes with tensor and pipeline parallelism rather than replacing them — it is scoped specifically to the data-parallel dimension of whatever combined parallelism strategy is already in place.
• Domain 7 is 14% of the NCP-GENL blueprint, and its own scope note explicitly calls out knowing how memory sharding differs from a parallelism axis as a study priority.

Memory sharding answers how to stop paying for redundant optimizer state once a data-parallel group already exists. It says nothing about a separate lever that changes how much memory each individual value costs in the first place, or how fast the arithmetic itself runs once memory is no longer the bottleneck. M7-04 picks that up next: mixed precision and Tensor Cores, where FP16 and BF16 compute cut the bytes-per-value and unlock dramatically higher throughput on modern GPUs, and where a different named trap — FP16 training silently underflowing without loss scaling — is this module's next distinction to get exactly right.