xavier-ramirez.com
STAGE 0 · 3 NODES

Distributed training

A frontier model is far too big for one GPU. This is how labs split the model — and the batch — across hundreds of chips, so a 7.5-billion-parameter model that needs 120 GB trains on 80 GB cards. Play with each split and watch the memory per GPU fall.

THE 3 SPLITS
01Data parallel02Tensor & pipeline03ZeRO / FSDP sharding
NODE 01 / 03

Copy the model, split the batch

The simplest split: put a full copy of the model on every GPU, feed each a different slice of the batch, then average their gradients so all copies stay identical.

  • Every GPU holds the whole model — same weights, same optimizer state, just different data.
  • Each GPU trains on its own batch shard, runs the forward and backward pass you met in How models learn, and gets its own gradients.
  • An all-reduce averages the gradients across GPUs — one network step that keeps every copy in sync. (all-reduce = each GPU sends its gradients and gets back the sum.)
  • Throughput scales, memory doesn't. N GPUs train ~N× faster, but each still needs room for the entire model — so data-parallel alone can't train a model bigger than one card.
GOOD TO KNOW · THE SYNC
all-reduce over NCCL
Averaging the gradients is a collective operation — NVIDIA's NCCL library runs it over fast GPU-to-GPU links (NVLink inside a node, InfiniBand between nodes). It's the main communication cost of data-parallel training, and why network bandwidth matters as much as raw compute.
DATA-PARALLEL GPU GRID
full model copybatch shard

Add GPUs — throughput climbs, but the memory per GPU stays pinned to one full copy.

Each replica is a full copy of the model. Adding replicas speeds up training but never shrinks the per-GPU model — that's the wall the next two splits break.

DATA_PARALLEL
GPUS4
GLOBAL BATCH32 samples
THROUGHPUT4×
MODEL / GPU120 GB
More replicas process more samples per step at higher throughput — but each GPU still stores one entire model.
NODE 02 / 03

Split one model across GPUs

When the model itself won't fit on a card, you cut it up. There are two ways, and a real run uses both at once.

  • Tensor parallel slices each layer. One layer's weight matrix is cut into vertical strips, one per GPU; each does part of the matrix multiply, then an all-gather stitches the pieces back.
  • Pipeline parallel slices the stack. Consecutive layers become stages — GPU 1 holds layers 1–4, GPU 2 holds 5–8 — and activations flow forward like an assembly line.
  • Pipelines have a bubble. While the first stage works, later stages sit idle; splitting the batch into micro-batches keeps them fed and shrinks the gap.
  • Both shrink the per-GPU model. Tensor parallel needs the fastest links (heavy chatter inside every layer); pipeline parallel is cheaper to communicate but harder to keep busy.
TENSOR vs PIPELINE
TENSOR PARALLEL
one layer, split across GPUs
PIPELINE PARALLEL
the stack, split into stages

Raise the degree — tensor parallel cuts each layer into more strips; pipeline parallel cuts the stack into more stages.

Degree = how many GPUs share the model. A higher degree means a smaller slice per GPU — and more communication to stitch the pieces back together.

MODEL_SPLIT
DEGREE3
TENSOR1/3 per layer
PIPELINE1/3 of layers
MODEL / GPU1/3
With degree N, tensor parallel puts 1/N of each layer on a GPU and pipeline parallel puts 1/N of the layers — either way, each GPU holds a fraction of the model.
NODE 03 / 03

Shard what each GPU stores

Data parallel wastes memory: every GPU keeps an identical copy of the params, gradients and optimizer state. ZeRO — and PyTorch's FSDP — stops the duplication, sharding those three across GPUs and gathering each piece only when it's needed.

  • Training memory is three things (from How models learn): the parameters, their gradients, and the optimizer state. With mixed-precision Adam that's about 16 bytes per parameter — 2 + 2 + 12.
  • ZeRO-1 shards the optimizer state, the biggest chunk (12 of the 16 bytes). Each GPU keeps only 1/N of it.
  • ZeRO-2 also shards the gradients; ZeRO-3 (FSDP) shards the parameters too — so no GPU holds a full copy of anything.
  • Same maths as data parallel, a fraction of the memory. A 7.5B model needs 120 GB replicated; sharded across 8 GPUs, ZeRO-3 drops it to 15 GB — it fits an 80 GB card with room to spare.
GOOD TO KNOW · 3D PARALLELISM
Real runs combine all three
A production run stacks the splits: tensor parallel inside a node (over fast NVLink), pipeline parallel across nodes, and data parallel with ZeRO on top. Splitting along all three axes at once is called 3D parallelism — it's how trillion-parameter models fit on thousands of GPUs.
MEMORY PER GPU
Params · 15 GBGradients · 15 GBOptimizer state · 90 GB

Step up the ZeRO stage and add GPUs — watch the bar drop under the 80 GB line.

Replicated (DDP) — plain data parallel. Every GPU stores all 120 GB: params, gradients and optimizer state in full. Nothing is sharded.

MEMORY_PER_GPU
STAGEReplicated (DDP)
SHARDED OVER8 GPUs
MODEL / GPU120 GB
FITS 80 GB?doesn't fit ✕
The 7.5B-parameter model needs 120 GB replicated. Each stage shards more of that across the GPUs; the bar and this number are the memory one GPU must hold.
EXPLAIN IT BACK
Data-parallel training puts a full copy of the model on every GPU. Why does that stop you from training a really big model — and what does ZeRO-3 change?
NEXT: MATRIX OPTIMIZERS

You've split the model across the cluster. Next: the optimizer that spends all that memory — from AdamW to Muon, the newer optimizers that update weights as whole matrices, not loose numbers, and train large models faster for the same compute.

← Mixed-precision trainingMatrix optimizers →
Language: English