Tensor and Pipeline Parallelism

Data-parallel training (see Data-Parallel Training) assumes every worker can hold a full copy of the model. That assumption breaks the moment the model itself is too large to fit in one device's memory — a modern large language model can have weights, gradients, and optimiser state that together need far more memory than any single accelerator carries. At that point splitting the data no longer helps, because there is nowhere on a single device to even load the model in the first place. The only option left is to split the model itself across devices, and there are two structurally different ways to do it.

The analogy: an assembly line versus a shared workbench

Pipeline parallelism is an assembly line. Station 1 bolts on the chassis, hands the car to Station 2 for the engine, which hands it to Station 3 for paint. Each station only ever needs the tools and parts for its own stage — a full car never has to fit inside one workshop, because no single station ever holds the whole thing. The catch: unless multiple cars are moving through at once, Station 2 sits idle while Station 1 is still working on the very first car, and Station 3 idles even longer.

Tensor parallelism is a shared workbench where several workers assemble a single engine together, each one machining a different slice of the same part simultaneously — one worker cuts the left half of a cylinder block, another the right half, and they must constantly compare notes because the two halves have to fit together perfectly. Nobody works on a separate stage; everybody works on the same operation at the same time, on their own slice of it.

The mechanics

Pipeline parallelism assigns different layers of the network to different devices. Device 1 holds layers 1–8, Device 2 holds layers 9–16, and so on; a training batch's forward pass flows through devices in sequence, and the backward pass flows back the other way. To keep every device busy rather than idle while waiting for its neighbour, the batch is split into smaller micro-batches that are fed through the pipeline in a staggered, overlapping schedule — while Device 2 processes micro-batch 1's forward pass, Device 1 is already starting micro-batch 2.

Tensor parallelism splits a single layer's matrix multiplication across devices. For a weight matrix $W$ split into column blocks $W = [W_1 \mid W_2]$ across two devices, each device computes its own partial output:

$$Y = XW = X[W_1 \mid W_2] = [XW_1 \mid XW_2]$$

In words: instead of one device multiplying the input by the whole weight matrix, each device multiplies the same input by its own slice of the columns, producing its own slice of the output — and because matrix multiplication distributes cleanly like this, concatenating the two output slices gives exactly the same result as one device doing the full multiplication. The cost is that intermediate results often need to be communicated between devices within a single layer's forward and backward pass, not just once per batch like data parallelism — far more frequent, smaller messages, which is why tensor parallelism is typically only practical between devices connected by very fast, low-latency links, such as GPUs on the same physical server.

Worked example 1: splitting a matrix multiply by hand

A tiny weight matrix $W$ (2 inputs $\to$ 4 outputs) is split by columns across two devices: $W_1$ takes the first two output columns, $W_2$ the last two.

$$W = \begin{pmatrix} 1 & 2 & 3 & 4 \\ 5 & 6 & 7 & 8 \end{pmatrix}, \qquad W_1 = \begin{pmatrix}1 & 2\\5 & 6\end{pmatrix}, \quad W_2 = \begin{pmatrix}3 & 4\\7 & 8\end{pmatrix}$$

Input $x = (1, 1)$. Device 1 computes $xW_1 = (1\cdot1 + 1\cdot5,\ 1\cdot2 + 1\cdot6) = (6, 8)$. Device 2 computes $xW_2 = (1\cdot3 + 1\cdot7,\ 1\cdot4 + 1\cdot8) = (10, 12)$. Concatenating gives $(6, 8, 10, 12)$. Check against the full multiply on one device: $xW = (1+5,\ 2+6,\ 3+7,\ 4+8) = (6, 8, 10, 12)$ — identical, confirming the split computation is mathematically exact, just distributed.

Worked example 2: pipeline idle time, a back-of-envelope count

A model is split into 4 pipeline stages, each taking 1 time unit to process one micro-batch's forward pass. With a single micro-batch (no splitting at all), the timeline is: Device 1 works during time 1, then sits idle for times 2–4 while the batch moves through the rest of the pipeline; the total wall-clock time to finish one batch is 4 units, but Device 1 was only busy 1 of those 4 units — 75% idle, sometimes called the "pipeline bubble."

Split the batch into 4 micro-batches instead. Device 1 processes micro-batch 1 at time 1, micro-batch 2 at time 2 (while Device 2 handles micro-batch 1's forward pass), and so on — every device is busy on every time step from time 4 onward. Total time to clear all 4 micro-batches through 4 stages is $4 + 4 - 1 = 7$ units (the last micro-batch still has to travel through all 4 stages), but each device is now busy roughly $4/7 \approx 57\%$ of the time instead of 25% — a large reduction in wasted idle time from splitting the batch finer, at the cost of extra bookkeeping and, per micro-batch, a smaller batch size to compute gradient estimates from.

What this means in practice

The two techniques are complementary and, at large scale, are used together with data parallelism in what's often called 3D parallelism: tensor parallelism within a server (where communication is fast enough to tolerate its frequent chatter), pipeline parallelism across servers (where the coarser, batch-boundary communication tolerates slower links), and data parallelism replicated across whole groups of servers to use the remaining hardware. Getting the split right is largely an engineering exercise in matching each parallelism strategy's communication pattern to the actual network topology available, and getting it wrong shows up as GPUs sitting expensively idle rather than as an outright failure — which makes it a common, quiet source of wasted compute budget in large training runs.