Distributed Training: DP, TP, PP, FSDP & Sequence Parallelism
A 70-billion-parameter model in fp32 needs 280 GB just for weights. Add Adam optimizer state and the bill climbs past 1.6 TB. An NVIDIA H100 has 80 GB. The arithmetic says you cannot train a frontier model on a single GPU, not now, not ever. The entire field of large-model training is the art of splitting that 1.6 TB of state across thousands of accelerators without letting communication eat your compute budget. This lesson opens the hood on how Meta, OpenAI, Anthropic, and DeepSeek actually do it.
#Why Parallelism Is Not Optional
The track-04 lesson on distributed training gave you the vocabulary — DDP, FSDP, gradient checkpointing, mixed precision. This lesson goes one layer deeper: how the actual frontier-lab parallelism stack (Megatron-LM, DeepSpeed, TorchTitan) composes data, tensor, pipeline, and expert parallelism into a single training job that spans tens of thousands of GPUs.
#The Three Walls
Before any parallelism strategy makes sense, we need the numbers. Take a 70B model — roughly the size of Llama 3.1 70B, Qwen 2.5 72B, or DeepSeek-V2-Lite at the small end of the frontier scale.
Memory wall. In fp32 (4 bytes/param), the weights alone are
70 × 10⁹ × 4 = 280 GB. Add gradients (another 280 GB) and Adam optimizer state (two moment buffers, 8 bytes/param) for 70 × 10⁹ × 8 = 560 GB. Plus a fp32 master copy if you're doing mixed precision. Plus activations. Total resident state with vanilla Adam: roughly 1.6 TB of GPU memory needed before you've even started training. An H100 has 80 GB. The model does not fit on 1 GPU, does not fit on 4 GPUs, does not fit on 8 GPUs. You need at least 20 H100s just to hold the state, and that is before you reserve memory for activations or any throughput headroom.Compute wall. Chinchilla-optimal training of a 70B model wants roughly 1.4 trillion tokens. At the standard estimate of 6N FLOPs/token (forward + backward), that is
6 × 70 × 10⁹ × 1.4 × 10¹² ≈ 5.9 × 10²³ total FLOPs. A single H100 at sustained 60% MFU on bf16 delivers about 0.6 × 989 × 10¹² ≈ 5.9 × 10¹⁴ FLOPs/sec. One H100 alone would take 5.9 × 10²³ / 5.9 × 10¹⁴ ≈ 10⁹ seconds ≈ 31 years. To train in 30 days, you need roughly 380 H100s running flat-out, and the linear-scaling math only works if you can keep all of them fed.Communication wall. The moment GPUs share work they must also share gradients, activations, or parameters. A single all-reduce of a 70B fp16 gradient is 140 GB of data flying across the cluster every step. At NVLink Gen4 bandwidth (~900 GB/s intra-node) this is hundreds of milliseconds; at InfiniBand HDR (~25 GB/s inter-node) it is seconds. Get the placement wrong and the network, not the matmul, becomes the bottleneck.
What Do You Think?
A 70B model in mixed-precision training with Adam needs roughly how much GPU memory for resident state (params + grads + optimizer state, no activations) before any sharding?
#Data Parallelism: The Baseline That Doesn't Scale
Data parallelism is the simplest strategy and the first one every team reaches for: replicate the full model on every GPU, hand each GPU a different batch shard, all-reduce gradients after backward, step optimizers in sync.
The math is dead simple:
This works beautifully until you remember that DDP does nothing about the memory wall. Every one of your 380 H100s holds the full 1.12 TB of resident state, which doesn't fit. DDP is the right answer when your model fits on one GPU and you just want to scale throughput. For frontier-scale models, it is the wrong layer of the stack to be solving the problem.
A subtle but real DDP gotcha is BatchNorm. Each rank by default computes BN statistics on its local shard only. If your effective per-GPU batch is small (which it usually is for big models), per-rank stats are noisy and training quality drops. The fix is
torch.nn.SyncBatchNorm.convert_sync_batchnorm(model), which adds an all-reduce of BN statistics across ranks. For transformers, the issue rarely matters because LayerNorm (operating on the feature axis, not the batch axis) doesn't have this pathology — one of the under-appreciated reasons transformers scaled cleanly while vision models needed SyncBN gymnastics.#ZeRO: Stop Replicating What You Don't Need
The DeepSpeed team's observation in 2020 was elegant: at any given moment in training, every rank is using only a small slice of the optimizer state. Why keep the rest resident in 99% of GPUs that are not currently consuming it?
ZeRO stands for Zero Redundancy Optimizer. Three stages, each more aggressive than the last:
- ZeRO-1 shards the optimizer state across data-parallel ranks. The 8 bytes/param of Adam moments live on only one rank each; everyone else carries
1/Kof the load. Each rank still keeps the full params and full gradients. - ZeRO-2 also shards gradients. After backward, each rank does a reduce-scatter so that rank
iends up owning the averaged gradient for parameter slicei(instead of every rank owning the full gradient). Params still full per rank. - ZeRO-3 (= PyTorch FSDP) also shards parameters. Each rank only holds
1/Kof the weight matrix. To run forward, ranks all-gather the full param for the current layer, compute, then immediately discard the gathered copy. Same dance again in backward.
The communication accounting matters. ZeRO-3 is bandwidth-equivalent to DDP in total bytes moved per training step (a few authors have proved this with comm-volume arithmetic), but it pays a latency tax — many small all-gathers instead of one big all-reduce. On a high-bandwidth, low-latency fabric (NVLink intra-node, fat-pipe InfiniBand inter-node), ZeRO-3 hits 90%+ of DDP throughput while opening the door to models 10x larger.
What Do You Think?
You have 8 H100 80GB GPUs and want to fine-tune a 30B model with AdamW. Which strategy is the smallest viable option?
#Tensor Parallelism: Slicing One Matrix Across GPUs
ZeRO/FSDP solves the resident-state problem by sharding across time: every rank holds a slice, and the full matrix is gathered just-in-time. But that gather still has to happen for every layer, and each gathered weight matrix still has to fit in one GPU's memory during compute. For really fat layers — think a 1.5T model with a 32k hidden dim — even a single feedforward matrix doesn't fit on one device.
Tensor parallelism (Megatron-LM, Shoeybi et al. 2019) solves this by slicing inside the matrix multiply itself, so no single GPU ever holds the full weight.
#The Megatron pattern: column-parallel + row-parallel pair
Consider a transformer feedforward block,
Y = GELU(X · W₁) · W₂. Megatron splits this in a clever way:Column-parallel for W₁ — split
W₁ ∈ R^(d × 4d) column-wise across T GPUs:Loading visualization...
Each GPU
i holds W₁ᵢ ∈ R^(d × 4d/T) and computes Hᵢ = GELU(X · W₁ᵢ). Notice we apply the nonlinearity element-wise before combining — this means no communication is needed between W₁ and the activation. Each GPU just has its own slice of the hidden activation.Row-parallel for W₂ — split
W₂ ∈ R^(4d × d) row-wise across the same T GPUs:Each GPU
i holds W₂ᵢ ∈ R^(4d/T × d) and computes the partial output Yᵢ = Hᵢ · W₂ᵢ. To get the final Y, you need to sum the partial outputs across ranks — a single all-reduce.Take a moment to internalize that matmul-as-einsum picture — it makes the column/row split obvious:
Loading visualization...
Reading
X · W as einsum("bd, df -> bf", X, W), the contracted axis is d (the inner dimension). When you column-split W along the output axis f, you literally are slicing the un-contracted index — no communication needed during the contraction itself. When you row-split W along the contracted axis d, the contraction becomes a partial sum, which is exactly what an all-reduce is computing. Tensor parallelism is the einsum lens applied at the GPU-mapping layer.#Attention parallelism: split the heads
For multi-head attention, Megatron splits the heads across TP ranks: rank
i holds nheads/T heads' worth of the Q, K, V projections and runs attention on its slice independently. The output projection Wo is row-parallel, so the final all-reduce sums partial outputs from each head group.This is why tensor parallelism's natural granularity is
nheads. Llama 3 70B has 64 attention heads — TP=8 splits evenly; TP=16 also evenly. TP=12 doesn't divide 64 cleanly, so it's avoided.#The intra-node constraint
Tensor parallelism's all-reduce happens on every transformer layer. For a 70B Llama-style model with 80 layers, that's 80 all-reduces in forward and another 80 in backward — 160 all-reduces per training step.
Each all-reduce moves
batch × seqlen × d_model × bytes_per_elem bytes across the TP group. For bs=8, seq=8192, d=8192, bf16, that's 8 × 8192 × 8192 × 2 = 1 GB per all-reduce. On NVLink Gen4 at ~900 GB/s, that's 1.1 ms. Multiplied by 160 it's 180 ms of pure comm per step — manageable if your matmul time is large enough to overlap.On InfiniBand HDR at ~25 GB/s, the same 1 GB all-reduce takes 40 ms. 160 × 40 ms = 6.4 seconds per step of pure communication, with limited overlap because the matmul time per layer is much smaller than the cross-node comm time. TP across nodes is a throughput catastrophe. This is why every production stack pins tensor parallelism to ≤ 8 ranks, exactly the size of a single NVIDIA HGX 8-GPU node connected by NVSwitch.
Quick check
Why is tensor parallelism almost always limited to intra-node (typically TP ≤ 8)?
#Sequence Parallelism: What Tensor Parallelism Misses
There's a subtle gap in tensor parallelism. Inside a transformer block you have matmuls (covered by TP), but also operations that act on the sequence dimension: dropout, LayerNorm, residual additions. These operations don't have a "hidden-dim split" — they apply identically to every position. Vanilla Megatron has every TP rank redundantly compute the LayerNorm and dropout on the full activation, which means each rank pays the full activation memory cost for these ops even though it only handles 1/T of the hidden dim elsewhere.
Sequence parallelism (Korthikanti et al. 2022, Megatron-LM v3) closes this gap by splitting the sequence axis across the same TP ranks for the LayerNorm/dropout/residual regions. The activation tensor of shape
(batch, seq, hidden) is now sharded along seq during the SP regions and along hidden during the TP regions, with all-gather/reduce-scatter pairs transitioning between the two views. The total comm volume per step is unchanged, but activation memory drops by another factor of T, which can be the difference between fitting and not fitting a long-context model.#Pipeline Parallelism: Splitting Layers Across Stages
So far DP/ZeRO and TP+SP both keep all layers on every rank in some form (DP replicates, TP shards within a layer). Pipeline parallelism instead carves the layer dimension: GPU 0 holds layers 1–10, GPU 1 holds layers 11–20, and so on. An activation from GPU 0 has to be sent over the network to GPU 1 to continue forward.
#The bubble problem
The naive PP schedule has a critical flaw. If you feed a single batch through a pipeline of
P stages:time →
GPU 0: F1 . . . . . . . B1
GPU 1: . F1 . . . . B1 .
GPU 2: . . F1 . B1 . .
GPU 3: . . . F1 . . .
^----fill----^^--drain--^
GPU 0 is idle for
P-1 time units at the end while it waits for backward to drift back. GPU P-1 is idle for P-1 time units at the start while it waits for the activation to arrive. The pipeline bubble is the fraction of wall-time GPUs sit idle. For P=8 stages with one micro-batch, the bubble is 2(P-1) / (2P) = 7/8 = 87.5% of total time — wildly bad.#GPipe: micro-batching to fill the pipe
GPipe (Huang et al. 2018) shrinks the bubble by chopping the batch into
M micro-batches that flow through the pipeline. After P micro-batches have entered, the pipeline is "full" and every stage runs every cycle. The bubble fraction becomes:#1F1B and interleaved scheduling
GPipe still has a memory issue: it runs all forward passes before any backward (so all
M activations are alive simultaneously on stage 0). PipeDream-1F1B (Narayanan et al. 2019/2020) and Megatron's interleaved 1F1B schedule alternate forward and backward steps as soon as they can — letting earlier activations be freed faster, reducing peak activation memory roughly P-fold.The "interleaved" variant further chops each pipeline stage into multiple "virtual stages" (e.g., stage 0 handles layers 1, 9, 17, 25; stage 1 handles 2, 10, 18, 26; etc.) which reduces the bubble further at the cost of more network traffic. This is what Megatron-LM v4 and TorchTitan ship today.
What Do You Think?
You have a model with 80 layers split across P=8 pipeline stages, and you're using 1F1B scheduling. Your batch is chopped into M=4 micro-batches. What fraction of the time will GPUs be idle in the bubble?
Let's actually simulate this. The playground below implements a tiny event-loop pipeline so you can see the bubble shrink as
M grows:Loading visualization...
Run that simulator: you'll see the bubble drop from ~85% at
M=1 to under 10% at M=64. This is the entire reason GPipe was a paper worth publishing.#Activation memory for PP
Pipeline parallelism trades a different resource than the other strategies. Each stage stores activations for every in-flight micro-batch, because backward for micro-batch
m needs the forward activations of m. With M micro-batches in flight and 1F1B, peak activation memory at stage 0 is M activations of one stage's worth — far less than the naive GPipe M × P.This is why PP is often combined with activation checkpointing at the per-layer or per-stage level: recompute activations during backward instead of keeping all
M copies alive. The compute overhead is ~30%, the memory savings are an order of magnitude, and the trade-off is almost always worth it for any model big enough to need PP.#Expert Parallelism: The MoE Special Case
Mixture-of-Experts models like Switch Transformer, Mixtral 8x7B, and DeepSeek-V2 break a single feedforward block into
E independent "expert" feedforward networks, with a gating network choosing the top-K experts for each token. If you have 8 experts and only 2 fire per token, you have 4x the parameter count for the same per-token compute.Expert parallelism is the natural sharding: put each expert (or a small group of experts) on a different GPU. Token routing then becomes a network operation — an all-to-all collective where every GPU sends each token to whichever GPU hosts the expert it was routed to, and receives back the activations for the tokens routed to its local expert.
The all-to-all dominates MoE training cost. Two big efficiency levers:
- Expert capacity factor. Cap how many tokens each expert handles per batch (say, 1.25x the mean). Overflow tokens are dropped or skipped — sacrificing a tiny bit of quality for predictable comm patterns.
- Locality-aware routing. DeepSeek-V2's auxiliary-loss design encourages tokens to route to experts on the same node when possible, slashing inter-node all-to-all.
#Composing the Stack: 3D and 4D Parallelism
In production, frontier labs combine all of the above into a single training job. The standard recipe is 3D parallelism = DP × TP × PP, with optional EP for MoE models.
#A real-world example: how DeepSeek-V2 sized 4D parallelism
DeepSeek-V2 is a 236B-total/21B-active MoE model trained on a 1024-GPU H800 cluster. The disclosed parallelism:
- EP = 64. 160 experts split across 64 expert-parallel groups (~2.5 experts/GPU).
- TP = 8. Tensor parallel within each 8-GPU node for dense attention and shared experts.
- PP = 8. 60 layers split into 8 pipeline stages.
- DP = 2. Only a 2-way data-parallel replica for redundancy.
The math: 64 × 8 × 8 × 2 = 8192, but the actual cluster was 1024 GPUs because EP and TP share GPUs (EP groups are the TP groups in their layout). The relevant decomposition is
1024 = TP × PP × DP = 8 × 8 × 16, with EP=64 layered on top of the FFN expert layers within the same physical GPUs. The all-to-all for EP is the largest single comm pattern in their training step, and the loss-balancing auxiliary they introduced explicitly aims to localize as much routing as possible to within-node experts.#Let's compute the memory budget for ourselves
Loading visualization...
Run it. You'll see the typical pattern: ZeRO-3 alone gets you a 70B model on 16 GPUs at ~2.4 GB per GPU of state — comfortably fitting. A 405B model needs the full 3D stack, and even then activation checkpointing is required.
#Tooling Landscape, 2026
Five frameworks own the production training stack:
| Framework | What it does best | Use when |
|---|---|---|
| PyTorch FSDP + DTensor | Native PT 2.0+ stack; FSDP=ZeRO-3, DTensor for TP/SP | Up to ~70B; you want clean PT and easy debugging |
| Megatron-LM | NVIDIA-maintained reference for TP+SP+PP, hand-tuned for H100 | You need maximum MFU on dense models at 100B+ |
| DeepSpeed | ZeRO + offload + pipeline; the venerable workhorse | Fitting huge models on small clusters via CPU/NVMe offload |
| TorchTitan | Clean composable native-PT recipe (Meta, 2024) | New projects that want all of 3D parallelism in PT idioms |
| HuggingFace Accelerate | Wrappers over FSDP/DeepSpeed for fine-tuning | Fine-tuning workflows, not from-scratch pretraining |
Colossal-AI (HPC-AI Tech) is a worthy fifth — particularly good at automatic parallelism search via its
auto-parallel engine. Alpa (UC Berkeley, now discontinued) pioneered the same idea earlier but didn't survive the transition to large-scale practice.#Diagnostics: Model FLOPs Utilization (MFU)
How do you know if your parallelism strategy is actually using the hardware well? The right metric is Model FLOPs Utilization (MFU) — introduced rigorously in Chowdhery et al. 2022 (PaLM paper).
MFU diagnoses where your stack is leaking efficiency. If MFU is 60% you are doing well — almost no time is going to comm or bubbles. If MFU is 25%, three suspects: (a) tensor parallelism is being run across nodes (kills it), (b) pipeline bubble is too large (raise micro-batch count), (c) activation recomputation cost wasn't accounted for (count rematerialization as overhead, not useful FLOPs).
Quick check
Your 70B training run is showing 22% MFU on a 256-H100 cluster. The profiler shows ~40% of step time in NCCL all-reduce. What's the most likely issue?
#Training Cost Arithmetic
Last piece of the puzzle. Given a target model and a Chinchilla-optimal token budget, how much does it actually cost?
total_flops = 6 × N_params × D_tokens
gpu_seconds = total_flops / (peak_flops_per_gpu × MFU)
gpu_hours = gpu_seconds / 3600
cost_usd = gpu_hours × cost_per_gpu_hour
At 2026 cloud prices (~$3/H100-hour on-demand, ~$1.50/hour spot for committed contracts), a Chinchilla-optimal 70B model:
- Total FLOPs:
6 × 70e9 × 1.4e12 ≈ 5.9 × 10²³ - H100 peak:
989e12FLOPs/sec, sustained at 50% MFU:4.95e14FLOPs/sec - GPU-seconds needed:
5.9e23 / 4.95e14 ≈ 1.2 × 10⁹= 332k GPU-hours - Cost: 332k × $3 ≈ $1.0 M on-demand, $500k on committed spot
For a 405B model at 15T tokens (Llama 3 scale), the same math gives ~8 million GPU-hours = $24M on-demand. Add fault-tolerance, restart overhead, hyperparameter search, ablations, and the actual headline number for training a frontier model in 2026 is comfortably $50M–$150M. Anthropic's Claude Opus 4, OpenAI's GPT-4 family, and Meta's Llama 3 405B are all in this band; o1's RL-post-training compute added significantly on top.
#Architect View: Why the Hardware Lottery Matters
#Key Takeaways
- Every parallelism strategy shards a different axis of the model. DP shards batch, TP shards hidden, PP shards layers, EP shards experts, SP shards sequence. The whole frontier-lab stack is a 4D matrix of these, composed.
- ZeRO-3 / FSDP is the memory baseline for any model > 13B. Without it you cannot fit Adam state for a 70B model on any realistic cluster. Use
transformer_auto_wrap_policy, not the size-based default. - Tensor parallelism is pinned at TP ≤ 8 because every transformer layer issues an all-reduce on the TP group, and only intra-node NVLink is fast enough. Inter-node TP is a throughput catastrophe.
- Sequence parallelism + TP is the modern Megatron-LM default, because activations on LayerNorm/dropout regions are otherwise replicated across the TP group, wasting memory at long context.
- Pipeline parallelism's bubble fraction is (P-1)/(M+P-1), so you need M ≫ P micro-batches. 1F1B scheduling shrinks the peak activation memory roughly P-fold over naive GPipe.
- Expert parallelism for MoE turns the FFN block into an all-to-all. The bandwidth-heaviest collective in modern training. DeepSeek-V2's auxiliary loss explicitly tries to keep this local.
- 3D parallelism = DP × TP × PP for dense models; 4D = + EP for MoE. Real training jobs at frontier scale are always at least 3D.
- MFU is the only honest efficiency metric. 50–60% on dense H100 training is good; below 25% something is broken in the parallelism layout.
#Quick Check
Quick check
Which parallelism strategy shards the same axis as the data parallelism replica, but does it inside the model state rather than across mini-batches?
The economics of building a frontier model are now an arithmetic problem: pick a parallelism shape that gives you 50%+ MFU on your cluster, multiply out the GPU-hours, and budget accordingly. The hard part isn't the math — it's that getting the parallelism wrong silently halves your MFU and doubles your $50M training bill. Every percentage point of MFU at frontier scale is real money.