Tensor, pipeline and context parallelism

When sharded data parallelism is not enough, you split the model graph itself. Tensor parallelism splits each big weight matrix across GPUs. Pipeline parallelism assigns consecutive layers to consecutive GPUs. Context parallelism splits the sequence, so each GPU handles a slice of the tokens. Each one communicates differently, and the choice depends on the cluster.

Tensor parallelism

Tensor parallelism cuts a weight matrix into pieces. In a column-parallel linear layer, each GPU holds a slice of the columns and computes its part of the output. In a row-parallel layer, each GPU holds a slice of the rows and the partial results are summed. Megatron-LM showed that attention and MLP blocks can be partitioned this way, with a small number of all-reduces per layer.

Those all-reduces happen inside every layer, so tensor parallelism needs high bandwidth and low latency. It is usually kept within one node, where GPUs share fast links such as NVLink. Crossing nodes with tensor parallelism on a slow network can cut the efficiency of the whole job.

Pipeline parallelism

Pipeline parallelism splits the depth. Stage 0 holds the first group of layers, stage 1 the next, and so on. Microbatches flow through the stages like items on an assembly line.

A naive schedule leaves stages idle while the pipeline fills and drains. That idle time is the bubble.

• GPipe feeds many microbatches so every stage stays busy most of the time.
• 1F1B (one forward, one backward) interleaves forward and backward passes, which shrinks the bubble and limits how many activations are held at once.

More microbatches shrink the bubble, but they also keep more activations in memory and make the schedule harder to manage.

Context parallelism

Long sequences make attention activations very large. Context parallelism shards the sequence dimension, so each GPU holds a slice of the tokens. Collective operations assemble the attention results across GPUs. It matters most for long-context pretraining.

Putting the axes together

Production runs combine the axes:

total GPUs ≈ DP × TP × PP        (context parallelism can add another axis)

For example, 64 GPUs can be laid out as DP = 8, TP = 4, PP = 2. Tensor parallelism stays inside nodes, pipeline stages span nodes, and data parallelism replicates the whole mesh for throughput.

Common mistakes

• Spanning tensor parallelism across slow links. It needs frequent, fast collectives, so slow links drag the whole step down.
• Too few microbatches for the pipeline. The bubble stays large, and the GPUs sit idle.
• Adding pipeline stages without checking memory. More microbatches in flight means more stored activations.

Check yourself

1. Why is tensor parallelism usually kept inside one node? It needs frequent, low-latency all-reduces inside each layer, which fast intra-node links provide.
2. With 64 GPUs laid out as TP = 4 and PP = 2, what is DP? 64 / (4 × 2) = 8.

Key takeaways

• Tensor parallelism splits weight matrices and needs the fastest links, so it usually stays within a node.
• Pipeline parallelism splits layers into stages, and the bubble shrinks with more microbatches.
• Context parallelism splits the sequence, which helps with very long inputs.
• Total GPUs is roughly the product of the data, tensor and pipeline degrees.

Further learning

The worked example in planning a training run puts these axes on a real cluster and checks the result.