Tensor Parallelism

Tensor parallelism is a model parallelism technique used to distribute the computational load of a single transformer layer across multiple GPUs. In LLMs, the most computationally intensive operations are matrix multiplications within attention and feedforward layers. Instead of computing the entire matrix multiplication on one GPU, tensor parallelism splits the weight matrices (e.g., the query, key, value projections, or the feedforward weight matrices) along their columns or rows across multiple devices. For example, a linear layer with weight W of shape [d_in, d_out] can be split into two shards each of shape [d_in, d_out/2] on two GPUs. Each GPU computes its portion of the output, and then results are combined via an all-reduce communication operation. This technique is critical for models that are too large to fit into the memory of a single GPU, such as 70B or 175B parameter models. For AI builders deploying on multi-GPU nodes (like 8x A100 or 8x H100 servers), tensor parallelism is often combined with pipeline parallelism. In vLLM and Hugging Face TGI, tensor parallelism is controlled by setting the `--tensor-parallel-size` flag, which determines how many GPUs are used for sharding. The main trade-off is communication overhead: each forward pass requires multiple all-reduce calls, which saturate GPU interconnects like NVLink and NVSwitch. Therefore, tensor parallelism is most effective on nodes with fast interconnects. It scales well up to 4-8 GPUs per node, beyond which communication becomes a bottleneck. For distributed training frameworks like Megatron-LM and DeepSpeed, tensor parallelism is a fundamental building block.