Collective operations from scratch

1. Intuition

Picture $N$ GPUs (called ranks), each holding a tensor of the same shape: for example its own gradients after a backward pass on its own data.

The key identity: all-reduce = reduce-scatter + all-gather. First each rank becomes responsible for summing one slice; then everyone shares their finished slice.

Parameter server. The obvious design sends every tensor to one rank, sums, and sends the result back. The server's link carries $N - 1$ tensors in and $N - 1$ out, so the time grows with $N$.

Ring. Arrange ranks in a circle. Split each tensor into $N$ chunks. In each step every rank sends one chunk to its right neighbour and receives one from its left, adding it to its own copy of that chunk. After $N - 1$ steps each rank has the complete sum of one chunk (reduce-scatter). Another $N - 1$ steps pass the finished chunks around (all-gather). Every link is busy every step and no rank is special.

Tree. For a broadcast, the rank with the data sends to one other, then both send to two more, then four send to four more. $N$ ranks are reached in $\log_2 N$ rounds. Trees take fewer steps, which matters when each message has a fixed startup cost.

2. Visualization

A model, not a measurement: synthetic link with 10 microseconds per message and 50 GB/s. Real libraries such as NCCL use pipelined and hierarchical variants whose crossover points depend on the hardware.

Try it (predict first, then run code/27-training-at-scale/collectives.py and edit it):

1. Set N, T = 8, 16. Predict the number of ring steps and the floats sent per rank before running.
2. Change alpha to 100e-6 (a slow network). Predict whether the ring-tree crossover moves to bigger or smaller tensors.
3. In reduce_scatter, predict which chunk index rank 2 ends up holding with $N = 8$. Check the printout.
4. Change tree_broadcast(8, ...) to 9 ranks. Predict the number of rounds and sends.

3. The math

Symbols

Ring all-reduce

Each chunk has $T/N$ bytes. Reduce-scatter takes $N - 1$ steps, each rank sending one chunk per step; all-gather takes another $N - 1$.

As $N$ grows the bandwidth term tends to $2T\beta$: it does not grow with the number of ranks. The latency term grows linearly with $N$.

Reduce-scatter and all-gather alone each cost half: $(N-1)\alpha + \frac{N-1}{N} T \beta$.

Parameter server and binary tree

The server's single link carries $N - 1$ full tensors each way. A simple binary tree reduces up and broadcasts down in $2 \log_2 N$ rounds, each moving the whole tensor. Trees win when $\alpha$ dominates (small messages), rings when $T\beta$ dominates (large messages). Production libraries use pipelined trees and hierarchical schemes that do better than this naive tree; NCCL picks an algorithm per call based on message size and topology.

Worked example: four ranks, eight numbers each

Ranks hold (synthetic integers):

Split into $N = 4$ chunks of 2. The column sums are [18, 27 | 21, 20 | 17, 21 | 23, 8].

1. Step $s = 0$. Rank $r$ sends its chunk $r$ to rank $r + 1$, which adds it to its own copy. Rank 1 receives rank 0's chunk 0, [8, 6], and adds it to its own [1, 8]: [9, 14].
2. Step $s = 1$. Rank $r$ sends chunk $r - 1$ (the partial sum it just updated). Rank 2 receives rank 1's chunk 0, [9, 14], and adds its own [6, 5]: [15, 19].
3. Step $s = 2$. Rank 3 receives chunk 0, [15, 19], and adds [3, 8]: [18, 27]. That is the full sum of chunk 0.
4. After reduce-scatter: rank 3 holds chunk 0 = [18, 27]; rank 0 holds chunk 1 = [21, 20]; rank 1 holds chunk 2 = [17, 21]; rank 2 holds chunk 3 = [23, 8]. With this schedule rank $r$ finishes chunk $(r + 1) \bmod N$; libraries relabel the chunks so rank $r$ owns chunk $r$.
5. All-gather: 3 more steps pass the finished chunks around. Every rank ends with [18, 27, 21, 20, 17, 21, 23, 8].
6. Traffic: $2(N-1) = 6$ steps, each sending one 2-float chunk, so $6 \times 2 = 12$ floats per rank $= \frac{2 \times 3}{4} \times 8 = 12$. ✓

Worked example: gradients of a 7B-style model

A 7B-style model has $P = 6{,}738{,}149{,}376$ parameters. Its bf16 gradients are $T = 2P = 13{,}476{,}298{,}752$ bytes $\approx 13.48$ GB. Ring all-reduce over $N = 8$ sends $\frac{2 \times 7}{8} \times 13.48 = 1.75 \times 13.48 = 23.58$ GB per GPU per step. At 50 GB/s per link (synthetic) that is $23.58 / 50 = 0.47$ s of pure transfer, unless it overlaps with the backward pass.

4. Implementation

Build it from scratch

We simulate $N$ ranks in one process. A Ring moves one array per rank to its right neighbour per step and counts what each rank sends, so we can check the formula.

class Ring:
    """N ranks on a ring. A step is synchronous: every rank sends one chunk to its right neighbour."""

    def __init__(self, n):
        self.n = n
        self.sent = [0] * n      # floats sent by each rank
        self.steps = 0

    def shift(self, outgoing):
        """outgoing[r] is what rank r sends to rank r+1. Returns what each rank receives."""
        for r, x in enumerate(outgoing):
            self.sent[r] += x.size
        self.steps += 1
        return [outgoing[(r - 1) % self.n] for r in range(self.n)]

Predict: in one shift, how many arrays cross the network in total?

Reduce-scatter: at step s rank r sends chunk (r - s) mod N and adds what it receives into chunk (r - s - 1) mod N. The chunk it sends is always the one it updated in the previous step, so partial sums travel around the ring and grow.

def reduce_scatter(ring, tensors):
    """Ring reduce-scatter. After N-1 steps rank r holds the full sum of chunk (r + 1) mod N."""
    n = ring.n
    chunks = [np.array_split(t.astype(np.float64), n) for t in tensors]   # chunks[rank][chunk]
    for s in range(n - 1):
        send_idx = [(r - s) % n for r in range(n)]
        received = ring.shift([chunks[r][send_idx[r]] for r in range(n)])
        for r in range(n):
            idx = (r - s - 1) % n                       # the chunk my left neighbour just sent
            chunks[r][idx] = chunks[r][idx] + received[r]
    return [((r + 1) % n, chunks[r][(r + 1) % n]) for r in range(n)]    # (chunk index, summed chunk)

Predict: if one rank's tensor contained a NaN in chunk 2, which ranks would end up with a NaN after reduce-scatter? After all-gather?

All-gather passes each finished chunk to the right $N - 1$ times. All-reduce is the two composed.

def all_gather(ring, owned, n_chunks):
    """Ring all-gather. owned[r] = (chunk index, data). After N-1 steps every rank has every chunk."""
    n = ring.n
    have = [{i: d} for i, d in owned]
    carry = list(owned)
    for _ in range(n - 1):
        ring.shift([d for _, d in carry])               # count the sends
        carry = [carry[(r - 1) % n] for r in range(n)]  # pass (index, data) to the right
        for r in range(n):
            have[r][carry[r][0]] = carry[r][1]
    return [np.concatenate([h[i] for i in range(n_chunks)]) for h in have]

def all_reduce(ring, tensors):
    """Ring all-reduce = reduce-scatter, then all-gather of the summed chunks."""
    return all_gather(ring, reduce_scatter(ring, tensors), ring.n)

Predict: all-gather does no arithmetic. Why does it still cost as many bytes as reduce-scatter?

The cost models are one line each:

def t_ring(n, T, alpha, beta):
    return 2 * (n - 1) * alpha + 2 * (n - 1) / n * T * beta

def t_tree(n, T, alpha, beta):
    """Reduce up a binary tree, then broadcast down: 2 log2(N) rounds, each moving the full tensor."""
    return 2 * np.log2(n) * (alpha + T * beta)

Predict: at $N = 64$, for which tensor size do the two times match? (Set them equal and solve for $T$.)

Use the library

On GPUs, torch.distributed with the NCCL backend runs these as optimized kernels over NVLink and InfiniBand:

import torch.distributed as dist
dist.init_process_group("nccl")                 # launched with torchrun --nproc-per-node=8
dist.all_reduce(grad, op=dist.ReduceOp.SUM)     # in place on every rank
dist.reduce_scatter_tensor(shard_out, full_in)  # shard_out has numel / world_size
dist.all_gather_into_tensor(full_out, shard_in)
dist.broadcast(weights, src=0)

You rarely call these directly. DistributedDataParallel groups gradients into buckets (bucket_cap_mb, 25 MB by default) and launches one all-reduce per bucket as soon as backward has filled it, so communication overlaps with the rest of backward. FSDP and DeepSpeed ZeRO issue the reduce-scatters and all-gathers for you.

Run it

uv run code/27-training-at-scale/collectives.py

Key lines:

  after reduce-scatter rank 0 holds chunk 1 = [21, 20]
  after reduce-scatter rank 3 holds chunk 0 = [18, 27]
  after all-gather every rank equals the sum: True
  steps = 6 (= 2(N-1)); floats sent per rank = [12, 12, 12, 12] (2(N-1)/N x T = 12)
  N = 1024: 2(N-1)/N = 1.9980 GB
3) Tree broadcast to 8 ranks: 3 rounds (log2 8), 7 sends, all have it: True
     0.004 MB: ring     1.260 ms  tree     0.121 ms  server     1.270 ms
  1000.000 MB: ring    40.635 ms  tree   240.120 ms  server  2521.260 ms

5. Engineering

Bucket small tensors. A model has hundreds of parameter tensors. One all-reduce each would pay the latency $\alpha$ hundreds of times. Bucketing turns them into a few large messages, where the ring's bandwidth term dominates.

Overlap with backward. Gradients of the last layers are ready first. DDP starts their all-reduce while earlier layers are still computing, hiding much of the communication. If communication still exceeds compute, the GPUs wait: check a profiler for idle gaps.

Topology matters. Inside a node, GPUs talk over NVLink at hundreds of GB/s; between nodes, over InfiniBand or Ethernet at roughly an order of magnitude less per GPU. Libraries run hierarchical collectives: reduce inside the node first, then across nodes. This is why tensor parallelism stays inside a node and data parallelism spans nodes.

Straggler effect. A ring is synchronous: every step waits for the slowest link. One slow GPU, a degraded cable, or a noisy neighbour slows every rank. Collective timeouts (NCCL_TIMEOUT, timeout= in init_process_group) turn a hang into an error you can debug.

Precision of the sum. Summing many bf16 gradients loses precision. Frameworks often reduce in fp32, or average rather than sum, to stay in range.

Common mistakes

• Thinking all-reduce traffic grows with $N$: per rank it approaches $2T$ and stops.
• Forgetting to divide by $N$ after a SUM all-reduce in data parallelism, which multiplies the effective learning rate by $N$.
• Calling collectives in different orders on different ranks (for example inside an if rank == 0 block): every rank waits forever. Every rank must make the same collective calls in the same order.
• Measuring communication cost with tiny tensors and extrapolating: small messages are latency-bound and give the wrong picture.
• Ignoring the latency term at large $N$: $2(N-1)\alpha$ grows with ranks, which is one reason hierarchical and tree algorithms exist.

6. Knowledge check

Practice exercise. You all-reduce 2 GB of gradients over 16 ranks with a ring. How many bytes does each rank send, how many steps are there, and how big is each message?

Implementation challenge. Write ring_all_reduce_mean(ring, tensors) that returns the average instead of the sum, and verify it against np.mean. Then add a check that raises a clear error if the ranks pass tensors of different shapes.

def ring_all_reduce_mean(ring, tensors):
    shapes = {t.shape for t in tensors}
    if len(shapes) != 1:
        raise ValueError(f"all ranks must pass the same shape, got {shapes}")
    return [x / ring.n for x in all_reduce(ring, tensors)]

out = ring_all_reduce_mean(Ring(4), tensors)
assert all(np.allclose(o, np.mean(tensors, axis=0)) for o in out)

Summary

• All-reduce gives every rank the sum; reduce-scatter gives each rank the sum of one slice; all-gather shares slices; broadcast copies from one rank.
• All-reduce = reduce-scatter + all-gather. A ring does each in $N - 1$ steps of $T/N$-sized chunks.
• Ring all-reduce sends $\frac{2(N-1)}{N}T$ per rank, under $2T$ for any $N$; a parameter server's link carries $N - 1$ tensors each way.
• The $\alpha$-$\beta$ model explains the choice: trees for small, latency-bound messages, rings for large, bandwidth-bound ones.
• Bucketing and overlapping with backward hide most data-parallel communication.

Next: Scaling and distributed training

Related: ZeRO and FSDP · Gradient descent

Go deeper, code-first: AI Engineering from Scratch, collective ops over real processes.

Parts of this lesson are adapted from AI Engineering from Scratch by Rohit Ghumare (MIT license, notice): Collective Ops From Scratch.

Interview angle