Splitting the Model Itself: Tensor and Pipeline Parallelism

Tensor parallelism splits an individual matrix multiplication across devices, so no device ever holds a whole weight matrix or computes a whole output.

There are two ways to cut a matrix, and they behave differently.

Split by columns. Give each device a vertical slice of the weight matrix. Every device needs the full input, and each produces a slice of the output columns. No communication is needed to compute, but the output arrives distributed and must be concatenated to be complete.

Split by rows. Give each device a horizontal slice. Now each device needs only the matching slice of the input, and each produces a partial sum over the full output shape. The true output is the sum of the partials, so an all-reduce is required.

Neither is attractive alone: the first needs a gather afterwards, the second needs a scatter beforehand and a reduce after. The insight that makes tensor parallelism practical is that chaining them in the right order cancels the intermediate communication entirely.

Take the feed-forward block of a transformer: an expanding matrix multiply, a nonlinearity, then a contracting matrix multiply back to the model dimension.

Split the first matrix by columns. Each device computes its own slice of the expanded hidden vector. Crucially, the nonlinearity is applied elementwise, so each device can apply it to its own slice with no knowledge of the others. No communication.

Now split the second matrix by rows. A row split needs exactly a column slice of its input, which is precisely what each device is already holding. So each device multiplies straight through and produces a partial sum.

One all-reduce sums the partials, and the block is complete.

The whole feed-forward block cost one collective operation. The intermediate hidden state, which is the largest tensor in the block, never had to be assembled anywhere.

The same pattern applies to attention: split the query, key and value projections by columns so each device owns whole heads, compute attention independently per head, then split the output projection by rows and all-reduce once.

A transformer layer under tensor parallelism therefore costs two all-reduces in the forward pass, one for attention and one for the feed-forward block, and two more in the backward pass.

That is a small number of collectives, but they happen per layer. An 80-layer model performs 160 all-reduces in a single forward pass, each one a blocking synchronisation across every participating device.

This is why tensor parallelism is confined to a single node. The collectives are frequent and sit directly on the critical path: no computation proceeds until the reduce completes, so there is nothing to hide the latency behind. On the intra-node fabric that cost is tolerable. Across a network it is not, and utilisation collapses.

The conventional degree is therefore the number of accelerators in one node, commonly eight.

A second constraint is divisibility. Splitting attention by heads requires the head count to divide by the tensor-parallel degree, and the hidden dimension must divide evenly too. Awkward model dimensions restrict your options before any performance question arises.

Tensor parallelism splits the matrix multiplies, which are most of the compute. It does not split everything.

Layer normalisation and dropout operate on the full model dimension and are not matrix multiplies, so under plain tensor parallelism every device redundantly computes them on the full activation tensor, and every device stores that full tensor. For long sequences those regions become a significant share of activation memory, and they are replicated across the entire tensor-parallel group.

Sequence parallelism splits those regions along the sequence dimension instead. Since layer normalisation is independent per token, device 0 can normalise the first slice of tokens while device 1 handles the next, with no interaction.

The two schemes alternate within each layer, and the transitions between them are all-gather in the forward pass and reduce-scatter in the backward pass, which together move the same volume as the all-reduce they replace.

So sequence parallelism removes the replicated activation memory at no additional communication cost. It is close to a free improvement, and standard in current implementations.

Pipeline parallelism takes the other obvious cut. Assign the first twenty layers to device 0, the next twenty to device 1, and so on. Activations flow forward from stage to stage, gradients flow back.

Communication is minimal: one activation tensor per stage boundary, not per layer, which is why pipeline parallelism tolerates crossing nodes.

The problem is that the stages are sequential. Device 1 cannot start until device 0 finishes, and device 0 has nothing to do afterwards until the backward pass arrives. Run one batch through a four-stage pipeline and three devices are idle at any moment.

The fix is micro-batching. Split the batch into m micro-batches and feed them in continuously, so that while stage 1 works on micro-batch 1, stage 0 is already on micro-batch 2.

Idle time remains at the start, while the pipeline fills, and at the end, while it drains. That idle fraction is the bubble, and it equals p minus 1, over m, where p is the number of stages.

The formula is small enough to reason with directly, and it settles most pipeline configuration questions.

Two readings follow. More stages is worse, since the pipeline takes longer to fill. More micro-batches is better, since the fixed fill-and-drain cost is amortised over more work.

The usual guidance is to keep m at least four times p, and preferably more. Note that the formula assumes m is at least p; with fewer micro-batches than stages the pipeline never fills and the expression stops being meaningful.

The constraint this creates is the important one. Micro-batches come from splitting the global batch, so a deep pipeline forces a large global batch, and global batch size is an optimisation decision, not a free parameter.

The original GPipe schedule runs all m forward passes, then all m backward passes. It is simple and it has a memory problem: every micro-batch's activations must be retained from its forward pass until its backward pass, so peak activation memory scales with m. Since large m is exactly what shrinks the bubble, the schedule fights itself.

The one-forward-one-backward schedule fixes this without changing the bubble at all. Once the pipeline is full, each stage alternates: one forward pass, one backward pass, repeatedly. A micro-batch's backward pass happens as soon as possible after its forward pass, so its activations are freed almost immediately.

Peak activation memory now scales with the number of stages rather than the number of micro-batches, which decouples the two and lets you raise m freely.

Interleaved 1F1B goes further on the bubble itself. Instead of one contiguous block of layers per device, each device holds several non-contiguous chunks. More chunks means a finer-grained pipeline and a smaller bubble, paid for with proportionally more stage-boundary communication.

Frontier training runs compose all three axes at once, and the layout is not arbitrary. It follows the bandwidth hierarchy.

Tensor parallelism goes innermost, inside a node. It communicates twice per layer on the critical path, so it needs the fastest links available and is capped at the node size, typically eight.

Pipeline parallelism goes next, across nodes. It communicates once per stage boundary, a small volume, so a slower link is acceptable.

Data parallelism, with ZeRO sharding, goes outermost. It communicates once per optimizer step, the most tolerant of all, and it is the axis you grow to use more hardware.

The total device count is the product of the three degrees. A run with tensor degree 8, pipeline degree 8 and data degree 64 occupies 4096 accelerators.

Reading the diagram bottom-up gives the tuning order too. Fix tensor parallelism to the node size, add just enough pipeline stages to make the model fit, and put everything else into data parallelism, since it is the axis that scales most cleanly.

A worked example makes the reasoning concrete. Suppose a 70B model, nodes of 8 accelerators with 80 GB each, and a cluster of 512 devices.

Start with tensor parallelism at 8, the node size. Each device now holds one eighth of every layer's parameters and one eighth of the per-layer activations.

Check whether the model fits with sharding alone at that tensor degree. Often it does, and if so you are finished: no pipeline stages, no bubble, and the simplest possible configuration.

If it does not fit, add pipeline stages, but only as many as needed. Each stage added raises the bubble and forces more micro-batches, which forces a larger global batch.

Whatever devices remain go to data parallelism. With tensor degree 8 and pipeline degree 4, the 512-device cluster leaves a data-parallel degree of 16.

The habit that matters most: change one degree at a time and measure. The interactions are not intuitive, and a configuration that looks better on paper frequently is not.

The three axes differ on every dimension that matters, which is why they compose rather than compete.

The row that catches people out is activations. Sharding does nothing for them, so a long-context run that keeps failing under stage 3 needs tensor or pipeline parallelism, or recomputation, and adding data-parallel ranks will not help however many are available.

The complementary trap is reaching for pipeline parallelism first because splitting layers is the most intuitive picture. It carries the bubble, constrains your global batch, and is the axis to add last.