Skip to content

Tensor Parallelism

Tensor parallelism is a way to run one neural network layer across multiple accelerators by splitting its weight matrices and having each device compute part of the same matrix multiplication. It reduces per-device memory and compute, but adds synchronising collective communication inside every transformer block.

Tensor parallelism exists because some models are too large, or too bandwidth-hungry, for one accelerator to serve efficiently. During decoding, every generated token has to pass through all transformer layers, touching large weight matrices and maintaining KV cache and workspace. Replicating the whole model on every device is simple but may not fit, and may waste memory bandwidth. TP lets several devices cooperate on the same request rather than each owning a complete copy.

The usual scheme splits linear layers in two complementary ways. For a projection Y equals XW, column parallelism divides W by output columns, so each rank produces a slice of Y and those slices can remain separate. Later, row parallelism divides W by input rows, so each rank multiplies its local input slice by its local weights. The partial outputs represent contributions to the same hidden state, so they must be summed with an all-reduce.

That all-reduce is the price people often understate. TP divides weights and much of the matrix multiply work across ranks, but the model graph now contains communication points at every layer, commonly after the attention output projection and after the MLP down projection. Those collectives are activation-sized and latency-sensitive. TP is therefore usually worthwhile inside one fast NVLink or NVSwitch domain, and often stops paying once the group spans slower inter-node networking.

Engineers meet tensor parallelism as a serving-time knob such as TP size in inference runtimes and fleet schedulers. It affects placement as much as kernel choice: the ranks in a TP group need to be near each other, with matched heads, memory, and bandwidth. It is commonly confused with data parallelism, but data parallelism runs independent requests on replicated models, while tensor parallelism splits one layer for one forward pass.

The honest tuning answer is that it depends on model width, batch shape, KV head layout, device memory, and network topology. Larger models and healthy batches can amortise communication because each rank avoids substantial weight traffic and compute. Small batches, narrow layers, or TP groups stretched across machines tend to become synchronisation-bound. Grouped-query attention also complicates sharding when there are fewer KV heads than tensor-parallel ranks.

Common questions

How is tensor parallelism different from pipeline parallelism?
Tensor parallelism splits the tensors and matrix multiplies inside a single layer, so all ranks cooperate on the same layer at the same time. Pipeline parallelism assigns different layers to different devices and passes activations between stages. TP pays frequent intra-layer collectives, while pipeline parallelism pays stage-to-stage transfers and can suffer from bubbles.
Why does tensor parallelism usually stay within one node?
Each transformer block creates collective communication that the next block must wait for. Fast local links can make that tolerable, especially when the saved matrix multiply and memory bandwidth are large. Across ordinary inter-node networking, the same all-reduces inherit higher latency, lower bandwidth, and more jitter, so the communication can erase the benefit.
Does tensor parallelism make inference faster?
Sometimes. It can speed inference when a model is too large for one device or when splitting weight reads and GEMMs saves more time than the collectives cost. It can slow inference for small batches, small models, or badly placed ranks. Treat TP as a memory and bandwidth tool with a communication bill, not an automatic speedup.
What is being all-reduced in tensor parallelism?
In the row-parallel part of a layer, each rank computes only a partial contribution to the output hidden state. The all-reduce sums those partial tensors so every rank has the same resulting activation for the next replicated part of the computation. In transformer blocks this commonly happens after attention output and MLP down projections.