Tensor Parallelism 101: Multi-GPU Inference for Large Language Models

You have a shiny new 70B parameter model and four A100 GPUs. You load it up, and your server throws an Out of Memory error. The math is simple but frustrating: the weights alone require more VRAM than a single card can hold. This is where most engineers hit a wall. They try to squeeze everything onto one device, fail, and then panic about buying more expensive hardware.

But there is a smarter way. It is called Tensor Parallelism. It does not just spread the memory burden; it splits the actual computation across multiple devices so they work in concert. Think of it less like copying a file to two computers and more like slicing a giant pizza so two people can eat it at the same time without fighting over the crust. If you are deploying large language models (LLMs) in production, understanding this technique is no longer optional-it is survival.

The Core Problem: Why Single GPUs Fail at Scale

Modern LLMs are massive. A model like Llama-3-70B requires roughly 140GB of VRAM just to store its weights in FP16 precision. Even the top-tier consumer cards or older enterprise chips cap out at 80GB or less. You physically cannot fit the model into one GPU's memory.

You might ask, "Why not just use Pipeline Parallelism?" That approach stacks layers vertically-GPU 1 handles layers 1-20, GPU 2 handles layers 21-40, and so on. While this saves memory, it creates "pipeline bubbles." While GPU 1 waits for GPU 2 to finish, it sits idle. For latency-sensitive applications like chatbots, that idle time kills performance. Tensor Parallelism solves this by splitting individual layers horizontally. Every GPU works on every layer simultaneously, keeping all hardware busy.

How Tensor Parallelism Actually Works

The magic lies in how matrix multiplication is handled. In a neural network, the heaviest lifting happens in linear layers, which are essentially big matrix multiplications. Tensor Parallelism takes these weight matrices and slices them along the feature dimension.

Imagine a weight matrix $W$ with dimensions $d \times d$. Instead of storing the whole thing on one GPU, we split it into two columns, $W_1$ and $W_2$, each $d \times d/2$. We place $W_1$ on GPU 1 and $W_2$ on GPU 2. When an input vector $x$ comes in, we replicate it on both GPUs. Each GPU computes a partial output: $y_1 = x \cdot W_1$ and $y_2 = x \cdot W_2$. Finally, we sum these partial outputs ($y = y_1 + y_2$) to get the final result.

This process relies on two specific communication patterns defined in the seminal Megatron-LM paper by NVIDIA Research:

• Column Parallelism: Used for QKV projections in attention layers. Inputs are replicated; outputs are gathered.
• Row Parallelism: Used for output projections. Inputs are split; outputs are summed via an All-Reduce operation.

By alternating between these two, the system minimizes data movement. The heavy lifting stays local; only small intermediate results travel between GPUs.

The Communication Bottleneck: NVLink vs. PCIe

Here is the catch: splitting computations means GPUs must talk to each other constantly. Every time a layer finishes its partial calculation, the GPUs need to synchronize. This is where your hardware infrastructure matters more than your code.

If you connect GPUs via standard PCIe 4.0, bandwidth is limited to about 32 GB/s per direction. For a 70B model, the synchronization overhead can consume 15-25% of your total inference time. You effectively pay a "communication tax" for every token generated.

NVIDIA’s NVLink changes the game. It provides up to 600 GB/s bidirectional bandwidth between connected GPUs. Benchmarks from AMD and NVIDIA show that using NVLink reduces communication overhead by nearly 35% compared to PCIe. If you are running tensor parallelism on a cloud instance without high-speed interconnects, you might find that adding a fourth GPU actually makes things slower due to latency.

Implementing TP with Modern Frameworks

You do not need to write custom CUDA kernels to get this working. Major inference engines have abstracted the complexity. As of late 2023 and early 2024, three frameworks dominate the landscape:

1. vLLM: Uses PagedAttention and supports tensor parallelism out of the box. It is highly optimized for throughput and is the go-to for many startups.
2. Hugging Face Text Generation Inference (TGI): Offers robust support for TP via the `--tensor-parallel-size` flag. It integrates seamlessly with the Hugging Face ecosystem.
3. NVIDIA TensorRT-LLM: Provides the highest performance by compiling models into optimized engines. It requires more setup but yields the best latency metrics.

A typical command to launch a 70B model on 4 GPUs using TGI looks like this:

tgi serve --model-id meta-llama/Llama-2-70b-chat-hf --tensor-parallel-size 4

Under the hood, the framework detects the number of available GPUs and automatically shards the model weights. However, you must ensure the number of attention heads in the model is divisible by the tensor parallel degree. If a model has 64 heads and you set TP=5, the split won't be even, leading to errors or inefficient padding.

When Not to Use Tensor Parallelism

Tensor Parallelism is powerful, but it is not a silver bullet. It shines in single-node deployments where GPUs are tightly coupled. Once you cross the node boundary-say, connecting two servers over Ethernet-the latency spikes. AWS documentation notes that EFA networking adds 1.2-2.5ms latency per synchronization point. For a model with 80 layers, that adds up quickly.

For multi-node setups, hybrid approaches are better. You might use Pipeline Parallelism across nodes (to reduce cross-node communication) and Tensor Parallelism within each node (to maximize intra-node speed). This combination is often referred to as 3D Parallelism when combined with Data Parallelism for batch scaling.

Also, consider Mixture-of-Experts (MoE) models like Mixtral. Here, Expert Parallelism is often superior to Tensor Parallelism. Instead of slicing every expert's weights across all GPUs, Expert Parallelism assigns complete experts to specific GPUs. This reduces cross-GPU communication by 40-60% because tokens only need to route to the specific GPU holding the relevant expert, rather than synchronizing every layer.

Troubleshooting Common Pitfalls

If you are debugging TP issues, look for these common culprits:

• All-Reduce Timeouts: If your GPUs hang during generation, check your NCCL timeout settings. Network hiccups can cause one GPU to wait forever for another. Increasing the timeout threshold often fixes transient stalls.
• Uneven Splits: Ensure your model architecture supports the chosen TP degree. If the hidden dimension isn't evenly divisible by the number of GPUs, you will face memory imbalances.
• Memory Fragmentation: Sometimes, even if total memory fits, fragmentation prevents allocation. Using quantization (like INT8 or FP8) alongside TP can alleviate this by reducing the size of intermediate activations passed between GPUs.

The Future: Hybrid and Automated Strategies

The industry is moving toward automation. Manually tuning TP degrees is tedious. New tools are emerging that profile your specific hardware and workload to recommend the optimal parallelism strategy. NVIDIA’s recent updates to TensorRT-LLM include communication compression techniques that reduce the volume of data sent between GPUs by up to 50% using FP8 quantization.

As models grow larger, pure Tensor Parallelism will likely evolve into context-aware hybrid systems. These systems will dynamically adjust how much of the model is sharded versus pipelined based on real-time request patterns. But for now, mastering basic TP is the foundation upon which all advanced LLM deployment rests.