Stacking the Axes: 3D Parallelism and Mixture-of-Experts

Why you compose rather than choose

Each parallelism axis hits a wall on its own. Tensor parallelism needs the fastest links and stops being efficient past a single server. Pipeline parallelism is cheap on bandwidth but its bubble grows if you stretch it too thin. Data parallelism scales batch but not model size. 3D parallelism composes all three into a single device grid: tensor-parallel within a server, pipeline-parallel across a handful of servers to span the model's depth, and data-parallel across the resulting model replicas to consume the dataset. A trillion-parameter run might use TP=8, PP=16, and DP=128 — and the product of those is the number of GPUs.

Pick TP to fit one layer's tensors in a server, bounded by the intra-node links.
Pick PP so the model's full depth, divided by TP, fits across your inter-node links with an acceptable bubble.
Let DP absorb whatever GPUs remain, scaling the global batch toward the gradient-noise-scale target.
Add activation recomputation and ZeRO-style optimizer sharding over the DP axis to close any remaining memory gap.

N_{\text{GPU}} = d_{p}\cdot t_{p}\cdot p_{p}

The total device count factors into the product of the data-, tensor-, and pipeline-parallel sizes — the 3D grid you compose rather than choose.

Overlap is what keeps the grid efficient

With three or four collectives running on different axes, the cluster only stays utilized if their traffic hides behind compute. communication–computation overlap becomes a scheduling problem across axes: the data-parallel reduce-scatter overlaps the backward pass, tensor-parallel all-reduces overlap nothing (they are on the critical path, which is why TP stays small and local), and pipeline point-to-point sends overlap the compute of adjacent micro-batches. Model FLOPs utilization (MFU) — the fraction of peak compute you actually convert into useful training — is the single number that tells you whether your 3D layout is sound; well-tuned trillion-scale runs report MFU in the 40–55% range.

A fourth axis: expert parallelism

Sparse Mixture-of-Experts (MoE) models replace a dense feed-forward layer with many expert sub-networks, of which a router activates only a couple per token. This decouples parameter count from per-token compute — but it means each layer holds dozens of experts that may not fit on one device. expert parallelism spreads the experts across devices, and because each token must reach whichever experts its router chose, the layer is bracketed by an all-to-all collective that shuffles tokens to their experts and back. Expert parallelism slots in as one more axis alongside the 3D grid, and the all-to-all becomes the dominant communication cost of an MoE layer.

y = \sum_{i \in \operatorname{TopK}(g(x))} g_i(x)\, E_i(x)

A sparse MoE layer's output is a top-k weighted sum over only the experts the router selects for that token.

Faster arithmetic: training in FP8

All this parallelism distributes work; FP8 training makes each unit of work cheaper. Modern accelerators run 8-bit floating-point matmuls at roughly double the throughput of bf16, so casting the big linear layers to FP8 directly buys speed. The catch is dynamic range: with only a handful of mantissa and exponent bits, gradients and activations can over- or underflow. The standard remedy is per-tensor scaling — track a scale factor for each tensor so its values land in FP8's representable window — applied to the matmul inputs while accumulating in higher precision and keeping a master copy of the weights. FP8 composes cleanly with every parallelism axis, since it changes only the numerics of the local compute, not the communication pattern.

x_{\text{fp8}} = \operatorname{cast}\!\left(\frac{x}{s}\right),\quad s = \frac{\max|x|}{448}

Per-tensor amax scaling rescales each tensor into FP8's E4M3 range (max ≈ 448) before the cheaper 8-bit matmul.