Distributed Training, Simulated · AI infra interview lesson

Why It Matters

Two ranks produce gradients 2 and 6; their global average is 4. Trace data on each rank before comparing replication, reduction and sharding. Separate communication from optimizer updates.

Two numbers first. Mixed-precision Adam training of a 7B model needs 2 (bf16 weights) + 2 (bf16 grads) + 4 (fp32 master) + 4 (m) + 4 (v) = 16 B of static memory per parameter, 16 * 7e9 = 112 GB, more than an 80 GB H100 before any activations. DDP all-reduces the bf16 gradients every step, 2 * 7e9 B = 14 GB, and with 8-GPU ring all-reduce each GPU sends 2 * (8 - 1) / 8 * 14 GB = 24.5 GB. ZeRO trades communication for memory, and TP and PP go further and split the weights and the compute.

Walk the Shapes First

N = 4 ranks, each holding a gradient of length 8 cut into 4 chunks of 2 elements. Ring all-reduce has two phases of N−1 = 3 steps each, and rank r only sends to rank (r+1) % 4:

Rules read from the table: ring all-reduce = reduce-scatter + all-gather; each rank sends one chunk per step, so per-GPU traffic is 2(N−1)/N copies of the data and barely grows with N. That is also why ZeRO can split the all-reduce into a reduce-scatter and an all-gather without adding traffic.

Running the examples: execute this page in order in one Python environment. Later snippets may reuse earlier definitions. CPU examples validate calculations and state changes, not GPU timing. Download the notebook for the complete exercise workspace.

Core Implementations

Step 1: collective primitives

New idea: xs: list[Tensor] represents the same tensor on N ranks, with xs[r] being rank r's copy. Six primitives:

• broadcast: every rank gets xs[src]
• reduce: only dst gets sum(xs)
• all_reduce: every rank gets sum(xs)
• reduce_scatter: rank r gets chunk r of sum(xs)
• all_gather: every rank gets cat(xs)
• all_to_all: chunk j of rank i goes to rank j, at position i (a transpose of the chunk grid)

Identity: all_gather(reduce_scatter(x)) == all_reduce(x).

import torch
torch.manual_seed(0)

def broadcast(xs, src=0):
    return [xs[src].clone() for _ in xs]

def reduce(xs, dst=0):
    out = [x.clone() for x in xs]
    out[dst] = torch.stack(xs).sum(0)
    return out

def all_reduce(xs):
    s = torch.stack(xs).sum(0)
    return [s.clone() for _ in xs]

def reduce_scatter(xs):
    s = torch.stack(xs).sum(0)
    return [c.clone() for c in s.tensor_split(len(xs))]   # rank r gets chunk r of the sum

def all_gather(xs):
    full = torch.cat(xs)
    return [full.clone() for _ in xs]

def all_to_all(xs):
    N = len(xs)
    chunks = [x.tensor_split(N) for x in xs]              # chunks[i][j]: rank i -> rank j
    return [torch.cat([chunks[i][j] for i in range(N)]) for j in range(N)]

N = 4
xs = [torch.randn(8) for _ in range(N)]
total = torch.stack(xs).sum(0)
assert all(torch.equal(b, xs[2]) for b in broadcast(xs, src=2))
assert torch.allclose(reduce(xs)[0], total) and torch.equal(reduce(xs)[1], xs[1])
assert all(torch.allclose(a, total) for a in all_reduce(xs))
rs = reduce_scatter(xs)
assert [c.shape for c in rs] == [torch.Size([2])] * 4
assert all(torch.allclose(a, b) for a, b in zip(all_gather(rs), all_reduce(xs)))   # AG(RS(x)) == AR(x)

# all_to_all is a transpose of the [rank, chunk] grid, so applying it twice is the identity
a2a = all_to_all(xs)
grid = torch.stack(xs).view(N, N, 2)                       # [src rank, chunk, elems]
assert torch.equal(torch.stack(a2a).view(N, N, 2), grid.transpose(0, 1))
assert all(torch.equal(a, b) for a, b in zip(all_to_all(a2a), xs))

# numel not divisible by N: tensor_split gives chunks of size 3, 3, 2, 2
ys = [torch.randn(10) for _ in range(N)]
assert [c.numel() for c in reduce_scatter(ys)] == [3, 3, 2, 2]
assert all(torch.allclose(a, b) for a, b in zip(all_gather(reduce_scatter(ys)), all_reduce(ys)))

Step 2: ring all-reduce and its traffic

New idea: all_reduce in Step 1 assumes a central node does the sum; in the ring algorithm each rank only talks to its right neighbour. All ranks send at the same time within a step, so the simulation must first collect every message of the step and then deliver them together; otherwise a rank would read a buffer already changed in this step.

• reduce-scatter step t: rank r sends chunk (r - t) % N, the receiver accumulates
• all-gather step t: rank r sends chunk (r + 1 - t) % N, the receiver overwrites
• per-rank traffic: 2 * (N - 1) / N * numel

import torch
torch.manual_seed(0)

def ring_all_reduce(xs):
    N = len(xs)
    bufs = [list(x.clone().tensor_split(N)) for x in xs]
    sent = [0] * N
    for t in range(N - 1):                                 # phase 1: reduce-scatter
        msgs = [(r, (r - t) % N, bufs[r][(r - t) % N].clone()) for r in range(N)]
        for r, c, data in msgs:                            # all sends of one step happen at once
            bufs[(r + 1) % N][c] = bufs[(r + 1) % N][c] + data
            sent[r] += data.numel()
    # now rank r holds the full sum of chunk (r + 1) % N
    for t in range(N - 1):                                 # phase 2: all-gather
        msgs = [(r, (r + 1 - t) % N, bufs[r][(r + 1 - t) % N].clone()) for r in range(N)]
        for r, c, data in msgs:
            bufs[(r + 1) % N][c] = data
            sent[r] += data.numel()
    return [torch.cat(b) for b in bufs], sent

for N, numel in [(4, 8), (8, 64), (2, 6), (4, 10), (3, 1)]:
    xs = [torch.randn(numel) for _ in range(N)]
    out, sent = ring_all_reduce(xs)
    total = torch.stack(xs).sum(0)
    assert all(torch.allclose(o, total, atol=1e-5) for o in out)
    assert sum(sent) == 2 * (N - 1) * numel                # every step moves each chunk exactly once
    if numel % N == 0:
        assert sent == [2 * (N - 1) * numel // N] * N      # 2(N-1)/N per rank

# per-GPU traffic for 7B bf16 gradients on 8 GPUs
grad_bytes = 2 * 7e9
assert abs(2 * (8 - 1) / 8 * grad_bytes / 1e9 - 24.5) < 1e-9

When numel is not divisible by N the per-rank traffic differs slightly, but the total stays 2(N−1) * numel. Besides ring, NCCL also has tree algorithms with lower latency for small messages.

Step 3: equivalence of DDP and gradient accumulation

New idea: in DDP each rank takes B/N samples, runs backward on its own, then all-reduces the gradients and divides by N. (Pro lesson)

Step 4: ZeRO sharded step and per-GPU memory

New idea: in DDP every rank stores a full copy of the Adam state, which is redundant. (Pro lesson)

Convert sharded states into per-rank bytes

ZeRO stages shard different objects: optimizer states, gradients and parameters. This estimates those static terms only; (Pro lesson)

Step 5: splitting a layer: TP MLP and MoE dispatch/combine

New idea: TP splits the weight matrix itself. For y = x @ W (here W is [in, out] ): column parallel: split W along out; (Pro lesson)

Route and combine MoE work

Compare with an all-expert reference, then dispatch tokens to selected experts and restore token order before weighted combination. (Pro lesson)

In Real Systems

PyTorch torch.distributed : all_reduce , broadcast , reduce_scatter_tensor , all_gather_into_tensor and all_to_all_single correspond to Step… (Pro lesson)

Common Pitfalls

Not dividing by N after the DDP all-reduce → the gradient is N times the single-GPU one → the learning rate is effectively multiplied by N a… (Pro lesson)

Say It in the Interview

I model N GPUs as a list of tensors, so every parallel scheme becomes local compute plus a collective, and I check it against single-GPU tra… (Pro lesson)

Hands-on Notebook

The notebook has 8 exercises, all simulating multiple GPUs with list[Tensor] in one process. (Pro lesson)