Edward Z. Yang has published an interactive explainer, hosted on ezyang.github.io, that turns the arithmetic of transformer parallelism into a set of live sliders. Instead of reading static formulas, an engineer can drag batch size, hidden dimension, layer count, or interconnect bandwidth and watch which parallelism strategy stays compute bound in real time. The guide adapts “How to Scale Your Model,” the training roofline chapters originally written by a Google DeepMind team that included Jacob Austin and Sholto Douglas, and extends the source material with a GPU fabric model and a section on expert parallelism.
The core distinction the guide establishes: data parallelism and its memory saving cousin, fully sharded data parallelism (FSDP, also called ZeRO-3), move gradients or weights around. Their communication cost per layer is fixed regardless of batch size, so a large enough per chip batch always hides the transfer behind compute. On a TPUv5p, the guide puts that floor at roughly 2,550 tokens per chip along one mesh axis, falling further as the collective spreads across more axes.
Tensor parallelism behaves differently. It shards activations, not weights, and its comms to compute ratio depends only on the feedforward width divided by network bandwidth, never on batch size. No amount of extra tokens per chip rescues a tensor parallel degree that has outrun the network; the bound is simply that the feedforward width must exceed the parallel degree times compute over bandwidth.
That asymmetry is the practical decision the guide hands an engineer planning a run. Reach for FSDP first when the batch can grow: it is nearly free on modern interconnects and its bandwidth wall recedes as more mesh axes carry the collective. Add tensor parallelism only once FSDP’s per chip batch has bottomed out, and only up to the sharding degree the fabric can carry, which the guide’s TPUv5p example puts in the tens of ways before communication overtakes the matmuls. Mixing both, the guide shows, drops the combined communication floor by roughly a factor of eight versus FSDP alone, in exchange for solving an optimal split between the two axes at a given chip count.
Pipeline parallelism inverts the tradeoff again. Its cross device traffic is a single activation hop per microbatch, cheap enough that the guide calls it close to free on the network, which is why it dominates on GPU clusters that lack TPU style dense interconnects. Its real cost shows up as idle time instead: the guide expresses the bubble fraction as pipeline stages minus one, divided by microbatches plus stages minus one, and notes that pipelining fights with FSDP because ZeRO-3 needs to gather weights on every microbatch, a step that stops paying off once each microbatch carries too few tokens to amortize the gather.
Expert parallelism gets a late addition to the guide because mixture of experts models break the assumption that FSDP’s weight traffic ignores batch size. When only a handful of experts activate per token, the effective communication load scales with the ratio of total to active experts. The guide walks through one open source MoE architecture where that ratio pushes the required batch per chip past 79,000 tokens across nodes unless the experts are sharded directly. Sharding along the expert axis, using an all to all exchange in place of an all reduce, is presented as the standard fix once that ratio gets large.
For a senior engineer, the actionable read is a checklist, not a lecture: measure your batch per chip against the FSDP floor before reaching for tensor parallelism at all, size any tensor parallel degree against your fabric’s bandwidth rather than your batch, and treat expert parallelism as mandatory rather than optional once your MoE’s total to active expert ratio makes FSDP’s implied batch requirement unreachable on your cluster.
Adapted from “How to Parallelize a Transformer for Training,” an interactive guide by Edward Z. Yang published on ezyang.github.io.
