Distributed Data Parallel and FSDP from Scratch

Data parallel training means each rank runs the same model on a different slice of the batch, then averages gradients before the optimizer step. FSDP (fully sharded data parallel) goes further: each rank holds a fraction of every parameter and reconstructs the full tensors layer by layer during the forward pass.

Here is when you reach for each. The model fits on one device, but the dataset is large and you want N times the examples per wallclock second. That is data parallel. If the model does not fit on one device either, that is FSDP.

The pain is bookkeeping. If parameters drift across ranks, the run is silently corrupt. If you average gradients but not the loss, the dashboard lies. If the collective backend cannot agree on a topology, the run hangs forever. My fix is to write the collectives by hand once, so you are not trusting a wrapper you cannot reproduce.

This lesson runs on CPU, with no CUDA assumed. The gloo backend ships with every PyTorch build and accepts torch.multiprocessing workers. The same code switches to nccl on a multi-GPU node without changing structure.

A collective is an operation that every rank in a process group takes part in, and the DDP contract is built from two of them. The sections below say which collective runs when.

The two collectives that matter

The DDP contract is broadcast at construction and all_reduce after backward, which MinimalDDP exposes as a sync_grads() method the training loop calls explicitly. The FSDP sketch adds all_gather before each layer's forward pass.

Gradient averaging matches single-process gradient

A model trained on a batch of B examples across N ranks must produce the same gradient as a single process training on a batch of N*B. Summing per-rank gradients and dividing by N gives the average loss gradient, which is what cross entropy with mean reduction would produce on the full batch. The code asserts this with max-abs-diff < 1e-4 between the manual all-reduce gradient and the reference single-process gradient. The looser 1e-3 threshold belongs to a different check, param_sum_spread across ranks (Exercise 1).

FSDP sketch

The memory win is exact: per-rank memory for parameters drops to 1/N. The cost is the gather, which is paid on each forward pass. Production FSDP overlaps the gather with the previous layer's compute, so the wallclock cost is much smaller than the naive accounting predicts. The code does the round-trip on every parameter and asserts with torch.allclose that the reconstruction matches the original.

CPU and the gloo backend

CUDA is the production target, but the same code paths exist on CPU. gloo is the CPU collective backend. It is slower than nccl on GPUs by orders of magnitude, but the API surface is identical. The process group is initialized with backend="gloo" and ranks are spawned with torch.multiprocessing rather than torchrun. Both end up at the same torch.distributed calls. On a multi-GPU node, the only changes are backend="nccl", device tensors, and torchrun to launch.

code/main.py is the runnable artifact.

Step 1: bring up the process group

os.environ["MASTER_ADDR"] = "127.0.0.1"
os.environ["MASTER_PORT"] = str(port)
dist.init_process_group(backend="gloo", rank=rank, world_size=world_size)

MASTER_ADDR and MASTER_PORT are the rendezvous: every rank dials the same port on the same host. The demo also points GLOO_SOCKET_IFNAME and TP_SOCKET_IFNAME at the loopback interface so the ranks talk over localhost. A bind-and-close trick picks a free port, which avoids collisions when several runs share a machine.

Step 2: broadcast at construction

MinimalDDP.__init__ walks every parameter and buffer and calls dist.broadcast(tensor, src=0). Rank 0's values become the canonical init. Skip this and each rank initializes from its own seed, so the ranks diverge from step one.

Step 3: all-reduce gradients after backward

def all_reduce_grads_(module, world_size):
    for p in module.parameters():
        if p.grad is None:
            p.grad = torch.zeros_like(p.data)
        dist.all_reduce(p.grad.data, op=dist.ReduceOp.SUM)
        p.grad.data.div_(world_size)

MinimalDDP.sync_grads() wraps this function, and the training loop calls it explicitly between loss.backward() and optimizer.step(). Every rank ends up with the same averaged gradient. The optimizer step then sees the same input on every rank, which is why the parameters stay in sync across the run.

Step 4: prove the equivalence

manual_all_reduce_matches_single_process builds the same model on rank 0 and compares the post-all-reduce gradient against the gradient a single process would compute on the concatenated input. The demo asserts the max-abs-diff stays well under the 1e-4 threshold.

Step 5: FSDP round trip

fsdp_round_trip_sketch flattens each parameter, pads to a multiple of world_size, slices, all-gathers, and unpads. Every rank's reconstruction matches the original under torch.allclose. That is the unshard step. The inverse (re-shard after the forward) is one slice off the gathered tensor.

Run it:

python3 code/main.py

Default world size is 2. Two CPU processes spawn, talk to each other through gloo, and exit zero. The output outputs/ddp-demo.json captures parameter sums per rank, the gradient norm after all-reduce, the FSDP round-trip result, and the manual-vs-reference gradient diff.

Production training stacks call the same primitives. PyTorch's DistributedDataParallel adds post-backward gradient hooks that overlap all-reduce with backward, bucketed all-reduce that combines several small gradients into one collective, and the no_sync context lesson 46 used.

PyTorch's FSDP adds a flat parameter view per layer, so each rank holds one contiguous buffer. It also overlaps the next layer's unshard with the current layer's compute, and offers optional CPU offload for the shards.

The shape stays the same: broadcast at startup, reduce after backward, shard parameters when they no longer fit.

outputs/skill-distributed-fsdp-ddp.md carries the recipe for a new training script. Spin up the process group with gloo for CPU and nccl for GPU. Wrap the model in a DDP shell that broadcasts at construction and reduces after backward. Optionally shard parameters with the all_gather pattern from the FSDP sketch.

1. Run with --world-size 4 and confirm the param spread stays under 1e-3 across the run.
2. Replace the manual averaging with dist.all_reduce(op=dist.ReduceOp.AVG) and time the difference.
3. Add a post-backward hook to the DDP wrapper so the all-reduce overlaps with the rest of the backward; measure the wallclock improvement.
4. Implement the FSDP re-shard step: after the forward pass, replace the full tensor with the local shard again. Confirm per-rank memory drops.
5. Switch the backend to nccl on a CUDA box. Note which environment variables change and which stay the same.

• PyTorch torch.distributed documentation for the collective semantics this lesson relies on.
• The gloo library's collective list, identical in shape to the CUDA-backed nccl primitives.
• Phase 19 lesson 46 for the gradient accumulation pattern that wraps the DDP all-reduce in no_sync.
• Phase 19 lesson 47 for the checkpoint layout that survives DDP and FSDP runs.
• PyTorch FSDP documentation for the production implementation of the parameter sharding sketched here.