Tensor parallelism splits individual weight matrices across several GPUs so that they compute one layer together, each handling part of the matrix and exchanging results. It is a form of model parallelism that operates within layers, as opposed to pipeline parallelism, which splits between layers. Introduced at scale by NVIDIA’s Megatron-LM work in 2019, it is now standard for training and serving very large models.
Key takeaways
- Each GPU holds a slice of every weight matrix, not whole layers.
- Devices must communicate within every layer, so links must be fast.
- It is usually kept inside one server, across NVLink or similar interconnects.
- It reduces both memory per GPU and per-token latency.
- Tensor-parallel size is typically 2, 4 or 8, matching GPUs per node.
Tensor parallelism explained in simple terms
Four accountants work on one enormous spreadsheet. Rather than each taking different pages, they each take a quarter of every page’s columns, compute their part, and combine totals at the end of each row. They must talk constantly, so they sit at the same table, not in different buildings.
How it works
Take a feed-forward layer that multiplies input by a large weight matrix.
- Split the matrix by columns across GPUs; each computes part of the output.
- Each GPU applies the activation function to its slice.
- The second matrix is split by rows, so each GPU produces a partial result for the full output.
- An all-reduce sums the partial results across GPUs.
- Every GPU now holds the complete layer output and proceeds to the next layer.
Attention layers split naturally too: different heads go to different GPUs, with an all-reduce after the output projection.
A worked example
A 70-billion-parameter model needs about 140 GB for 16-bit weights, more than a single 80 GB GPU can hold. With tensor parallelism across 4 GPUs:
- Each GPU stores about 35 GB of weights.
- Each holds a quarter of every matrix.
- Each layer requires communication, so the GPUs must be tightly connected.
The model now fits, and each token is produced faster than one GPU could manage, since four devices share the work.
Why fast interconnect matters
Tensor parallelism communicates at least twice per transformer layer. In a model with 80 layers, that is well over a hundred synchronisations per token. Over NVLink inside a server, this is fast enough to be worthwhile. Over ordinary Ethernet between servers, the communication cost usually outweighs the benefit, which is why tensor parallelism stays within a node and pipeline or data parallelism is used across nodes.
Tensor vs pipeline parallelism
| Aspect | Tensor parallelism | Pipeline parallelism |
|---|---|---|
| Splits | Inside each layer | Groups of layers |
| Communication | Frequent, small, every layer | Less frequent, activations between stages |
| Interconnect need | Very high | Moderate |
| Latency effect | Reduces per-token latency | Can add latency through stages |
| Typical scope | Within a server | Across servers |
At inference
Serving engines such as vLLM and TensorRT-LLM support tensor parallelism with a setting for how many GPUs to use. It is the standard way to serve models that exceed single-GPU memory, and it also improves tokens per second for a single request, which pipeline parallelism does not.
Limitations
- Communication overhead grows with the number of GPUs.
- Efficiency falls sharply without high-bandwidth links.
- Some dimensions must divide evenly by the parallel size.
- More complex deployment and debugging than single-GPU serving.
