Collective Ops From Scratch

Four collectives do the communication in distributed training. Allreduce leaves every rank holding the same reduced tensor. Broadcast copies one rank's tensor to the others. Allgather gives every rank every other rank's shard. Reduce_scatter leaves each rank with the sum of one chunk only.

Allreduce is the one that hurts when it's done naively. A naive version over N ranks sends N times the tensor to a root and broadcasts N times back. Bandwidth scales as O(N) per rank, the root becomes a bottleneck, and the wall-clock floor is the slowest link times N. Ring allreduce flattens that into 2(N-1) chunks of size T/N, so per-rank bytes drop to 2T(N-1)/N independent of cluster size. Tree allreduce wins on small N and high-latency links because depth is log2(N) hops instead of 2(N-1). Pick the wrong topology for the cluster shape and the slowest GPU dictates step time.

The frameworks you'll read in this track lean on these four. PyTorch DDP synchronises gradients with one allreduce per parameter bucket. ZeRO shards optimiser state by reduce_scatter and broadcasts updated parameters by allgather. FSDP turns the full forward into allgather plus reduce_scatter. Pipeline parallel needs broadcast for activations across stage groups. Without the four collectives in your head, it's hard to say why training stalls, why a gradient mismatch shows up at rank 3, or why the pipeline bubble doubles when you swap topologies.

Ring allreduce in two passes

Ring allreduce computes the elementwise sum of every rank's tensor and leaves the result on all ranks, using only sends between neighbours on a cycle.

Split the tensor into N equal chunks indexed 0..N-1. Rank r starts by sending the chunk whose index equals its rank. If the element count does not divide by N, ring_allreduce zero-pads the tensor up to the next multiple of N and trims the result back at the end. Pass 1, reduce-scatter, runs N-1 steps. At step s, rank r sends chunk (r - s) mod N to rank (r + 1) mod N and receives chunk (r - s - 1) mod N from rank (r - 1) mod N, accumulating the received chunk into its local copy. After N-1 steps, rank r owns the full sum for chunk (r + 1) mod N. Pass 2, allgather, runs another N-1 steps and rotates the finished chunks around the ring until each rank holds the full sum for every chunk.

The table lists tree allreduce for comparison only. code/main.py implements ring allreduce, and its broadcast is a doubling tree, but there is no tree allreduce.

Queue mesh as a stand-in for NCCL

NCCL runs over PCIe and NVLink with hardware-offloaded reductions, and a CPU has no equivalent. Here every ordered pair of ranks gets its own multiprocessing.Queue (a full NxN grid), and the ring algorithms use only the neighbour edges. Each queue gives ordered point-to-point delivery with a single producer and a single consumer. The reduction happens in user space, so you pay Python overhead, but the wire pattern matches NCCL's ring allreduce. Check correctness on the queue version and the schedule carries over to a cluster. The timings don't.

Verify against gloo

verify_against_gloo runs each primitive on the queue mesh and through torch.distributed initialised with the gloo backend, on the same tensors and world size, and main() calls it for all four. The check passes when the max absolute difference is below 1e-5. It is a tolerance rather than byte equality because float addition is not associative, so two implementations that sum in a different order can differ in the last bits. If your ring allreduce diverges from gloo by more than that, the check fails. A reference implementation catches the bugs that make a primitive only look correct. Without it, a bug can survive until step 10000 of a real training run.

code/main.py implements:

• Mesh dataclass, one rank's view of a full NxN grid of multiprocessing.Queues. It holds rank and world_size and exposes send(dst, tensor) and recv(src). recv times out after RECV_TIMEOUT_S = 30.0. The ring algorithms use only the neighbour edges.
• ring_allreduce(mesh, tensor) running the two-pass algorithm, zero-padding to a multiple of world_size.
• broadcast(mesh, tensor, src) over a doubling tree.
• allgather(mesh, tensor) using N-1 rotations.
• reduce_scatter(mesh, tensor) as the first half of allreduce. It pre-rotates the chunks so rank r returns the sum of chunk r, matching the torch.distributed contract, and raises ValueError when the element count is not divisible by world_size.
• gloo_reference(op, world_size, per_rank_tensors) that runs the same inputs through torch.distributed with gloo, and verify_against_gloo(op, world_size, per_rank_tensors) that returns (match, max_abs_diff) with a 1e-5 threshold.

Run it:

python3 code/main.py

The output is a per-primitive verification table comparing queue-mesh and gloo outputs, followed by a per-rank byte counter that checks the 2T(N-1)/N scaling.

Three patterns harden the primitives enough to ship.

Bucket gradients before allreduce. A 1B-parameter model has tens of thousands of gradient tensors. One allreduce per tensor pays the latency floor N times. DDP buckets gradients into ~25 MB chunks and issues one allreduce per bucket, so the small tensors ride along with the big ones. Skip bucketing and the latency overhead dominates the step.

Overlap communication with computation. Backward computes gradients layer by layer in reverse order. Once the last layer's gradient is ready, start its allreduce while the next layer keeps computing. PyTorch DDP wires this with bucket-ready hooks. When the network has slack, the overlap can halve visible communication time.

Pick ring or tree by message size. NCCL ships a topology detector that picks ring for messages above ~1 MB and tree below. The crossover is bandwidth versus latency: above 1 MB, the bandwidth term 2T(N-1)/N dominates and ring wins, and below 1 MB the log2(N) hop count wins. Hard-coding one topology costs throughput at the wrong message size.

Where the primitives show up:

• PyTorch DDP. Calls dist.all_reduce on bucketed gradients after backward. The bucket size is tunable, and the 25 MB default is reasonable for 100Gbit Ethernet.
• DeepSpeed ZeRO. Issues reduce_scatter to shard gradients and allgather to reconstruct full parameters before forward. Those are the same primitives you build here.
• FSDP. Forward begins with allgather to unshard the layer, computes, then reduces with reduce_scatter and discards the unshard. Same primitives, different schedule.

Lessons 77-81 build on these queue-mesh primitives. Lesson 77 wires allreduce into DDP, lesson 78 wires reduce_scatter into ZeRO, and lesson 79 wires broadcast into pipeline activations. Lesson 81 composes all four into the end-to-end demo.

1. Add a tree allreduce variant and switch between ring and tree by message size. Measure the crossover.
2. Mesh.recv already hardcodes RECV_TIMEOUT_S = 30.0. Make it configurable with a recv_timeout_ms argument so a stalled rank surfaces a deadline error you choose instead of waiting a fixed 30 seconds.
3. Replace multiprocessing.Queue with TCP sockets for the four primitives. Keep verify_against_gloo and use a real wire.
4. Add a bandwidth instrumentation hook so the per-rank byte counter logs to JSONL.
5. Compare wall-clock time of ring versus tree on 4 ranks for tensors of size 1KB, 1MB, 16MB. Defend the crossover empirically.

• PyTorch Distributed: NCCL collectives
• Horovod ring allreduce paper
• NCCL topology and algorithm selection
• Patarasuk and Yuan, Bandwidth optimal allreduce algorithms
• Phase 10 Lesson 05 - distributed training overview
• Phase 19 Lesson 77 - DDP wired on top of these primitives