← Concept library

Training Infrastructure

Tensor and Pipeline Parallelism

How frontier labs split a model across thousands of GPUs by sharding within layers (tensor parallel) and across layers (pipeline parallel), and how to pick the split.

advanced · 10 min read · Premium

Data parallelism replicates the whole model on every GPU. Once the model is too large for that to be viable, you have to actually split the model itself. Tensor parallelism (TP) shards within a layer - matrix multiplies are partitioned across GPUs. Pipeline parallelism (PP) shards across layers - different GPUs hold different parts of the depth. Both introduce their own pathologies, and the largest training runs in the world stack DP, TP and PP into "3D parallelism" to dodge each one's failure mode.

Tensor parallelism (Megatron style)

A transformer's dominant cost is two matmuls per block: the attention projection and the MLP. Megatron-LM shards them along complementary axes so the synchronisation can be deferred until the end of the block.

For the MLP Y = GeLU(X A) B:

  1. Shard A column-wise across t GPUs. Each GPU computes GeLU(X A_i) on its slice. No comms needed because the activation is element-wise.
  2. Shard B row-wise across the same t GPUs. Each GPU computes GeLU(X A_i) B_i, which is its partial sum of the output.
  3. AllReduce the partial sums at the end of the block to get Y.

The pattern for attention is analogous (shard heads across GPUs, AllReduce after the output projection). Each transformer block costs one AllReduce in forward and one in backward, all of them inside the block.

Why TP is bandwidth-hungry

That AllReduce happens on the activations, not the gradients - so it scales with batch and sequence length, not parameter count. For long contexts and large batches, TP comms can exceed compute time unless the interconnect is NVLink-class. The unwritten rule across frontier labs: TP degree of at most 8, contained within a single node, so the AllReduces stay on NVLink / NVSwitch.

Pipeline parallelism

PP splits the model's layers across GPUs. GPU 0 holds layers 0-7, GPU 1 holds 8-15, and so on. A forward pass flows activations through the pipe; the backward pass flows gradients back.

The naive schedule (GPipe-style "fill and drain") leaves most GPUs idle most of the time. With p pipeline stages and one batch, the bubble (idle fraction) is roughly (p-1)/p. Micro-batching reclaims most of that.

Stages:    0   1   2   3
        +---+---+---+---+
batch 0 | F | F | F | F |          forward fills the pipe
batch 0 |   |   |   | B |          backward drains it
batch 0 |   |   | B |   |
batch 0 |   | B |   |   |
batch 0 | B |   |   |   |

With m micro-batches per global batch and p stages, the bubble shrinks to (p-1)/(m + p - 1). The trick is making m large without blowing up activation memory (each in-flight micro-batch holds activations).

Keep reading with Pro.

You're reading the preview. Unlock the full concept plus the library, study plans, the AI mentor, and daily emails.

Sign in to save and react.
Share Copied