Distributed Data Parallel and FSDP from Scratch
Multi-rank training is two collectives and one rule. Broadcast the parameters at startup, average the gradients after backward, never let the ranks disagree about what step they are on. Bring up a process group across N ranks with the gloo backend, no special hardware. Implement a minimal DDP wrapper that broadcasts parameters at construction and all-reduces gradients after backward. Prove that the all-reduce of per-rank gradients matches a single-process gradient on the concatenated input. Sketch FSDP parameter sharding: each rank holds a slice, the full tensor is gathered for the forward pass and dropped after. The model fits on one device. The dataset does not. The optimization budget says you want to see N times the examples per wallclock second. The first lever is data parallel: each rank runs the same model on a different slice of the batch, then averages gradients before the optimizer step. The second lever is FSDP: the model does not fit on one device either, so each rank holds a fraction of every parameter and reconstructs the full tensors layer by layer during the forward pass. The pain is the 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.…
Distributed Data Parallel and FSDP from Scratch: Multi-rank training is two collectives and one rule. Broadcast the parameters at startup, average the…
This free lesson is part of the AI Engineering from Scratch curriculum. Read the full explanation, run the lesson code, and verify the result in the interactive reader or from the repository source.
Browse the complete course catalog or open this lesson on GitHub.