Distributed Training, End to End

A 70-billion-parameter model needs 1.1 TB just to hold its weights, gradients, and Adam states — fourteen times what fits on one GPU. Compute the bytes, then install every fix in order: data parallelism with ring all-reduce, FSDP/ZeRO sharding, tensor and pipeline parallelism with their bubbles, activation checkpointing, and the interconnect that decides which of them you can afford — drawn, computed, and animated.

System design · AI / ML.

A free, interactive, animated visual explainer of Distributed Training, End to End — built to be understood, not skimmed.

Questions

Why does a large model not fit on a single GPU?
Because training holds far more than the weights. With mixed-precision Adam the resident state is 16 bytes per parameter: 2 bytes for the fp16 weights, 2 for the fp16 gradients, and 12 for the fp32 optimizer states (a master copy of the weights plus Adam’s first and second moments). The ZeRO paper puts it exactly: "Mixed-precision Adam has K=12. In total, this results in 2Ψ+2Ψ+KΨ=16Ψ bytes of memory requirement." A 70-billion-parameter model is therefore about 1,120 GB before a single activation — roughly fourteen 80 GB GPUs just to hold it, which is why training is distributed at all.
What is the difference between data, tensor, and pipeline parallelism?
They split different axes. Data parallelism replicates the whole model on every GPU and splits the batch, synchronising gradients with an all-reduce each step — simple, but every GPU still needs room for the full model. Tensor parallelism splits individual matrix multiplications across GPUs (Megatron shards the MLP’s first GEMM by columns and its second by rows, then all-reduces), cutting per-GPU memory but demanding very high bandwidth, so it lives inside one node. Pipeline parallelism splits the layers into stages on different GPUs and streams micro-batches through them, paying a "bubble" of idle time at the fill and drain. Large runs combine all three — 3D parallelism.
What is FSDP and how does it relate to ZeRO?
FSDP (Fully Sharded Data Parallel) is PyTorch’s implementation of ZeRO Stage 3. Plain data parallelism replicates the optimizer states, gradients, and parameters on every GPU; ZeRO shards them instead. Stage 1 shards the optimizer states (4× memory reduction), Stage 2 adds the gradients (8×), and Stage 3 adds the parameters so per-GPU memory falls linearly with the number of GPUs. FSDP "unshards (via all-gather) before the forward, reshards after the forward, unshards before the backward … and reshards after," then reduce-scatters the gradients — trading a little extra communication for a large drop in memory.
What is the pipeline bubble and how do you shrink it?
The bubble is the idle time while a pipeline fills and drains: the last stage sits waiting for the first micro-batch to arrive, and the first stage sits waiting for the backward pass to return. For P stages and M micro-batches the bubble fraction is (P−1)/(M+P−1) — GPipe writes it as "O(K−1/M+K−1) amortized over the number of micro-steps M." You shrink it by feeding more micro-batches (GPipe found it "negligible when M≥4×K"); 1F1B scheduling keeps the same bubble but far lower activation memory, and interleaved schedules cut the bubble further.
Why does the interconnect decide which parallelism you use?
Because each parallelism moves a different amount of data at a different frequency. Tensor parallelism all-reduces activations twice per layer, so it needs the fastest link — NVLink at 900 GB/s per H100 — and is kept inside a single node. Data parallelism all-reduces gradients once per step (about twice the gradient size, independent of GPU count for a ring), so it tolerates InfiniBand at ~50 GB/s across nodes. Pipeline parallelism only passes activations between adjacent stages, the lightest traffic, so it spans the slowest links. The bandwidth hierarchy of your cluster is what forces the layout.

Related explainers