Skip to content

Tensor Parallel#

concept

Tensor parallelism (TP) splits the matrix multiplications inside each layer across GPUs. Every GPU holds a slice of every layer's weights and computes a slice of every layer's output. It cuts weight and activation memory per GPU, but it has to communicate inside every layer, on the critical path.

How Megatron-style TP splits a transformer block#

A transformer MLP is two linear layers: Y = GeLU(X·A) then Z = Y·B.

                     column-parallel                       row-parallel
 X (replicated) ──►  A = [A1 | A2 | … | At]  ──► Yi = GeLU(X·Ai) ──► Zi = Yi·Bi ──► all-reduce(Σ Zi) ──► Z
                     each GPU i holds Ai          no communication     B split by rows    one collective
  • First linear, column-parallel: each GPU multiplies the full input by its own column slice. GeLU is element-wise, so it runs locally with no communication.
  • Second linear, row-parallel: each GPU multiplies its partial activation by its row slice, producing a partial sum.
  • One all-reduce combines the partial sums.

Attention is split the same way: attention heads are divided across GPUs (column- parallel QKV projection, row-parallel output projection), followed by one all-reduce.

Communication per layer#

Pass Collectives per transformer layer Size of each
Forward 2 all-reduces (attention + MLP) batch × seq × hidden activations
Backward 2 all-reduces same

These collectives are blocking: the next operation needs their result. Unlike DP gradient all-reduce, they cannot hide behind other compute (some recent kernels overlap them partially, but the dependency remains).

Sequence parallel: TP's companion#

Layer norm and dropout sit between the TP regions, and in plain TP their activations are replicated on every GPU. Sequence parallel (SP) splits those activations along the sequence dimension, and replaces each all-reduce with a reduce-scatter plus an all-gather of the same total volume. The result is the same communication cost and much less activation memory. Megatron enables it together with TP.

For a 70B-class model (hidden size 8192), at batch 1 and sequence 8192 in BF16, each all-reduce moves about 8192 × 8192 × 2 B ≈ 134 MB. That is four times per layer, for 80 layers, every step.

Path Bandwidth per GPU Time for one 134 MB all-reduce (8 GPUs, ideal)
NVLink 4 (H100, NVSwitch) ~450 GB/s per direction ≈ 0.5 ms
400G RDMA NIC ~50 GB/s per direction ≈ 4.7 ms

At fabric speed, TP communication would dominate the step. TP is therefore sized to the NVLink domain: TP ≤ 8 on HGX H100/H200/B200 nodes, and up to the size of the rack-scale NVLink domain on NVL72 systems.

Effect on the rest of the stack#

  • Control plane: a TP group must land on one node or one NVLink domain. Partial placement across nodes is a silent performance failure, not an error.
  • L8 interconnect: TP performance is NVLink/NVSwitch performance. nvidia-smi topo -m must show NV# between all GPUs in the group.
  • L9 accelerator and memory: TP lowers per-GPU weight memory linearly, which is why it is also the first tool for fitting large models in inference.

When to use#

  • Training: when a single layer's weights or activations are too large, usually combined with PP and DP. Typical values are 2, 4 or 8.
  • Inference: to fit a model that exceeds one GPU's memory and to cut per-token latency. vLLM and SGLang expose it as --tensor-parallel-size.
  • Avoid: across nodes on a standard RDMA fabric.