Guide

Scaling out

Tensor parallelism

Split each layer across GPUs so every token uses every rank. Column-parallel and row-parallel matmuls plus all-reduces do the work.

76 of 138 · 3 min read

Tensor parallelism (TP) cuts individual linear layers across GPUs so every rank takes part in every token. It is the usual way to fit a model that is too big for one GPU, and to cut decode latency by raising the memory bandwidth that one step can use.

Concept

A column-parallel multiply Y = X W shards W along the output dimension. A row-parallel multiply shards W along the input dimension, then all-reduces the partial products so every rank sees the same result.

The big idea

Imagine eight people reading one book by taking a strip of every page. Each page finishes only when all eight strips are done. That is TP: every token waits on every GPU.

The picture breaks at communication. After each layer, the ranks must exchange activations. If that exchange is slower than the math you removed, more GPUs make the step slower.

Walking through a transformer block

The Megatron-LM pattern is the workhorse:

  • In attention, shard heads (or groups of heads under GQA) across TP ranks.
  • In the MLP, the up-projection is column-parallel and the down-projection is row-parallel, so one all-reduce lands at the end of the MLP.
  • A typical layer does on the order of two all-reduces of an activation vector of hidden size H (attention residual path and MLP residual path).

Bytes moved per rank for a ring all-reduce of H elements in dtype B bytes is about 2 × (T−1)/T × H × B.

H = 8192, BF16 (B = 2), T = 8, one residual all-reduce
Payload per rank          = 8192 × 2 = 16 KiB
Ring volume per rank      ≈ 2 × 7/8 × 16 KiB ≈ 28 KiB
Two collectives × 80 layers ≈ 80 × 2 × 28 KiB ≈ 4.4 MB
per token per decode step

On NVLink (order 900 GB/s bidirectional aggregate on H100 NVSwitch nodes) a few megabytes is tens of microseconds, often fine next to a memory-bound decode matmul. Across nodes on InfiniBand, the same collective can dominate time per output token. Keep TP inside a node unless you have no choice.

TP also shards the KV cache along heads. That lowers per-GPU memory for long context, but the full cache still exists across the replica.

Trade-off

Higher TP lowers per-GPU weight memory and can cut decode latency. Lower TP reduces collective overhead and leaves room for a larger batch. Prefer the smallest TP that fits the model and meets the latency target.

Common mistakes

  • Stretching TP across InfiniBand. Collectives then dominate the step. Use pipeline parallelism across nodes instead.
  • Adding TP when the SLO is already met. Extra ranks make the replica larger and harder to schedule, and they buy nothing if communication is not cheaper than the work you remove.
  • Reading NCCL timeouts as a model bug. They are usually topology: ranks spanning sockets without NVLink, or a bad NCCL_P2P_LEVEL.

Check yourself

  1. A step does two all-reduces per layer. If you raise TP from 4 to 8 and time per token gets worse, what likely happened? Answer: collectives overtook the bandwidth win. Profile NCCL vs matmul time, then drop TP and add a replica.
  2. Does TP shrink the total KV cache of one replica? Answer: no. It shards the cache along heads, so each GPU holds less, but the replica still stores the full cache.

Key takeaways

  • TP shards every layer so every token uses every rank.
  • Column-parallel then row-parallel layers, with all-reduces on the residual paths, are the Megatron pattern.
  • Keep TP inside the NVLink domain. A few MB per token is cheap on NVLink and expensive on InfiniBand.
  • Pick the smallest TP that fits and meets time per output token.

Further learning

Next, see how pipeline parallelism splits layers across nodes. The Megatron-LM paper is the original write-up: arxiv.org/abs/1909.08053.

Parallelism for inference

Tensor, pipeline, data and context parallelism: splitting a model across GPUs so weights and KV fit, and single-request latency stays under SLO.

Pipeline parallelism

Assign blocks of layers to different GPUs or nodes. Inter-stage traffic is small, but decode leaves pipeline bubbles that waste GPU time.

Distributed training

How to train models too big for one GPU, using data, sharding, tensor, pipeline and context parallelism. Covers the memory math that forces the choice and how to judge the result.