I am scaling my distributed training setup for a new Transformer-based model and I am hitting a wall with VRAM limits on my A100 nodes. I’ve heard that tensor parallelism (TP) is more efficient than standard data parallelism for these massive architectures because it shards the actual weight matrices. However, I’m worried about the communication overhead between GPUs during the all-reduce steps. How exactly does tensor parallelism improve training efficiency at this scale, and are there specific configurations to minimize latency?
3 answers
Tensor parallelism is essential for models that simply cannot fit on a single GPU's memory. By sharding the attention heads and the feed-forward network (FFN) layers—using column-parallel and row-parallel linear layers—you effectively divide the memory footprint of the weights by the number of GPUs in your TP group. For instance, in a 175B parameter model, no single GPU can hold the parameters, activations, and optimizer states simultaneously. TP solves this by ensuring each device only processes a fraction of the matrix multiplication. To minimize the latency you're worried about, it’s best to keep TP groups within a single node to leverage high-speed NVLink interconnects, as the frequent all-reduce operations will saturate standard PCIe or ethernet links very quickly.
That’s a solid explanation of the memory benefits, but how do you handle the "bubbles" or synchronization stalls if your model has a deep pipeline? Does increasing the tensor parallel degree eventually reach a point of diminishing returns where communication outpaces the compute speedup?
In my recent benchmarks with Llama-3 architectures, implementing tensor parallelism allowed us to increase our sequence length significantly because the activation memory was also sharded across the GPUs.
I agree with Matthew; the sequence length flexibility is a huge hidden win. I’d also recommend looking at "Sequence Parallelism" alongside TP, as it helps shard the LayerNorm and dropout parts which often remain replicated in basic TP setups.
That is exactly what happens if you over-shard, Barbara. When your matrix shards become too small, the GPU kernels can't reach peak TFLOPS, and you spend more time waiting for the All-Reduce than actually calculating. In my experience, for a standard 8-GPU node, a TP degree of 4 or 8 is the limit. If you need to scale further, you should look into 3D parallelism, which combines Tensor, Pipeline, and Data parallelism to balance the load without choking the interconnect.