3D and 4D Parallelism

"I am tensor-parallel to my left, pipeline-parallel to my front, and data-parallel to my distant cousins three racks over. I contain a sixty-fourth of one layer, and I have never been more sure of my place."

A Device That Found Its Coordinates

Each preceding section of this chapter gave you one tool. Tensor parallelism cut a matrix multiply across devices so a layer too wide for one accelerator could still run. Pipeline parallelism cut the layer stack into stages so a model too deep for one device could still fit. Sharded data parallelism (Section 16.5) split the optimizer state so a replica's memory footprint shrank. Sequence and context parallelism (Section 16.7) cut the activation along the token axis so a long sequence could still be held. Used alone, each tool runs out of room: tensor parallelism stops scaling once the all-reduce per layer leaves the fast intra-node link, pipeline parallelism wastes devices in its startup bubble once the stage count grows, and pure data parallelism cannot help at all when a single replica does not fit in memory. The frontier-model answer is not to pick one tool but to combine them so each covers the others' limits.

The combination is not arbitrary. The defining idea of this section is that the device grid is laid out to match the cluster's physical topology, the same topology-aware placement principle introduced for collectives in Section 4.9. Communication on each axis is mapped onto the interconnect tier that suits its volume and frequency. Get that mapping right and a thousand-device run keeps its accelerators busy; get it wrong and the same thousand devices spend most of their time waiting on the network.

1. Stacking Three Axes Into a Grid Intermediate

Picture the devices not as a flat list but as a three-dimensional grid. Along one axis sits the tensor-parallel group: a small set of devices, typically the eight inside a single node, that jointly hold one layer and all-reduce partial results every single layer of the forward and backward pass. That traffic is large and incessant, so it must ride the fastest link the cluster owns, the intra-node NVLink-class fabric. Along a second axis sits the pipeline: each point on this axis is a different stage of the layer stack, and stages exchange only the activations at their boundaries, a few times per micro-batch. That traffic is lighter, so the pipeline axis can span the medium-speed node-to-node links. Along the third axis sits data parallelism: each point is a full replica of the entire tensor-by-pipeline arrangement, and replicas talk only once per step, when they all-reduce gradients. That is the rarest communication of all, so the data axis is placed outermost, across the slowest cross-rack fabric.

The rule is a single sentence: put the tensor-parallel group on the fastest links, the pipeline across slower links, and data parallelism outermost. It falls directly out of matching communication frequency to bandwidth. Figure 16.9.1 draws the grid and the tier mapping; the demo in Section 4 will then put numbers behind why the ordering is the cheap one.

2. The Configuration Constraint and the Combinatorial Space Intermediate

A 3D configuration is a triple of parallel degrees: the tensor degree $t$, the pipeline degree $p$, and the data degree $d$. Because every device occupies exactly one cell of the grid, the degrees are bound by a hard identity. With $D$ total devices,

This is a partition of $D$ into ordered factors, and it constrains the search sharply: only factorizations of $D$ are legal. When you add expert parallelism with degree $e$ for a mixture-of-experts model (the subject of Chapter 17) and sequence parallelism with degree $s$ for long context (Section 16.7), the grid gains axes and the constraint generalizes to

which is why practitioners speak of 4D and 5D parallelism. The number of legal configurations is the number of ordered factorizations of $D$ into the chosen count of axes, and it grows quickly. For the three-axis case, the count is bounded by summing over divisors, and for a device count with many factors it reaches into the dozens or hundreds. Each configuration has a wildly different memory footprint and communication cost, and only some are feasible at all, because a configuration that leaves each device holding more than its memory can fit simply does not run.

So the engineer faces a discrete optimization: among all factorizations of $D$ that fit in memory, find the one that minimizes time per step. There is no closed form, because the cost of each axis depends on the model shape, the batch size, and the exact bandwidths of the cluster. This is why configuration search has become its own subfield, treated as the frontier in Section 5 and as the central concern of Section 16.10, the next section, which turns the search into a tuning methodology you can apply by hand.

3. From 3D to 4D and Beyond Advanced

The third dimension is rarely the last. Two further axes appear in frontier stacks often enough that 3D is best read as a floor, not a ceiling. The first is expert parallelism. A mixture-of-experts layer replaces one feed-forward block with many experts and routes each token to a few of them; placing different experts on different devices turns the routing into an all-to-all collective and adds an axis whose degree $e$ multiplies into the device-count identity. Chapter 17 develops this axis in full, including why its all-to-all is the most placement-sensitive collective of them all. The second is sequence or context parallelism, which splits a single long sequence's activations across devices so that context lengths far beyond one device's memory become trainable; Section 16.7 introduced it, and in a 4D stack it sits as yet another factor in the product.

What matters conceptually is that none of these axes is special. Each is just another way of cutting the work, each carries a characteristic collective (tensor parallelism an all-reduce, pipeline a point-to-point send, data parallelism an all-reduce, expert parallelism an all-to-all, sequence parallelism an all-gather), and each must be mapped onto the interconnect tier its collective can afford. The device-count identity grows a factor; the placement rule stays the same. A frontier training stack is the disciplined composition of all of them, which is exactly the arrangement that Chapter 19 assembles end-to-end when it puts a real foundation-model run together.

4. A Configuration Search You Can Run Intermediate

The cleanest way to feel the combinatorial space is to enumerate it. The program below takes a device count $D$, walks every legal factorization into $(d, p, t)$, estimates each one's per-device memory and per-step communication with the simple cost models this chapter has built, discards the infeasible shapes that would exceed device memory, and reports the cheapest survivor. It is a miniature of what production auto-parallelizers do, with the cost models kept transparent enough to read line by line. The tensor-parallel bandwidth is deliberately downgraded once the group spills past one node, which is the single most important piece of topology awareness in the model.

from itertools import product

D            = 64          # total devices (must equal dp * pp * tp)
P_params     = 70e9        # model parameters
bytes_state  = 16          # bytes/param for params+grad+Adam state (mixed precision)
L_layers     = 80          # transformer layers (limits useful pipeline depth)
act_per_dev  = 3.0e9       # activation bytes one device would hold (1 stage, full batch)
mem_per_dev  = 80e9        # HBM per device (e.g. 80 GB accelerator)
GPUS_PER_NODE = 8          # devices sharing the fast NVLink domain inside one node
microbatches = 16          # pipeline micro-batches per step (sets the bubble size)
t_stage      = 8.0e-3      # compute time per pipeline stage per micro-batch, sec

BW_fast      = 600e9       # NVLink-class, bytes/sec (only WITHIN a node)
BW_pipe      = 100e9       # fast node-to-node
BW_data      = 25e9        # cross-rack fabric

def tensor_bw(tp):
    # tensor-parallel all-reduce stays on NVLink only while the group fits a node;
    # spilling past GPUS_PER_NODE drops it onto the slow node-to-node link.
    return BW_fast if tp <= GPUS_PER_NODE else BW_pipe

def divisors(n):
    return [k for k in range(1, n + 1) if n % k == 0]

def memory_per_device(dp, pp, tp):
    state = bytes_state * P_params / (dp * pp * tp)   # ZeRO-style sharding over dp
    act   = act_per_dev / (pp * tp)                   # pp cuts layers, tp cuts within
    return state + act

def comm_cost(dp, pp, tp):
    layer_bytes = bytes_state * P_params / L_layers
    t_tensor = (tp - 1) * layer_bytes / (tp * tensor_bw(tp)) if tp > 1 else 0.0
    t_pipe   = (pp - 1) * act_per_dev / (pp * BW_pipe)       if pp > 1 else 0.0
    t_bubble = (pp - 1) / microbatches * t_stage            if pp > 1 else 0.0
    grad_bytes = bytes_state * P_params / (pp * tp)
    t_data   = (dp - 1) * grad_bytes / (dp * BW_data)       if dp > 1 else 0.0
    return t_tensor + t_pipe + t_bubble + t_data

ranked = []
for tp, pp in product(divisors(D), divisors(D)):
    if D % (tp * pp) != 0:
        continue
    dp = D // (tp * pp)
    if pp > L_layers:                         # cannot have more stages than layers
        continue
    mem = memory_per_device(dp, pp, tp)
    if mem > mem_per_dev:                      # infeasible: would OOM
        continue
    ranked.append((comm_cost(dp, pp, tp), dp, pp, tp, mem))
ranked.sort()

print(f"devices D                 : {D}")
print(f"feasible configurations   : {len(ranked)}")
cost, dp, pp, tp, mem = ranked[0]
print(f"best (dp, pp, tp)         : ({dp}, {pp}, {tp})   product = {dp*pp*tp}")
print(f"  memory / device         : {mem/1e9:6.2f} GB   (cap {mem_per_dev/1e9:.0f} GB)")
print(f"  comm time / step        : {cost*1e3:6.2f} ms")
print("\ntop 4 feasible shapes (dp, pp, tp -> comm ms, mem GB):")
for cost, dp, pp, tp, mem in ranked[:4]:
    print(f"  ({dp:2d}, {pp:2d}, {tp:2d})  ->  {cost*1e3:6.2f} ms   {mem/1e9:6.2f} GB")
three_d = [r for r in ranked if r[1] > 1 and r[2] > 1 and r[3] > 1]
if three_d:
    cost, dp, pp, tp, mem = three_d[0]
    print(f"\nbest genuine 3D shape     : (dp={dp}, pp={pp}, tp={tp})  ->  {cost*1e3:6.2f} ms")

devices D                 : 64
feasible configurations   : 28
best (dp, pp, tp)         : (1, 8, 8)   product = 64
  memory / device         :  17.55 GB   (cap 80 GB)
  comm time / step        :  50.17 ms

top 4 feasible shapes (dp, pp, tp -> comm ms, mem GB):
  ( 1,  8,  8)  ->   50.17 ms    17.55 GB
  ( 1, 16,  4)  ->   53.12 ms    17.55 GB
  ( 1, 32,  2)  ->   56.23 ms    17.55 GB
  ( 1, 64,  1)  ->   61.03 ms    17.55 GB

best genuine 3D shape     : (dp=2, pp=4, tp=8)  ->  744.42 ms

The winning shape, $(d, p, t) = (1, 8, 8)$, is exactly the textbook arrangement: a tensor-parallel group of eight devices confined to one node where the per-layer all-reduce stays on the fast link, and a pipeline of eight stages spanning the slower node-to-node fabric. The moment the search tries to widen the tensor group past a node, the tensor_bw downgrade fires and the cost climbs. The moment it introduces a data axis, the gradient all-reduce over the slow cross-rack fabric dominates everything, which is why the only genuine three-axis shape costs roughly fifteen times more per step here. That is not a flaw in 3D parallelism; it is the search correctly telling you that at this device count and model size, the extra axis is not yet earning its communication. Add more devices or a larger model and the balance shifts. The point is that the answer is a search result, not a fixed recipe.

# Run with: torchrun --nproc_per_node=8 --nnodes=8 thisfile.py   (64 devices)
from torch.distributed.device_mesh import init_device_mesh

# Name the three axes; the product of the sizes must equal the world size.
mesh = init_device_mesh("cuda", (1, 8, 8), mesh_dim_names=("dp", "pp", "tp"))

tp_group = mesh["tp"].get_group()   # the 8 devices that all-reduce each layer
pp_group = mesh["pp"].get_group()   # the 8 pipeline stages
dp_group = mesh["dp"].get_group()   # the data-parallel replicas (size 1 here)

5. Letting a Machine Choose the Shape Advanced

Code 16.9.1 brute-forced 28 candidates because a 3D space over 64 devices is small. Real stacks face a far larger space: five or more axes, dozens of model-shape choices, micro-batch counts, recomputation policies, and overlap schedules, multiplied across hundreds or thousands of devices. Enumerating it by hand stops being possible, and the search itself becomes the engineering problem. This is the domain of auto-parallelization, and it is an active frontier precisely because a good configuration can mean the difference between a run that finishes in a week and one that finishes in two.

With the composition rule, the device-count constraint, and the search in hand, the remaining question is practical: faced with a specific model and a specific cluster, how do you actually pick and tune the degrees without running every candidate to convergence? That is the methodology of the next section. Section 16.10 turns the cost-model intuition of Code 16.9.1 into a step-by-step strategy for choosing and tuning a parallelism configuration, the capstone that closes this chapter.