Splitting the Model Itself: Tensor, Pipeline, and Sequence Parallelism

Tensor parallelism: split a matmul across devices

Data parallelism and sharding both keep each layer's math whole on one device. tensor parallelism breaks that assumption: it partitions the weight matrices inside a layer so that the matrix multiply itself is computed jointly by several devices. For a transformer feed-forward block `Y = GeLU(X W1) W2`, you split `W1` by columns and `W2` by rows; each device computes a partial result, and a single all-reduce at the end sums them into the correct output. The attention block splits naturally along heads. This is the core trick of Megatron-LM, and it is what lets a layer too big for one device run across, say, eight GPUs inside a server.

Tensor parallelism splits the weight matrix so each device holds only a slice; the partial products Y_iB_i must then be summed by an all-reduce that sits in the critical path.

Sequence parallelism: shard the leftover activations

Tensor parallelism shards the matmuls, but the operations between them — layer norm, dropout, the residual adds — are still replicated, and each device holds the full activation for those regions. At long sequence length that replicated activation becomes the new bottleneck. sequence parallelism closes the gap by sharding those activations along the sequence dimension across the same tensor-parallel group, converting the all-reduce at the region boundaries into an all-gather plus reduce-scatter pair that moves the same volume of data while keeping memory sharded throughout. In practice it composes with tensor parallelism almost for free and is essential for long-context training.

Pipeline parallelism: split the layers into stages

Tensor parallelism splits within a layer; pipeline parallelism splits across layers. You assign, say, layers 1–8 to stage A, 9–16 to stage B, and so on down a chain of devices. A batch flows forward through the stages and gradients flow back, like an assembly line. The danger is the pipeline bubble: while stage A computes the first micro-batch, stages B, C, D sit idle with nothing to work on yet. The cure is to chop the batch into many micro-batches and keep them in flight at once, so that once the pipeline is full, every stage is busy. The bubble fraction shrinks as you increase the number of micro-batches relative to the number of stages.

# 1F1B schedule: interleave forward and backward to bound memory
for step in pipeline_schedule(num_microbatches, num_stages):
    if step.is_forward:
        act = stage.forward(recv_from_prev())
        send_to_next(act)
    else:                          # backward of an earlier microbatch
        grad = stage.backward(recv_from_next())
        send_to_prev(grad)

The pipeline bubble: the idle 'fill and drain' fraction shrinks as the number of microbatches m grows relative to the number of stages p.

Recomputation: the universal memory discount

Whichever split you use, activations dominate memory at scale, and the lever that always works is activation recomputation (also called gradient checkpointing). Instead of storing every intermediate activation for the backward pass, you store only a few checkpoints — typically the input of each transformer block — and recompute the rest on the fly during backprop. The trade is explicit: roughly one extra forward pass of compute (about a 30% slowdown if you checkpoint every block) in exchange for cutting activation memory by an order of magnitude. Selective recomputation, which only recomputes the cheap-but-bulky operations, gives most of the saving for a fraction of the cost.

Gradient checkpointing trades compute for memory: keeping √N checkpoints instead of all N activations cuts activation memory to O(√N) at the cost of one extra forward pass.