If you learned distributed training through DDP, you probably carry two instincts: after the backward pass, all-reduce the gradients; and if only one place has the freshest weights, broadcast them out. I carried both into FSDP and they cost me real confusion, because both are wrong there. Not slightly wrong, wrong in a way that means the mental model underneath is wrong. Working out why fixed my understanding of FSDP more than anything else, so this post is that explanation: what all-gather and reduce-scatter actually do, why reduce-scatter specifically is the right collective after backward, and why broadcast and all-reduce are answers to questions FSDP never asks.

Before the story starts, the two operations themselves, in plain terms, no FSDP attached. All-gather: every rank contributes its piece of a tensor, and afterwards every rank holds the complete tensor. Reduce-scatter: every rank contributes a full size tensor, the tensors get combined element wise (averaged, for our purposes), and each rank keeps only its own slice of the result. One assembles pieces, the other merges disagreeing copies and deals out the shares. That’s the entire vocabulary of this post (the formal definitions live in NCCL’s collective operations docs). Everything that follows is about why FSDP uses exactly these two, at the moments it does, and not the collectives you might reach for instead.

Here’s the same vocabulary as a picture. Two GPUs, four numbers, A responsible for the first half and B for the second. Notice the mirror: one op goes small in, big out; the other goes big in, small out.

all-gather reduce-scatter before before A 1 2 B 3 4 all-gather after A 1 2 3 4 B 1 2 3 4 small in, big out no math, pure assembly A 8 0 4 2 B 0 4 8 6 reduce-scatter avg = [4,2,6,4], in flight only after A 4 2 B 6 4 big in, small out merges disagreeing copies, deals out slices

Solid cells are what a rank contributed, pale cells are what arrived over the wire, and dashed cells hold nothing. The same numbers show up again below, when these two ops go to work inside FSDP.

One scope note as well: everything below describes plain one dimensional full sharding, FSDP2’s default, written against PyTorch 2.11; the low level details, especially gradient scaling and how the collectives get scheduled, can shift between releases. Hybrid sharding adds a replica dimension on top, and with it extra communication (including, yes, an all-reduce). That’s a different post.

If you only take three lines from this post:

  1. Parameter shards are complementary pieces of one true weight, so using them takes an all-gather.
  2. Gradients are full size but different answers, computed from different data, so they need averaging: a reduce.
  3. Each rank only updates its own slice, so reduce-scatter does the reduction and the delivery in one op.

Two different worlds

In DDP, every GPU permanently stores the entire model. 8 GPUs means 8 full copies that have to stay bit identical forever. All of DDP’s communication exists to keep those copies in sync.

In FSDP, the model exists exactly once, chopped into W pieces. GPU k permanently owns piece k of every weight tensor, and of its gradient and optimizer state too. No full copy of anything exists anywhere at rest. Full size tensors only show up as short lived photocopies during compute, and then they get shredded. (If you’re already asking “why doesn’t gathering full tensors blow up memory?”, good question. Hold it until the next section.)

Once this picture is in your head, every “which collective goes here?” question answers itself. You just ask: in this world, who is allowed to permanently hold what? Both of the DDP instincts above are symptoms of the same bug: imagining full copies that need to be kept in sync, in a world that deliberately has none.

All-gather: everyone shows their piece

Each of W ranks contributes its piece, and afterwards everyone holds the concatenation of all the pieces.

Say we have two GPUs and one 4 element weight. A owns [w1,w2] = [1, 2] and B owns [w3,w4] = [3, 4]:

before:   A: [1, 2, ., .]        B: [., ., 3, 4]
all-gather ------------------------------------
after:    A: [1, 2, 3, 4]        B: [1, 2, 3, 4]

FSDP needs this because a matmul touches every entry of the weight. So right before a layer runs, its custodians pool their slices into a temporary full copy. Compute, shred, move on to the next layer. The whole forward pass is just gather, use, shred, repeated per layer.

Notice the direction: each rank starts with 1/W of the data and ends with all of it. Small in, big out. And there’s no arithmetic anywhere, it’s pure assembly.

I’ll keep saying “layer” because it reads better, but strictly the unit is the FSDP communication group: whatever you wrapped in one fully_shard() call. Wrap per transformer block, the common setup, and “group” and “layer” mean the same thing.

This is also where the OOM question from earlier gets its answer. The photocopies don’t blow up memory because they never all exist at once: only the layer currently computing is unsharded, plus the next one, fetched early to hide latency. For a 5B model split into 24 blocks, one block’s full bf16 weights are about 0.4 GB, so with the prefetched block in flight call it under a gigabyte alive at any moment, against about 10 GB if every block stayed gathered. So the memory spike scales with your largest block or two, not with the model. And “short lived” is literal: a block’s photocopy exists for the few milliseconds its compute takes, then the buffer is recycled for the next block. The fine print is that this guarantee comes from how you wrap. Call fully_shard() only on the root and there’s one group, the whole model becomes one giant photocopy, and that can absolutely OOM.

Why backward needs a different collective

Here’s the question that unlocked this for me. FSDP really has two communication jobs, attached to two different things. Parameters get all-gathered whenever compute needs them in full: before a layer’s forward, and, because the photocopy gets shredded right after forward, usually again before that layer’s backward. Gradients get reduce-scattered once backward has produced them. So the real split isn’t “forward vs backward”, it’s parameters vs gradients: why do parameters gather while gradients reduce? Where does the “reduce” suddenly come from?

Look at what the ranks are holding in each case.

When parameters move, the shards are complementary pieces of one true weight. A’s [1, 2] and B’s [3, 4] don’t disagree about anything; they’re different chapters of the same book. Assembling them takes concatenation and nothing else. No arithmetic, so no reduce. All-gather.

When gradients move, the situation is completely different. Each GPU ran the same weights on different data, so each holds a full size gradient and the copies disagree:

A's grad:  [8, 0, 4, 2]
B's grad:  [0, 4, 8, 6]
average:   [4, 2, 6, 4]   <- computed in flight, never assembled on any GPU

Disagreeing copies can’t be concatenated, they have to be combined. That combining step is the “reduce”. So the rule that generalizes: reduce shows up exactly when the per-rank copies disagree and must be merged. Parameters never disagree, there’s one true weight living in pieces. Gradients disagree in general, because each rank saw different data. That’s the whole reason the two use different collectives.

And the “scatter” half? After averaging, each rank only needs its own slice. A will only ever update w1 and w2, so shipping it the averaged gradient for w3 and w4 would be spending network bandwidth on numbers it throws away. Reduce-scatter does both at once: averages everyone’s full gradients and delivers each custodian just its slice. A gets [4, 2], B gets [6, 4], and the full averaged gradient never exists on any single GPU.

Direction-wise this is the exact mirror of all-gather: big in, small out. And the reduction you want is an average, not a sum. For bf16 and fp32, NCCL’s AVG op folds the divide by W into the collective itself, no separate division kernel. Other dtypes take slightly different routes (fp16 splits the divisor across pre and post scaling to avoid overflow), but every route ends the same place: each rank holds its shard of the averaged gradient. All of this is visible in PyTorch’s FSDP2 collectives source if you want to see the machinery.

NCCL’s bf16 reduction also accumulates in bf16 along the way, which starts to get lossy as the world size grows. That’s why FSDP2’s mixed precision policy lets you compute in bf16 but reduce in fp32 (reduce_dtype), trading twice the reduce-scatter bytes for numerical safety. That trade is a bandwidth story, and it belongs to the next post.

A terminology note, since “scatter” and “sharding” sound interchangeable: sharding is a state, scatter is an action. Think of a card game. Dealing a card to each player is a scatter. Each player holding their own hand is being sharded. FSDP’s weights are sharded (the standing layout); reduce-scatter is the verb that re-establishes that layout for gradients, with an average folded in.

DDP’s all-reduce, chopped in half

There’s a connection here: all-reduce = reduce-scatter + all-gather. “Everyone ends up with the full averaged tensor” breaks into “everyone gets their averaged slice” followed by “everyone shows their slice”.

DDP needs the full all-reduce on gradients, since every GPU stores full weights and therefore needs the full averaged gradient. FSDP runs only the first half after backward. Each rank updates only its slice, so reduce-scatter is enough. The second half isn’t skipped though. It moves to the next forward pass, where an all-gather was needed anyway to build the photocopy.

The same connection also tells you the cost. The standard ring all-reduce is literally these two ops run back to back, so stopping at reduce-scatter moves half the bytes: FSDP’s gradient sync costs half of DDP’s all-reduce on the wire. The other half of the traffic comes back later as the parameter all-gather, paid at the moment it’s useful. A bandwidth win as well as a memory one.

So FSDP is DDP’s all-reduce chopped in half, with each half moved to where the data is actually needed. Nothing new gets invented. The pieces just run at different times.

To be precise, the second half doesn’t carry the same tensor. The optimizer steps in between, so the later all-gather moves updated parameter shards, not the gradient shards that came out of reduce-scatter. What gets chopped in half is the communication pattern, not one particular tensor.

Why not broadcast?

Broadcast means “one rank has the truth, copy it to everyone”. The DDP instinct says: the optimizer just updated the weights, other ranks need them, broadcast. But broadcast exists to fix stale copies, so count the copies. After the optimizer updates w3, how many permanent copies of w3 exist? Exactly one, on its custodian, freshly updated. There is nothing to be stale. The other ranks don’t hold an outdated w3. They hold nothing at all, because their photocopy was shredded after backward. They’ll get the fresh w3 automatically at the next forward’s all-gather, straight from the one rank that owns it.

Broadcast is DDP thinking. It assumes replicas that can drift apart. When there’s one original per weight, staleness isn’t a thing you have to prevent. It just can’t happen.

Why not the other collectives?

Each collective answers a specific question. “Which one goes here?” really means “which question is FSDP asking right now?”

Collective The question it answers Does FSDP ask it?
broadcast one rank knows, everyone needs a copy no, there are no replicas to sync
scatter one rank holds everything, deal out the pieces no, pieces never start centralized
gather collect all pieces onto one rank no, every rank needs the full layer, not just one
all-gather everyone has a piece, everyone needs the whole yes, before every layer’s compute
all-reduce everyone has a full version, everyone needs the full average no, that’s DDP
reduce-scatter everyone has a full version, each rank needs its slice of the average yes, after gradients

Look at the “no” rows. They all either assume a central rank (scatter, gather, broadcast) or assume full replicas (broadcast, all-reduce). FSDP’s world has neither. Every rank is a custodian of equal standing, and there’s exactly one original of everything.

Run it yourself

Don’t take my word for the numbers. This reproduces every value in this post with the torch.distributed API directly, no FSDP involved:

import torch
import torch.distributed as dist

use_cuda = torch.cuda.is_available()
dist.init_process_group("nccl" if use_cuda else "gloo")
rank = dist.get_rank()
if use_cuda:
    torch.cuda.set_device(rank)
dev = torch.device("cuda", rank) if use_cuda else torch.device("cpu")

# all-gather: A contributes [1,2], B contributes [3,4]
shard = torch.tensor([1.0, 2.0] if rank == 0 else [3.0, 4.0], device=dev)
full = torch.empty(4, device=dev)
dist.all_gather_into_tensor(full, shard)
print(f"rank {rank} after all-gather:     {full.tolist()}")

# reduce-scatter: A contributes [8,0,4,2], B contributes [0,4,8,6]
grad = torch.tensor([8.0, 0.0, 4.0, 2.0] if rank == 0 else [0.0, 4.0, 8.0, 6.0], device=dev)
mine = torch.empty(2, device=dev)
if use_cuda:
    dist.reduce_scatter_tensor(mine, grad, op=dist.ReduceOp.AVG)
else:
    dist.reduce_scatter_tensor(mine, grad, op=dist.ReduceOp.SUM)
    mine /= dist.get_world_size()
print(f"rank {rank} after reduce-scatter: {mine.tolist()}")

dist.destroy_process_group()

Save it as collectives_demo.py and run:

torchrun --nproc_per_node=2 collectives_demo.py

Output:

rank 0 after all-gather:     [1.0, 2.0, 3.0, 4.0]
rank 1 after all-gather:     [1.0, 2.0, 3.0, 4.0]
rank 0 after reduce-scatter: [4.0, 2.0]
rank 1 after reduce-scatter: [6.0, 4.0]

It runs on two GPUs over NCCL, or on plain CPU over gloo, and the fallback branch is a small lesson in itself: ReduceOp.AVG is NCCL only, so on CPU you sum and divide yourself. Verified on PyTorch 2.11.0.

The check that fixed my mental model

When I’m not sure which collective belongs somewhere, I stop and count who permanently holds what. If my answer requires a full copy of anything sitting on a GPU at rest, I’ve slipped back into DDP world. In FSDP world:

  • permanent state is shards only: weights, grads, optimizer moments, all 1/W
  • full tensors are photocopies that live for one layer’s compute
  • there’s one original of every number, so “keeping copies in sync” isn’t a concept

And if the copies flowing through a collective disagree with each other, expect a reduce in its name. If they’re complementary pieces of one thing, expect a gather.

Everything else in FSDP, the CUDA streams, prefetching, reshard_after_forward, overlap, is engineering on top of one follow up question: the gathers cost time, can we hide them behind compute? That’s a coming post, with real profiler traces.

References

The mechanism claims in this post are checked against the PyTorch 2.11.0 source, and the memory numbers come from my own 8 GPU H100 runs (scripts will ship with the next post).