Hybrid Parallelism Math

Why one axis is never enough

Each parallelism axis solves one problem and creates another. Data parallelism (DP) replicates the whole model and splits the batch; it scales throughput but does nothing for memory — a model that does not fit on one device still does not fit under pure DP. Tensor parallelism (TP) shards the weight matrices of each layer, cutting per-device memory, but it communicates activations on every layer, so its bandwidth demand is brutal. Pipeline parallelism (PP) splits the model by layers into stages, cheap on communication, but introduces a bubble of idle time at the start and end of each batch.

Because each axis is strong exactly where another is weak, the answer is not to pick one — it is to combine them. TP shrinks the per-device footprint enough to fit a layer; PP fits the depth that TP cannot; DP then replicates that whole TP×PP unit across the remaining devices for throughput. Hybrid parallelism is the disciplined stacking of these axes so the cluster is both feasible (it fits) and efficient (it stays busy).

The device mesh: N = DP x TP x PP

The organizing idea is a device mesh: arrange your N accelerators as a logical grid whose dimensions are the parallelism axes. If you have N = 512 GPUs, a mesh of DP=8 × TP=8 × PP=8 assigns every GPU a coordinate (d, t, p). The only constraint is that the factors multiply to the device count:

N = DP × TP × PP  (× EP for MoE)
512 = 8  ×  8 ×  8

Each axis defines a communication group: TP ranks that share a layer form one group, PP ranks that pass activations stage-to-stage form another, DP ranks that all-reduce gradients form a third, and a given GPU belongs to one group per axis at once. The entire art of hybrid parallelism is choosing this factorization and then placing the groups onto the physical network so the chattiest axis lands on the fastest links. Get the factorization right and everything fits; get the placement right and it runs fast.

The communication-volume ranking

To place the axes you must rank them by how much they communicate — the ordering is stark, and it drives every layout decision.

Tensor parallelism is the heaviest. Each transformer layer does a pair of all-reduces on the activation tensor in the forward pass and another pair in the backward pass. The message size is roughly batch × seq × hidden per collective, and it happens for every one of the model’s dozens of layers, every microbatch. That is an enormous, latency-sensitive volume.

Data parallelism is next. It all-reduces the gradients once per optimizer step; the volume equals the parameter count, but it is amortized over the whole step and overlaps with backward compute.

Pipeline parallelism is the lightest. It sends only the boundary activations between adjacent stages — point-to-point, not a collective, and only at stage seams. The ranking TP » DP > PP is the single most important fact in the subject: it tells you which tier of your network can afford to host each axis.

Placing axes onto the topology

Real clusters are hierarchical: 8 GPUs inside a node share ultra-fast NVLink (hundreds of GB/s), while nodes talk over InfiniBand or Ethernet an order of magnitude slower. Match the axis ranking to this hierarchy:

This is why TP is almost always set to the node size (8) and never larger — the moment a tensor-parallel group spans two nodes, every layer’s all-reduce crosses the slow fabric and throughput collapses. Pipeline parallelism, sending only thin activation slices between neighbours, is the natural axis to stretch across nodes, and data parallelism — communicating once per step with compute to hide behind — sits happily on the outermost, slowest tier.

Pipeline parallelism and the bubble

Pipeline parallelism splits the layer stack into P stages placed on different device groups; a microbatch flows stage 1 → 2 → … → P and its gradients flow back. The catch is startup: while stage 1 processes the first microbatch, stages 2..P sit idle, and the same drains at the end. With m microbatches and the 1F1B schedule the wasted fraction is:

bubble fraction = (P - 1) / (m + P - 1)

With P=8 stages and only m=8 microbatches the bubble is 7/15 ≈ 47% — nearly half the pipeline idle. Push to m=64 microbatches and it falls to 7/71 ≈ 10%. So PP buys memory (each device holds only 1/P of the layers) at the cost of needing a large microbatch count — and thus a big global batch — to stay efficient. That is the central tension PP adds to the hybrid design: more stages fit deeper models but demand more microbatches to keep the bubble small.

Data parallelism and ZeRO/FSDP as the outer axis

Once TP and PP have made one model copy fit across a TP × PP block of devices, data parallelism replicates that block to use the rest of the cluster and multiply throughput. Classic DP holds a full copy of the model states on every replica, which is wasteful. ZeRO (and its PyTorch form, FSDP) instead shards the model states across the DP group: ZeRO-1 shards optimizer states, ZeRO-2 adds gradients, ZeRO-3 adds the parameters themselves, gathering each layer’s weights just-in-time for its forward and backward.

The consequence is large: DP stops being purely a throughput axis and becomes a third memory-saving axis, dropping per-device optimizer memory by another factor of DP. The cost is extra parameter all-gathers during the step (ZeRO-3), which is why DP still sits on the outer, slower tier where that communication can overlap with the compute of a full forward-backward pass.

Expert parallelism for MoE

Mixture-of-Experts models add a fourth axis. An MoE layer holds many expert FFNs but routes each token to only one or two, so the parameter count balloons while the compute per token stays modest. Expert parallelism (EP) places different experts on different devices; a routing step then performs an all-to-all to send each token to the device holding its chosen expert, and a second all-to-all to bring the results back.

EP composes with the others by carving the experts out of the existing DP × TP space rather than adding fresh devices: GPUs that were data-parallel replicas of a dense FFN instead each own a slice of the experts. Its distinctive cost is that all-to-all traffic, which is sensitive to routing balance — if a popular expert draws too many tokens, its device becomes a straggler. So EP adds a load-balancing concern on top of the usual tradeoff, and its all-to-all wants fast links much as TP does.

A worked memory budget

Make it concrete with a 70B-parameter model in mixed precision. Adam needs, per parameter, a bf16 weight and gradient (2 + 2 bytes) plus fp32 master weight, momentum, and variance (4 + 4 + 4 bytes) — about 16 bytes. So the model states alone are:

70e9 params × 16 bytes ≈ 1.12 TB of model state

No single 80 GB device comes close. Now shard it on the DP=8 × TP=8 × PP=8 mesh. TP and PP each cut the resident parameter memory by their factor, and ZeRO across DP shards the optimizer states:

per-device model state ≈ 1.12 TB / (TP × PP × DP)
                       = 1.12 TB / 512  ≈ 2.2 GB

That leaves the bulk of each 80 GB device free for activations, whose memory grows with batch, sequence, and the layers held on that device — and which checkpointing trades against recompute. The lesson: the three axes multiply their memory savings, and it is that product, not any one axis, that turns an impossible 1.12 TB into a comfortable per-device budget.

Choosing the split for a cluster

Given a real cluster, the factorization follows a fixed order of reasoning rather than guesswork. Work from the tightest constraint outward:

1. Set TP to the node size. Choose the largest TP that stays inside one NVLink domain — usually 8. This is your first, cheapest memory cut, and going past the node is forbidden by the bandwidth ranking.

2. Set PP to fit the depth. Increase pipeline stages until model_layers / PP worth of parameters and activations fit on a device after the TP cut. Keep PP as small as it can be, because every extra stage worsens the bubble and demands more microbatches.

3. Let DP absorb the rest. Whatever devices remain, DP = N / (TP × PP), become data-parallel replicas, with ZeRO for the extra memory relief.

4. Size the global batch so the microbatch count keeps the pipeline bubble under ~10%. This is where the axes interact: a bigger PP needs a bigger batch, which changes the DP math and the optimizer dynamics.

Common pitfalls

The recurring mistakes all come from violating the ranking or the mesh. Setting TP larger than a node is the classic one: it fits the model but every layer’s all-reduce now crawls across InfiniBand and throughput falls several-fold. Too few microbatches leaves the pipeline bubble dominating — a P=16 pipeline with m=16 is idle almost half the time. Forgetting activation memory is common: a mesh that fits the weights can still OOM on activations at long sequence length unless you checkpoint.

Others: assuming DP scales linearly forever when the gradient all-reduce eventually saturates the outer network, and letting MoE routing go unbalanced so one expert’s device stalls the all-to-all. Frameworks like Megatron-LM, DeepSpeed, and FSDP encode these rules, but they still ask you for the mesh dimensions — and a bad factorization silently wastes half the cluster.