Putting It Together: 3D Parallelism
Lesson 5 of 5 in Distributed Training.
No large run picks one axis — it stacks all three, because each solves a problem the others cannot. The standard composition, usually called 3D parallelism, matches each axis to the network tier it can afford: tensor parallelism inside each server on the fastest links, pipeline parallelism across servers to hold the full depth of the model, and data parallelism (often ZeRO-sharded) replicating that whole arrangement to soak up the rest of the cluster and the training data.
The degrees multiply. Choose 8-way tensor, 4-way pipeline, and 32-way data parallelism, and you have described 8 × 4 × 32 = 1,024 devices and, implicitly, the entire cluster topology: groups of 8 sharing a server, chains of 4 servers forming one model replica, and 32 such replicas averaging gradients. When a systems paper states its configuration in one terse line, this is the decoding — and models with Mixture of experts (MoE) layers add a fourth axis on the same principle, expert parallelism, placing different experts on different devices.
| Constraint you hit | Axis that relieves it | What it costs | Where it runs |
|---|---|---|---|
Run too slow; model fits on one device | Data parallelism (plain) | Gradient all-reduce once per step; global batch grows with replica count | Across the whole cluster — tolerates the slowest links |
Model states don’t fit; single layers still fine | Gather/scatter traffic woven through each step (≈1.5× plain DP at full sharding) | Across the data-parallel group | |
One layer’s math or activations overwhelm a device | Tensor parallelism | All-reduces inside every layer, forward and backward — the chattiest axis | Inside one server, on the fastest links |
Too many layers per device; need depth across servers | Pipeline parallelism | Bubbles — idle ramp time amortized by micro-batches; activation storage or recompute | Across servers — only boundary activations cross the wire |
One more thing the composition buys you: an honest read of training reports. When a logbook mentions lost nodes, restarts from a Checkpoint, or throughput dips after a topology change, you can now locate the pain on a specific axis — a dead server takes out one pipeline stage of one replica; a slow link between servers drags one data-parallel group; a mis-set micro-batch count shows up as bubble. Teams that published their training logbooks — OPT’s is a well-known public example (Zhang et al. 2022, arXiv:2205.01068) — describe hardware failures and restarts as routine facts of cluster-scale life, which is why checkpoint cadence is a first-class design decision, not an afterthought. What failure does to the numbers — loss spikes, precision, recovery — is the next module’s territory.
In production
Assembling 3D parallelism by hand is exactly what managed training platforms increasingly abstract: cluster provisioning, device placement, health monitoring, restart-on-failure. But the abstraction has a floor — the axes, their degrees, and the batch and checkpoint decisions remain yours.
AWS
Amazon SageMaker runs managed distributed-training jobs and provides distributed-training libraries for sharded data parallelism and model parallelism; managed cluster features monitor node health and resume runs from checkpoints. What it does not decide for you: parallelism degrees, global batch size, checkpoint cadence, and whether your job’s communication pattern matches the cluster’s network layout — placement of tensor-parallel groups within a server remains the difference between compute and stalls.
Azure
Azure Machine Learning launches distributed PyTorch jobs across GPU clusters built with high-bandwidth interconnects between nodes, and integrates DeepSpeed — Microsoft’s training library whose ZeRO implementation is the one this module described. The platform schedules and supervises; choosing the ZeRO stage, tensor/pipeline degrees, and micro-batch count — and having a restart plan when a node dies mid-run — is still your engineering.
Google Cloud
Vertex AI custom training jobs distribute across GPU clusters, and TPU pods shard along the same conceptual axes through the XLA compiler stack — you annotate how tensors and data split, and the compiler emits the collectives. Different toolchain, identical vocabulary: something must still declare how the model splits across devices, and checkpoint frequency still bounds how much work a failure can destroy.
Key terms: Data parallelism, Tensor parallelism, Pipeline parallelism, ZeRO, Fully sharded data parallel (FSDP), Mixed precision
Interactive checkpoint quiz (1 questions) — open this page in a browser to take it.