All-Gather and Reduce-Scatter: The Primitives Behind ZeRO and FSDP

"I only ever hold one eighth of the weights. Right before the layer runs, I borrow the other seven eighths, do the math, and politely give them back."

A Shard That Believes It Is the Whole Model

In Section 4.3 we used all-reduce as a single indivisible operation: every worker arrives with a gradient vector, and every worker leaves with the elementwise sum. Section 4.4 then opened the box and showed that the bandwidth-optimal ring algorithm never moves the whole vector to one place; it works in two passes, a reduce phase and a broadcast phase, each touching only $1/K$ of the data per step. Those two phases are not improvised. Each one is a named collective in its own right, useful far beyond all-reduce, and this section gives them their names and their independent uses. The payoff is immediate: once you can reduce-scatter and all-gather separately, you can hold only a fraction of a model on each device and still train it, which is the central trick of memory-efficient large-model training.

1. Two Primitives, and the Identity That Links Them Intermediate

Fix a group of $K$ ranks. Suppose every rank $k$ holds a vector $g_k$ of length $P$, and we want the elementwise sum $s = \sum_{k=1}^{K} g_k$. All-reduce leaves the full vector $s$ on every rank. The two primitives of this section split that goal cleanly in half by deciding who ends up holding what.

Reduce-scatter performs the same elementwise reduction but then keeps each rank's appetite small: it partitions the length-$P$ result into $K$ contiguous slices and delivers slice $k$ to rank $k$ only. After a reduce-scatter, no single rank holds the whole sum; rank $k$ holds just $s$ restricted to its slice, a vector of length $P/K$. Writing $s^{(j)}$ for the $j$-th slice of the summed vector,

All-gather is purely about movement, with no arithmetic. Each rank $k$ starts holding one slice $a_k$ of length $P/K$, and the operation makes every rank end with the full concatenation of all slices in rank order:

Now compose them. Start with the vectors $g_k$, reduce-scatter so that rank $k$ holds $s^{(k)}$, then all-gather those owned slices. Each rank contributes its $s^{(k)}$ and receives the concatenation $[s^{(1)} \Vert \cdots \Vert s^{(K)}]$, which is exactly the full sum $s$. Every rank now holds $s$. That is the definition of all-reduce, so

This is not a loose analogy; it is an equality of results, and it is precisely the two-phase structure of the ring all-reduce in Section 4.4. The reduce phase of the ring is a reduce-scatter (each rank ends owning the fully reduced value of one chunk), and the broadcast phase is an all-gather (each rank's owned chunk is propagated to everyone). Both phases move $P(K-1)/K$ elements per rank, so an all-reduce costs about twice a single all-gather, a fact we will spend later in the FSDP cost argument.

2. Verifying the Identity in Code Intermediate

The cleanest way to trust the identity is to implement both primitives on simulated shards and check that their composition reproduces a direct all-reduce, number for number. The code below represents the $K$ ranks as rows of a single array (so we can see everything at once on one machine), implements reduce-scatter and all-gather as plain array operations, and compares reduce-scatter then all-gather against a direct sum. It then plays the FSDP cycle in miniature: shard a full weight vector, all-gather it back to compute, and reduce-scatter a batch of per-rank gradients so each rank keeps only its slice.

import numpy as np

rng = np.random.default_rng(0)
K, P = 4, 12                      # workers, parameter-vector length
g = rng.standard_normal((K, P))   # g[k] = the gradient computed on worker k

# --- reduce-scatter: sum across workers, then each rank keeps only its slice ---
def reduce_scatter(g):
    K, P = g.shape
    assert P % K == 0
    total = g.sum(axis=0)                       # the summed (reduced) full vector
    slices = total.reshape(K, P // K)           # slice s belongs to rank s
    return slices                               # out[k] = rank k's owned shard

# --- all-gather: each rank contributes its shard; everyone ends with the concat ---
def all_gather(shards):
    return np.concatenate(shards, axis=0)       # length-P vector, identical on all ranks

# --- the identity: all-reduce == reduce-scatter then all-gather ---
owned = reduce_scatter(g)                       # K shards, one per rank
allgathered = all_gather(owned)                 # full summed vector, on every rank
allreduce_ref = g.sum(axis=0)                   # direct all-reduce (SUM)

print("workers K, params P        :", K, P)
print("reduce-scatter shard shape :", owned.shape, "(K shards of P/K)")
print("max abs diff vs all-reduce :", f"{np.max(np.abs(allgathered - allreduce_ref)):.2e}")
print("identity holds (RS+AG==AR) :", np.allclose(allgathered, allreduce_ref))

# --- FSDP cycle in miniature: shard -> all-gather to compute -> reduce-scatter grads ---
full_w = rng.standard_normal(P)                 # the layer's true full weight
w_shards = full_w.reshape(K, P // K)            # each rank stores only its shard
gathered_w = all_gather([w_shards[k] for k in range(K)])
print("FSDP all-gather rebuilds w :", np.allclose(gathered_w, full_w))
per_rank_grad = rng.standard_normal((K, P))     # each rank's grad for the FULL layer
kept = reduce_scatter(per_rank_grad)            # after backward, keep only your shard
print("each rank keeps P/K grads  :", kept.shape[1], "of", P)

workers K, params P        : 4 12
reduce-scatter shard shape : (4, 3) (K shards of P/K)
max abs diff vs all-reduce : 0.00e+00
identity holds (RS+AG==AR) : True
FSDP all-gather rebuilds w : True
each rank keeps P/K grads  : 3 of 12

The difference is not merely small; it is exactly zero, because both primitives only add and copy floating-point numbers that the direct all-reduce also adds in the same grouping. The last two lines are the heart of the next part: each rank can store a third of a weight, briefly borrow the rest to compute, and after the backward pass discard everything except its own third of the gradient. Nothing about correctness changed; only the peak memory each rank must hold changed, and dramatically.

3. Sharded Data Parallelism: ZeRO and FSDP Advanced

Plain data parallelism (Chapter 15) replicates the entire model on every worker: every rank holds a full copy of the parameters, the gradients, and, expensively, the optimizer state. For an Adam-trained model in mixed precision, the optimizer state and master weights can be several times the size of the parameters themselves, so replication wastes most of the memory on every device holding identical copies. The insight of ZeRO (Zero Redundancy Optimizer) and its PyTorch sibling FSDP (Fully Sharded Data Parallel) is that this redundancy is unnecessary: with $K$ ranks, shard each of those tensors into $K$ pieces and have rank $k$ permanently own only piece $k$. No rank holds a full copy of anything large at rest.

The cost of holding only a shard is that the full weights of a layer do not exist on any single device when that layer needs to run. The two primitives of this section supply them on demand, and the rhythm is a tight cycle repeated layer by layer:

1. All-gather the weights. Just before a layer's forward pass, the group all-gathers that layer's parameter shards so every rank momentarily holds the full layer weights and can compute the layer's output. The gathered copy is freed as soon as the layer finishes, so only one layer's full weights ever sit in memory at once.
2. All-gather again for backward. The same all-gather reconstructs the layer's full weights during the backward pass so each rank can compute the gradient with respect to the complete weight tensor.
3. Reduce-scatter the gradients. After backward, each rank holds a full-length gradient for the layer (summed contributions from its own microbatch). A reduce-scatter sums these across ranks and leaves rank $k$ holding only slice $k$ of the averaged gradient, exactly the slice whose optimizer state and parameters rank $k$ owns. Rank $k$ then updates its shard locally.

Figure 4.5.1 draws both halves of the story: the identity on top (reduce-scatter plus all-gather reconstituting an all-reduce) and the FSDP gather-compute-scatter cycle below. Notice that plain data parallelism's single gradient all-reduce per step has been replaced by, per layer, two all-gathers (forward and backward) plus one reduce-scatter. Since an all-reduce already equals one all-gather plus one reduce-scatter, FSDP communicates roughly one-and-a-half times the volume of plain data-parallel training. That extra traffic is the price, and in return each rank's resident memory for parameters, gradients, and optimizer state falls by a factor of about $K$.

# Run with: torchrun --nproc_per_node=4 thisfile.py
import torch, torch.distributed as dist
from torch.distributed.fsdp import FullyShardedDataParallel as FSDP

dist.init_process_group("nccl")
rank, world = dist.get_rank(), dist.get_world_size()

# --- raw primitives ---
shard = torch.ones(3, device="cuda") * rank          # this rank's slice
full  = [torch.empty(3, device="cuda") for _ in range(world)]
dist.all_gather(full, shard)                          # every rank now has all slices

grad  = torch.randn(3 * world, device="cuda")        # full-length gradient
out   = torch.empty(3, device="cuda")
dist.reduce_scatter_tensor(out, grad, op=dist.ReduceOp.SUM)  # keep only my slice

# --- or let FSDP schedule both around every layer ---
model = FSDP(my_model)                                # shards params/grads/optimizer
# forward/backward now all-gather weights and reduce-scatter grads under the hood

4. When the Trade Pays, and When It Does Not Intermediate

Sharding is not free, and the same cost discipline from Chapter 3 decides whether it helps. FSDP adds about half again the communication of plain data parallelism, so its benefit appears only when memory is the binding constraint. If a model already fits comfortably with room for the optimizer state, plain replication is faster because it communicates less. As the model grows past one device's memory, replication simply stops working, and FSDP's extra traffic becomes the price of training at all rather than a slowdown over a working alternative. The interesting regime is the middle, where the model nearly fits: there, partial sharding (shard the optimizer state but replicate the parameters, the lightest ZeRO stage) often wins by cutting memory enough to fit while adding the least communication.

Two levers govern the trade. The first is interconnect speed: the all-gathers and reduce-scatters are bandwidth-bound, so a fast intra-node fabric makes sharding nearly free while a slow cross-node link makes it punishing, which is why FSDP shines within a node and is often combined with other parallelism strategies across nodes (Chapter 16). The second is overlap: because each layer's weight all-gather can be issued before the layer's compute begins, a good implementation prefetches the next layer's weights while the current layer runs, hiding most of the communication behind computation, a theme Section 4.10 develops for gradients.

We now have all-reduce decomposed into its two working halves and have seen those halves carry the entire weight of memory-efficient large-model training. Reduce-scatter and all-gather move data along the contiguous-slice partition we chose; the next section asks what happens when the destinations are not slices of one vector but entire messages that must each go to a different rank, the all-to-all pattern that routes tokens to experts in a mixture-of-experts model. That is the subject of Section 4.6.