3D parallelism explorer

42 Book labs · Distributed training

Splitting a model across 16 GPUs

From Chapter 46 of Large Language Models from the Ground Up. No single seam suffices at scale, so real runs stack three: data, tensor, and pipeline parallelism. Toggle each on and watch 16 GPUs re-partition a model — and watch the memory per GPU and the network traffic change with it.

Interactive

Toggle the three seams; degrees auto-factor to 16 GPUs

Each GPU tile is coloured by pipeline stage (hue) and tensor slice (shade); the outlined boxes are data-parallel replicas. The badge shows DP×TP×PP = 16. Memory and communication are computed for a 16-billion-parameter model at 16 bytes/parameter.

DP×TP×PP = 16×1×1 = 16

pipeline stage = hue tensor slice = shade outlined box = data-parallel replica

Memory per GPU (weights + optimizer)

16 GB

Model size N16 B params

Bytes / param16 B

Params held per GPU1 B

Shard factor (TP×PP)16×

16 B/param = bf16 weight 2 + bf16 grad 2 + fp32 master 4 + Adam m 4 + Adam v 4. Data parallelism replicates the model, so it does not shrink per-GPU memory (that is ZeRO/FSDP's job); only tensor and pipeline splits do.

What crosses the wire

Data: all-reduce gradients across replicas — once per step, ~2× the gradient size per GPU.

Tensor: all-gather / all-reduce activations every layer — only survivable on the fastest in-node links.

Pipeline: point-to-point activations between adjacent stages — only megabytes, so it tolerates slower links.

The chapter's rule of thumb: TP innermost (in a node), PP next, DP outermost — following each method's appetite for bandwidth.

What to notice

Turn on only Data parallel and every GPU is a full copy: 256 GB of training state each — far past an 80 GB GPU, which is precisely why 70B "can't fit" and why ZeRO/FSDP shards the copies. Turn on Tensor or Pipeline and the model is genuinely cut into pieces: params-per-GPU falls as 1/(TP×PP), and the memory drops with it. But each seam bills the network differently — DP all-reduces gradients once per step, TP synchronises every layer (so it stays inside one node), PP passes only small activations between stages (so it spans machines). The book's canonical fleet stacks all three — TP=8, PP=16, DP=80 = 10,240 GPUs on one model — nesting the chatty seam inside the fast links and the frugal one across the slow ones. This explorer shows the same logic at 16 GPUs, where the factors stay small enough to see.

← Back to all labs