Tensor Parallelism: Why TP=1→TP=4 Gave 2.1× Instead of 4×
The 2.1× problem
A compute platform team at a Series B AI company upgraded their training cluster from TP=1 to TP=4 before a long model training run. The math looked straightforward: four GPUs share the weight matrices, each handling one quarter of the computation, so throughput should improve by approximately 4×.
The actual improvement: 2.1×.
The team ran the upgraded configuration for three days before reviewing the profiling data. The model training was completing faster than before — just not by the expected margin. The on-call ML engineer opened Nsight Systems and pulled the timeline. The GPU compute cores were busy. The NVLink interconnect between GPUs was also busy. Sometimes both at once, sometimes in sequence. The pattern was clear once visible: the GPUs were spending 47% of each training step waiting for AllReduce communications across NVLink.
The team had optimised the wrong thing. Before the TP=4 upgrade, they had spent two weeks tuning their attention kernels. The kernels were fast. But at TP=4, each forward pass required four AllReduce operations — one per tensor-parallel layer — and each AllReduce had to complete before the next layer could start. The bottleneck had moved from compute to communication. The kernel optimisations were contributing zero to the actual wall clock time.
The specific issue: their NVLink topology was not fully connected. The cluster used NVLink 3.0 with a switch fabric that provided full bandwidth between GPUs 0-1 and between GPUs 2-3, but routed cross-pair traffic (GPU 0 to GPU 2) through the switch at reduced bandwidth. At TP=2, all communication stayed within pairs. At TP=4, half the AllReduce traffic crossed the switch.
The team fixed it by rearranging GPU assignments to keep TP=4 groups within fully-connected NVLink domains. Training throughput jumped from 2.1× to 3.4×. The remaining gap from the theoretical 4× was accounted for by the AllReduce overhead that was now unavoidable: at TP=4, there is always some communication overhead. But 3.4× instead of 2.1× recovered the majority of the missing performance.
What the team was missing: a model of which GPUs could communicate at full NVLink bandwidth before choosing the tensor parallelism degree.
Why tensor parallelism breaks the “more GPUs = more speed” assumption
The intuition behind tensor parallelism is compelling and partially correct. Splitting a giant weight matrix across four GPUs means each GPU handles one quarter of the computation. Four times the compute hardware should yield four times the throughput. This intuition holds in a vacuum. Real clusters are not vacuums.
Consider a construction crew analogy. Four workers each build one quarter of a wall simultaneously. They finish in one quarter of the time — as long as they never need to talk to each other. Now add a requirement: after every ten bricks, all four workers must stop and report their progress to each other before continuing. The synchronisation point changes everything. If the reporting takes as long as laying ten bricks, the four-worker crew is only twice as fast as one worker, not four times.
Tensor parallelism splits weight matrices column-wise or row-wise across GPUs. The forward pass computes partial results on each GPU, then combines them via AllReduce before passing to the next layer. AllReduce is the synchronisation point. Its cost is determined by the NVLink bandwidth between GPUs, the message size (proportional to the layer’s hidden dimension), and the number of participating GPUs.
Here’s the key thing: tensor parallelism efficiency is not determined by the number of GPUs but by the ratio of compute time to AllReduce time — and that ratio changes with every layer, every model architecture, and every NVLink topology.
Tensor parallelism communication pattern
Github Link:
https://github.com/sysdr/AISysDesign-Labs/tree/main/distributed_training/issue-18
