Lab 01 — Ring All-Reduce Over Real Sockets (Python)

Phase: 08 — Distributed & HPC | Difficulty: ⭐⭐⭐⭐☆ | Time: 5–8 hours Language: Python (stdlib socket, multiprocessing) | Hardware: none (multi-process on one machine)

Concept primer: ../WARMUP.md Ch. 2–4.

Run

python solution.py --world 4          # 4 worker processes, ring all-reduce over TCP
python solution.py --world 8 --size 1000000

0. The mission

Implement the algorithm NCCL uses, over real OS sockets across real processes — so the communication is genuine, not simulated. Ring all-reduce in two phases (WARMUP Ch. 3):

  1. Reduce-scatter (N-1 steps): ring-pass chunks, accumulating, until each rank owns one fully-reduced chunk.
  2. All-gather (N-1 steps): ring-pass the completed chunks until everyone has all of them.

Then prove the cost: each rank moves ≈2(N-1)/N × V bytes — measured by counting actual socket bytes — and beats the naive all-gather-then-sum (N×V).

1. What the output shows

ring all-reduce: world=8, vector=80000 doubles
  correctness: PASS (every rank's result == sum 1..8 = 36)
  bytes/rank moved : ring=1,120,000   naive_allgather=4,480,000
  per-rank volume vs 2(N-1)/N*V bound: 1.00x  (1.00 = bandwidth-optimal)
  ring moves 4.0x fewer bytes than naive all-gather
  • Correctness: every rank's vector equals the true element-wise sum across all ranks — the collective is exact.
  • Bytes/rank: the ring moves ≈2(N-1)/N × V per rank, hitting the bandwidth-optimal bound. The key result — run --world 4 then --world 8: the ring's per-rank volume stays ≈2V (the factor barely changes from 1.5V to 1.75V) while naive all-gather grows as (N-1)·V, so the ring's advantage widens with N (2× at N=4, 4× at N=8). That N-independence is why ring all-reduce made large-scale data-parallel training practical (WARMUP Ch. 3).

2. Reading order (solution.py)

  1. worker — each rank: connects to its ring neighbors over TCP, gets its data shard, runs the two phases.
  2. reduce_scatter — the N-1 send/recv-and-add steps; note the chunk indexing (rank r sends chunk (r-step) % N).
  3. all_gather — the N-1 forward-only steps.
  4. The byte counters and the correctness check against a reference sum.

3. Extensions

  1. Tree all-reduce: implement reduce-up/broadcast-down; benchmark latency vs ring on small messages (1 KB) and large (10 MB) — find the crossover (WARMUP Ch. 4).
  2. Bucketing/overlap: split the vector into buckets, all-reduce each as it's "ready," overlapping with simulated backward compute; measure the hidden time (WARMUP Ch. 7).
  3. Failure demo: have one rank sys.exit() mid-collective; watch the others hang in recv. Then add a socket timeout that turns the hang into a detected failure with the culprit rank (WARMUP Ch. 8).
  4. Bandwidth-optimality check: vary world size 2/4/8/16; confirm per-rank volume stays ≈2V (N-independent) while a naive all-gather-sum grows as N·V.

4. Common pitfalls

  1. Deadlock from ordered send/recv: if every rank sends-then-recvs to the same neighbor synchronously, you can deadlock on socket buffers. The lab uses a send thread (or non-blocking) — keep that structure.
  2. Chunk index off-by-one: the rank-to-chunk mapping per step is the crux; get it wrong and the reduction is incorrect (the correctness assert catches it, but the index math is the lab's real content).
  3. Float serialization: use a fixed format (struct/array) consistently across ranks, or you'll get garbage sums.
  4. Counting bytes wrong: count payload bytes actually sent on the wire, not the logical vector size.

5. What this lab proves about you

You've implemented and measured the algorithm at the heart of every data-parallel training run, and you can derive its bandwidth-optimal cost from the byte counts you collected — not recite it. NCCL is no longer magic: it's this, with topology detection and GPUDirect RDMA on top.