CapstonePython

Data Parallel DDP From Scratch

DistributedDataParallel (DDP) is a thin wrapper around allreduce. Wrap a model, broadcast the initial parameters from rank 0 so every rank starts identical, and allreduce the gradients with an explicit call after backward. The rest is plain gradient descent, and the whole pattern fits in about 220 lines.

TIME: ~90 minPREREQUISITES: Phase 19 Track C lessons 42-49

Learning Objectives

Wire a `DistributedDataParallel`-shaped wrapper that broadcasts initial parameters and allreduces gradients after backward.
Spawn N CPU ranks with `mp.get_context("spawn")` and `ctx.Process` over the gloo backend with file-based rendezvous.
Prove gradient-sync correctness by training the same model on the same data sequentially and showing matching per-step losses and a matching final parameter norm.
Defend the use of buckets (gradient fusion) and overlap (comm during backward) as the two changes that turn a working DDP into a production DDP.

The Problem

Data parallel splits the batch across N ranks. Each rank runs forward and backward on its own shard, and at each step the ranks' gradients are summed so all N copies stay identical. The optimiser steps on that summed gradient.

The reason to bother: a 1-billion-parameter model with 12 GB of activations does not fit on one consumer GPU, and even when a model fits, training can take weeks.

Skip the gradient sync and the N replicas diverge after the first update. You no longer have one model trained on more data, you have N separate models that happen to share initial weights. Do the sync badly (one allreduce per parameter, no overlap, no bucketing) and the network becomes the bottleneck while the GPUs idle waiting on the wire. The craft of DDP is making gradient sync nearly free relative to compute. Canonical PyTorch DDP gets there by bucketing gradients, overlapping allreduce with the next layer's backward, and using NCCL on NVLink. The code here is the naive version on CPU with gloo, one allreduce per parameter with no buckets and no overlap, and exercise 1 adds bucketing so you can measure the difference.

Warm-up0/3

1. Four workers each compute a gradient on a different quarter of one batch. How do you combine them to match the gradient from the full batch under a mean loss?Warm-up

2. Two runs use identical code, the same random seed and the same data. What do you expect for the freshly initialised weights of the model in each run?Warm-up

3. A script starts four separate operating-system processes on one machine. What do the processes not share by default?Warm-up

Check: The Problem0/4

1. What does "overlap" mean in production DDP?Define

2. If the gradient sync step is left out, what happens to the N replicas as training proceeds?Explain

3. Why does the lesson call the code in main.py the naive version of DDP?Explain

4. Profiling a large transformer with one allreduce per tensor and no overlap shows GPUs idle while the network is busy. Which two changes does the lesson say fix this?Analyze

The Concept

Mermaid diagram

The three operations DDP needs

Stage	Collective	Why
Init	broadcast from rank 0	Every rank starts with the same parameters
After backward	allreduce of each grad	The mean gradient is what the optimiser steps on
Sometimes	broadcast of buffers	Batchnorm running stats stay synchronised (production DDP only; `main.py` has none because MiniMLP has no buffers)

Why mean and not sum

The mean gradient is allreduce-SUM divided by world_size. It keeps the same magnitude as world_size changes, so a learning rate tuned at one rank works at four. Leave out the division and you retune the learning rate every time the cluster size changes. DDP wraps the SUM and divides, and the code here does the same.

Why bucket gradients

A bucket is a group of gradient tensors reduced with a single allreduce. A transformer has thousands of parameter tensors, and one allreduce per tensor pays the gloo latency floor thousands of times. DDP groups gradients into ~25 MB buckets and issues one allreduce per bucket. The same total bytes cross the wire, but the latency is amortised over the bucket. The code here does not bucket: sync_grads issues one allreduce per parameter tensor (six for MiniMLP), and exercise 1 adds buckets.

Why pin the seed

Every rank calls torch.manual_seed(SEED) before building the model, so initial parameters start identical, and the broadcast from rank 0 guarantees it. Data is not shuffled per rank. Each rank slices its own shard out of one shared synthetic dataset by offset (step * world_size * batch + rank * batch), so ranks see different rows without needing a rank-specific seed. Seeding parameter init per rank would make the replicas start apart, and only the broadcast would repair that. Skip both the shared seed and the broadcast and the ranks start from different weights, so their step-0 losses already disagree with the reference. (Rank 0 keeps the reference seed if the others only vary theirs, so it first drifts at step 1.)

Check: The Concept0/9

1. Four ranks each hold a gradient tensor and call allreduce with SUM. What does each rank hold afterward?Define

2. What is a bucket in production DDP?Define

3. Why does DDP broadcast parameters at construction instead of relying only on the same seed across ranks?Explain

4. Why does DDP divide the allreduce-summed gradient by world_size?Explain

5. Bucketing sends the same total bytes as one allreduce per parameter. What does it save?Explain

6. You tuned a learning rate on one rank, then moved to 4 ranks using allreduce-SUM with no division. What do you expect?Apply

7. In _ddp_worker with world_size=4 and batch=8, which rows of the dataset does rank 2 slice at step 3?Apply

8. After a refactor, DDP rank 0 matches the reference loss at step 0 but differs from step 1 onward, and the final norms disagree. Which change fits these symptoms?Analyze

9. A teammate changes the worker to torch.manual_seed(SEED + rank) but leaves the broadcast in DistributedDataParallel. What do you expect?Analyze

Build It

code/main.py implements five pieces:

MiniMLP: a 3-layer MLP small enough to converge in seconds, large enough to expose the wiring.
DistributedDataParallel(model, world_size): a thin wrapper, not an nn.Module. It broadcasts params at construct time, and its sync_grads() divides the allreduce-summed grads by world_size. The training loop calls ddp.sync_grads() explicitly after loss.backward(); there are no backward hooks.
_ddp_worker(rank, world_size, init_file, iface, steps, batch, lr, out_queue): the full training loop, with torch.distributed init over gloo, then forward, backward, sync, step.
run_ddp(...): starts the ranks with mp.get_context("spawn") and ctx.Process, using a file-based rendezvous, and collects each rank's losses and final parameter norm.
reference_single_process(...): trains the same model on the same data sequentially in one process, walking every rank's micro-batch in rank order each step. The printed table compares its losses to the DDP run.

Run it:

bash

python3 code/main.py

The output is a per-step table with the reference loss, the DDP loss on rank 0, the DDP loss on rank 3, and the drift between those two ranks, followed by the final parameter norm for the reference and for rank 0. The two paths produce matching loss curves and a matching final norm to float epsilon, which shows the gradient sync is correct. The code does not compare parameters after every step.

Check: Build It0/10

1. What does the DDP wrapper in this lesson do after loss.backward() and before optimizer.step()?Define

2. Why does the training loop call ddp.sync_grads() explicitly after loss.backward()?Explain

3. What does run_ddp use to start the ranks and let them form one process group?Explain

4. Why does the reference implementation walk every rank's micro-batch in rank order within one process?Explain

5. How many allreduce calls does sync_grads issue per step for MiniMLP in main.py?Apply

6. Four ranks build MiniMLP with different weights, then each calls DistributedDataParallel(model, 4). Which weights does every rank hold afterward?Apply

7. In reference_single_process, each micro-batch loss is divided by world_size before backward(). After the four backward calls of one step, what does each parameter's .grad hold?Apply

8. run_ddp starts each rank with args (r, world_size, init_file, iface, steps, batch, lr, out_queue). Which of these differs from rank to rank?Apply

9. The table shows rank_drift of 6.22e-01 at step 0, yet the final parameter norms match. A reviewer calls it a sync bug. What is the right reading?Analyze

10. You write a bucketed version of sync_grads and need evidence that it is correct. Which test follows the lesson's approach?Design

Production patterns in the wild

Three patterns harden DDP enough to ship.

Find unused parameters. Some forward paths skip parameters conditionally (early exit, a mixture-of-experts router). Those parameters get no gradient, but DDP's bucket-ready hook still waits for them and the allreduce deadlocks. find_unused_parameters=True makes DDP check which params got gradients before reducing. The cost is a graph walk per step, so leave it off unless your forward branches.

Static graph optimisation. When the forward is stable across steps, static_graph=True lets DDP precompute the bucket schedule. Precomputing saves a few ms per step, and that compounds across 10000 steps.

Gradient accumulation needs care. Accumulating gradients over K microbatches without syncing each one is a 10x throughput win. DDP exposes no_sync() as a context manager that pauses the post-backward allreduce. Forget the manager and you allreduce K times for nothing, and throughput drops to the floor.

Check: Production patterns in the wild0/6

1. What does find_unused_parameters=True do in production DDP?Define

2. What does the no_sync context manager do in DDP?Define

3. What condition makes static_graph=True appropriate, so DDP can precompute the bucket schedule?Explain

4. You accumulate gradients over 8 microbatches and throughput is far below expectation because every microbatch triggers an allreduce. What is the fix?Analyze

5. Your model has an early-exit branch, so some layers get no gradient on some steps. static_graph is meant for forwards that are stable across steps. Which setup fits?Design

6. You need gradient accumulation over K microbatches without paying K allreduces, and you want proof it matches single-process training. Which design fits?Design

Use It

Where this pattern shows up in practice:

PyTorch DDP. The canonical implementation. torch.nn.parallel.DistributedDataParallel(model) wires up bucketing, overlap, and the no_sync context.
HuggingFace Accelerate. Adds a launcher that handles torchrun env vars and the model wrap. Same DDP under the hood.
Megatron-LM data parallel. Combines DDP with tensor parallel for large models. The data-parallel piece is the same allreduce-after-backward pattern.

Ship It

Lesson 78 (ZeRO sharding) swaps the per-parameter allreduce for reduce_scatter, so each rank stores only its shard of the optimiser state. Lesson 81 composes DDP with ZeRO into the end-to-end demo.

Exercises

Add gradient buckets of configurable size and measure the speedup vs one-allreduce-per-parameter on a deeper model.
Implement no_sync() as a context manager and verify gradient accumulation matches a single-process baseline over K microbatches.
Add a find_unused_parameters mode where the forward sometimes skips one of the MLP layers; without the flag the run should deadlock.
Replace gloo with torch.distributed.barrier()-only synchronisation to feel the difference between allreduce-based and barrier-based sync.
Measure the gradient-sync overhead as a fraction of step time for batch sizes 1, 16, 256 and explain the scaling.

Key Terms

Term	What people say	What it actually means
DDP	"Data parallel"	Wrapper that broadcasts params and allreduces grads each step
Bucket	"Fuse grads"	Group N small allreduces into one large one (not in `main.py`; exercise 1)
Overlap	"Hide comm"	Issue allreduce while later layers still computing backward (production DDP only)
no_sync	"Accumulate"	Skip the post-backward allreduce for gradient accumulation
find_unused	"Branchy forward"	Detect parameters with no grad before reducing

Author's Note

My bet is that the reference loop deserves a closer read than the DDP wrapper. The wrapper is a broadcast and an allreduce-and-divide. The reference walks each rank's micro-batch in rank order on one process, and matching it is the correctness test here.

I'd also hold off on calling this a production DDP. sync_grads in code/main.py issues one allreduce per parameter, so bucketing is what exercise 1 adds, and overlap is something the problem section describes but the code does not implement. Measure that gap before you trust the claim that sync is nearly free.

If the DDP run and the reference disagree, I'd start with sync_grads. It is where the code sums gradients across ranks and divides by world_size, so that is where I'd look for a wrong divisor or a skipped parameter.