GPT-4 trained on ~25,000 GPUs for several months. No single GPU can hold even 1% of its parameters. The engineering trick that makes trillion-parameter models possible isn't smarter math — it's slicing the model and gradients across thousands of machines so each one only sees a fraction. Data parallelism, tensor parallelism, pipeline parallelism, ZeRO, FSDP. Pick the wrong split for your bottleneck and a $10M training run becomes $50M.
Learning Objectives
After this lesson, you will be able to:
Choose the right parallelism strategy — data parallel, pipeline parallel, tensor parallel, or sharded — based on whether your bottleneck is compute, memory, or communication
Use Fully Sharded Data Parallel (FSDP) and DeepSpeed ZeRO to fit a 100B parameter model on hardware where DDP would OOM in seconds
Apply gradient checkpointing, mixed precision (fp16/bf16), and gradient accumulation to halve memory usage with predictable compute tradeoffs
Profile a distributed training run with PyTorch Profiler to distinguish a compute-bound bottleneck from a comm-bound one — and pick the right intervention for each
Don't worry if "distributed training" sounds like a separate field -- once you understand the three flavors of parallelism (data, model, pipeline) and the memory math behind ZeRO, it all clicks. The vocabulary looks daunting but the ideas are simple.
The distributed-training literature compresses to three orthogonal techniques. Real systems combine them.
#Data Parallelism: Replicate the Model, Split the Batch
Every GPU holds a full copy of the model. The batch is split into per-GPU shards. Each GPU runs a forward and backward pass on its shard, computing local gradients. Then an all-reduce averages gradients across all GPUs before the optimizer step.
∇Lglobal=N1i=1∑N∇Liwhere each GPU i computes ∇Li on shard Bi
PyTorch has two classes for this — and one is a trap
nn.DataParallel (DP): one Python process spawns threads — Python's GIL bottlenecks everything. Don't use it.
nn.parallel.DistributedDataParallel (DDP): one process per GPU, no GIL contention, ~95% scaling efficiency. Use this.
DDP works great until the model itself doesn't fit in one GPU. Then you need sharding.
Put the first 6 transformer layers on GPU 0, layers 7-12 on GPU 1, and so on. Forward pass: GPU 0 finishes layers 1-6 on a micro-batch, hands its activations to GPU 1, then immediately starts layers 1-6 on the next micro-batch. The "bubble" — the time GPUs sit idle waiting for the pipeline to fill — is the cost.
#Tensor Parallelism: Split a Single Layer Across GPUs
Even a single transformer layer can be too big. Tensor parallelism (Megatron-LM, Shoeybi 2019) splits the matmul within a layer: the weight matrix W of shape (d_model, d_ffn) becomes W = [W_0 | W_1 | ... | W_k] sharded column-wise, with each GPU holding one slice and computing partial outputs that get all-reduced.
If W∈Rd×h,W=[W0∣W1∣⋯∣Wk−1],xW=[xW0∣xW1∣⋯]
Tensor parallelism is the most communication-heavy of the three. It only works when you have fast intra-node interconnect (NVLink, InfiniBand) — running it across slow Ethernet kills throughput.
The three parallelism strategies above all leave a memory problem: with data parallelism, every GPU stores the full model + full optimizer state. For Adam, optimizer state alone is 2x the model size. A 13B-parameter model in fp32 + Adam on 8 GPUs replicates 104GB of optimizer state per GPU that nobody benefits from.
ZeRO (Rajbhandari 2020) and PyTorch's FSDP (2021) shard the model itself across data-parallel workers. Each GPU only holds a slice. When a layer's weights are needed, an all-gather pulls the full layer to every GPU; immediately after the layer is used, the gather is freed.
ZeRO has three stages, each saving more memory at the cost of more communication:
Memory per GPU (Adam, fp16):DDP: 16N bytes (params + grads + optimizer state)ZeRO-1: 4N+12N/k(shard optimizer state across k GPUs)ZeRO-2: 2N+14N/k(also shard gradients)ZeRO-3 / FSDP: 16N/k(also shard parameters)
Decision rubric
Model fits in single GPU + you want speed → DDP
Model barely fits, optimizer state is the killer → ZeRO-1
Optimizer + gradients won't fit → ZeRO-2
Even parameters don't fit → ZeRO-3 / FSDP (or DeepSpeed ZeRO-3)
Model is so big a single layer doesn't fit → add tensor parallelism on top
What Do You Think?
You have a 13-billion-parameter model and 8 NVIDIA A100 80GB GPUs (~640GB total). You want to fine-tune with AdamW. Which strategy?
DDP looks like it should fit, but Adam needs to store 8 bytes per parameter for the moment estimates plus 4 bytes for the master fp32 copy and 2 bytes each for params and gradients in mixed precision. Total: 16 × 13B = 208GB per GPU. FSDP shards this across 8 GPUs, dropping it to ~26GB per GPU — the only realistic option. Pipeline and tensor parallel add complexity without buying you the linear sharding that ZeRO-3 gives for free.
The "16 bytes per parameter" rule is the load-bearing arithmetic of every distributed-training decision. Let's open the hood and write it down precisely, because every "will this fit?" question reduces to one inequality. The accounting (Rajbhandari et al. 2019, "ZeRO: Memory Optimizations Toward Training Trillion Parameter Models") tracks four buffers per parameter during vanilla mixed-precision training with AdamW:
Parameters (fp16/bf16). What the forward pass reads: 2 bytes
Gradients (fp16/bf16). What backward writes: 2 bytes
AdamW optimizer state. First moment m and second moment v, both kept in fp32 for numerical stability (4 + 4 = 8 bytes) plus a fp32 master copy of the parameters used for the actual weight update (4 bytes). Total: 12 bytes
Sum: 16 bytes per parameter. For a 7B model, that's 112 GB before you've stored a single activation.
Activations are the second budget. The Megatron memory paper (Korthikanti et al. 2022, "Reducing Activation Recomputation in Large Transformer Models") gives the peak activation memory per transformer layer as approximately:
The 34·b·s·h term covers the LayerNorm outputs, QKV projections, attention output, MLP intermediate, and dropout masks. The 5·b·s·h · (num_heads·s/h) term is the attention-score matrix — quadratic in seq_len, which is why long-context training is so brutal.
weights and activations along feature dim divided by G
The playground below computes this for any model size and GPU count. Edit the variables at the top — N_params, G, batch_size, seq_len, hidden_size, num_layers, num_heads — and re-run.
Loading visualization...
What Do You Think?
You have a 70B-parameter model and eight 80GB H100 GPUs. You want vanilla AdamW + bf16 mixed precision. Does plain DDP fit on this hardware?
The math here is unforgiving and exactly what you should reach for every time someone asks "will this run on N GPUs?". Numbers, not vibes.
The communication cost is the other half of the trade. ZeRO-3 / FSDP must all-gather every layer's weights immediately before forward and again before backward, then reduce-scatter the gradients. Communication volume per training step is O(N_params · num_layers) — proportional to the number of times you cross the network. For a 70B model with 80 layers and bf16 weights, that's roughly 2 · 80 · 140 GB = 22 TB of inter-GPU traffic per step. On 600 GB/s NVLink this is ~37 ms of pure comm; on 100 Gb/s Ethernet it is 30+ minutes. Pick your interconnect before you pick your sharding strategy.
Half-precision (fp16 or bf16) cuts memory and compute in half. fp16 has narrow dynamic range — gradients underflow without GradScaler. bf16 has the same exponent range as fp32 but lower mantissa precision; the modern default for transformers because no scaler is needed.
pythonreference · read-only
1
2
3
4
5
6
7
8
9
10
from torch.cuda.amp import autocast, GradScaler
scaler = GradScaler()
for batch in loader:
with autocast(dtype=torch.bfloat16): # or fp16 with scaler
out = model(batch)
loss = criterion(out, target)
scaler.scale(loss).backward() # for fp16; bf16 doesn't need scaling
scaler.step(optimizer)
scaler.update()
Activations stored during forward consume the bulk of memory in deep networks. Gradient checkpointing (Chen 2016) doesn't store them — it recomputes activations during backward. Trade: ~30% extra compute for ~10x activation-memory savings.
Need an effective batch size of 256 but your GPU only fits batch 16? Run 16 micro-batches, accumulating gradients without stepping the optimizer, then step once.
pythonreference · read-only
1
2
3
4
5
optimizer.zero_grad()
for i in range(accum_steps):
loss = model(batch[i]) / accum_steps
loss.backward()
optimizer.step() # one step per (accum_steps × micro_batch) effective batch
Three orthogonal parallelism strategies — data, pipeline, tensor — solve different walls. Data parallel scales compute, pipeline parallel splits layers, tensor parallel splits within a layer. Real systems combine all three for trillion-parameter models.
DDP duplicates everything; FSDP shards everything. DDP is fine until your optimizer state alone exceeds GPU memory. FSDP / ZeRO-3 shards parameters, gradients, and optimizer state across data-parallel workers, dropping per-GPU memory by Nx for N GPUs.
Mixed precision is free 2x speedup if you use bf16. Bf16 has the dynamic range of fp32 with half the bits; fp16 needs GradScaler to avoid underflow but is faster on older hardware.
Gradient checkpointing trades ~30% compute for 10x activation-memory savings. Mandatory above 7B parameters, optional below. The single most useful trick for fitting larger models on the same hardware.
Profile before optimizing. PyTorch Profiler tells you whether you're compute-bound (use bigger model or smaller batch), memory-bound (add checkpointing or shard with FSDP), or comm-bound (faster fabric or fewer all-reduces).
You have 8 GPUs and a 70B-parameter model. The model itself doesn't fit on a single GPU. What strategy?
This wraps the foundations of deep learning: tensors, networks, optimization, regularization, normalization, the architectural ladders for vision and sequences, embeddings, autoencoders, and now production-scale training. The next track shows how transformers reshape this entire toolkit for language — turning the same gradient descent into ChatGPT.