Tensor Parallelism

Technical Overview of Tensor Parallelism

Tensor Parallelism (TP) is an intra-model parallelization technique that partitions weight matrices within individual neural layers across multiple accelerators (GPUs).

When modern LLMs (e.g. 70B parameters or larger) exceed the VRAM capacity of a single GPU, developers use Tensor Parallelism to partition individual matrix multiplication operations. This allows the model to split the computational workload and memory footprint of single transformer blocks across a GPU group.

Key Architecture & Implementation

Tensor Parallelism splits operations inside self-attention blocks and feed-forward networks (FFN) using column-wise and row-wise divisions:

Column Parallelism:   [ Input X ]  x  [ W_col1 | W_col2 ]  ===>  [ Y1 | Y2 ]
Row Parallelism:      [ Input Y1 | Y2 ]  x  [ W_row1 ]     ===>  [ Z ] (requires All-Reduce)
                                            [ W_row2 ]

1. Multi-Layer Perceptron (MLP) Split:
   • The first linear layer is split column-wise ($W_{col} = [W_1, W_2]$). Each GPU calculates $XW_i$ independently.
   • The second linear layer is split row-wise ($W_{row} = [W_1; W_2]$). The outputs are combined at the end of the layer using an All-Reduce operation.
2. Attention Layer Split:
   • The Query ($Q$), Key ($K$), and Value ($V$) projections are split column-wise matching attention heads.
   • The output projection layer is split row-wise, followed by an All-Reduce step before passing to the next block.

Core Parameters

• Interconnect Dependency: Extremely high. Row-parallel and column-parallel layers require multiple synchronization passes (All-Reduce) per block, making high-bandwidth interconnects (like NVLink at 900 GB/s) mandatory.
• Scaling Limit: Typically bound to a maximum of 8 GPUs (one physical server node) to avoid ethernet bottlenecks.

Real-world Applications

• Core component of Distributed frameworks like Megatron-LM, vLLM, and DeepSpeed.
• Crucial for running low-latency inference on models like Llama-3-70B on 8x H100 GPU clusters.

The Communication Cost

Tensor parallelism requires GPUs to exchange partial results with each other after every split operation — for example, after splitting a matrix multiplication across GPUs, the partial outputs need to be combined (via an all-reduce communication operation) before the next layer’s computation can proceed. This makes tensor parallelism considerably more communication-intensive than data parallelism, which only needs to synchronize gradients periodically rather than after every single layer’s computation. As a result, tensor parallelism is typically applied only within a single physical server where GPUs are connected via extremely fast interconnects (like NVIDIA’s NVLink), rather than across GPUs in different servers connected by comparatively slower networking, where the communication overhead would erase any speed benefit from splitting the computation in the first place.