Tensor Parallelism: Splitting the Math

Lesson 3 of 5 in Distributed Training.

Sharded data parallelism still assumes one device can at least execute a layer once its weights are gathered. What if a single layer’s matrices — or the activations flowing through them — are themselves too large, or one device is simply too slow for the math? Then you split the matrix multiplication itself: Tensor parallelism.

The canonical recipe is Megatron-LM (Shoeybi et al. 2019). Its starting observation: a Transformer block is dominated by a handful of big matrix multiplies — two in the MLP, and the projections around Attention — and a matrix multiply is easy to cut. Split a weight matrix into column blocks and give each device one block: each device computes its slice of the output using the same input, in parallel, no communication required. The cleverness is in how the pairs of matrices are cut so that an entire MLP or attention sub-block runs with only one synchronization at the end rather than one after every matrix.

One transformer MLP, split across two GPUs (Megatron-style)

  1. Input X — full copy on both GPUs

    Both devices hold the same incoming activations for this layer.

  2. GPU 0: GeLU(X·A₁)

    A₁ is the first half of A’s columns. GPU 0 produces its half of the hidden activations — no communication yet.

  3. GPU 1: GeLU(X·A₂)

    A₂ is the other half of A’s columns. Runs simultaneously with GPU 0.

  4. GPU 0: multiply by B₁

    The second matrix B is split by rows to match the column split of A. Each GPU now holds a partial sum of the true output.

  5. GPU 1: multiply by B₂

    Same shape of work on the other half — still no communication.

  6. All-reduce: add the partial outputs

    The single synchronization point of the whole MLP: partial results sum into the exact output a single device would have produced.

  7. Output Y — identical on both GPUs

Now the price tag. Data parallelism synchronizes once per step; tensor parallelism synchronizes inside every layer, in the forward pass and the backward pass. Those all-reduces carry activation-sized tensors, they sit directly on the critical path of the math, and every participating device must wait for the slowest link. Compute units stall whenever communication cannot keep up — so the interconnect’s bandwidth and latency directly set how much of your FLOPs budget turns into useful work, which you will meet on the bill as MFU.

This is why tensor parallelism almost always lives inside a single server, where accelerators are wired to each other with the fastest links a machine can offer, and why the tensor-parallel degree is usually kept small — the number of devices sharing one server, not the whole cluster. Stretch tensor parallelism across ordinary datacenter networking between servers and the stalls eat the speedup.

Why the column-then-row split — and what happens in attention

The MLP computes Y = (GeLU(X·A))·B. Why must A be split by columns and B by rows, and not the other way around? Because GeLU is nonlinear. Split A by columns and each device computes complete columns of X·A — the nonlinearity applies elementwise to values each device fully owns, so GeLU(X·A₁) is exactly correct locally. Split A by rows instead and each device would hold a partial sum of every entry; you cannot apply GeLU to partial sums, because GeLU(a + b) ≠ GeLU(a) + GeLU(b) — you would be forced to all-reduce before the nonlinearity, adding a synchronization point. The column-then-row pairing (Shoeybi et al. 2019, §3) defers all communication to a single all-reduce after B.

Attention splits even more naturally: heads are independent by construction, so the Q, K, V projections are split column-wise (each device owns a subset of heads and computes their attention entirely locally), and the output projection is split row-wise — the same one-all-reduce pattern. This is also why tensor-parallel degrees tend to divide the head count evenly.

Count the traffic: each transformer block costs a small, fixed number of all-reduces in the forward pass and the same again in the backward pass, every one of them carrying batch × sequence × d_model (model dimension)-sized activations. Multiply by depth and by every step of a months-long run, and the demand for extreme intra-server bandwidth stops looking like an implementation detail and starts looking like the whole design constraint.

Interactive checkpoint quiz (1 questions) — open this page in a browser to take it.