← Back to DevBytes
Developer Tutorial August 18, 2026 9 min read

Understanding Tensor Parallelism for Multi-GPU LLM Serving

Understanding Tensor Parallelism for Multi-GPU LLM Serving

As large language models (LLMs) grow beyond the memory capacity of a single GPU, developers face a critical challenge: how do you serve a 70B or 175B parameter model when no single accelerator has enough VRAM to hold it? The answer lies in distributed inference strategies, and among the most effective is tensor parallelism (TP). This tutorial explains what tensor parallelism is, why it matters for production LLM serving, how to implement it in practice, and the best practices you should follow.

What Is Tensor Parallelism?

Tensor parallelism is a model-parallel distributed computing technique where individual tensors (the weight matrices of a neural network) are split across multiple GPUs. Unlike data parallelism, where each GPU holds a full copy of the model and processes different batches, tensor parallelism divides the model's parameters themselves so that each GPU computes a fragment of every operation.

In a transformer-based LLM, the two main components that benefit from tensor parallelism are:

The key insight is that these splits require communication between GPUs only at specific synchronization points — typically after the attention concatenation and after the FFN's row-parallel reduction. This keeps the communication overhead manageable compared to naive approaches.

Why Tensor Parallelism Matters for LLM Serving

Serving LLMs in production introduces several constraints that tensor parallelism directly addresses:

It is worth noting that tensor parallelism is not free. The all-reduce operations required after each transformer block introduce communication overhead, which is why TP is typically confined to a single node with high-bandwidth interconnects like NVLink rather than spread across slower network links.

How Tensor Parallelism Works Under the Hood

To understand TP concretely, consider a linear layer Y = XW where X is the input activation and W is the weight matrix. There are two primary partitioning strategies:

Column parallelism: The weight matrix W is split along its output dimension. Each GPU i holds W_i and computes Y_i = XW_i. Since all GPUs share the same input X, no communication is needed before the computation. The outputs Y_i are different slices of the final result.

Row parallelism: The weight matrix is split along its input dimension. Each GPU holds W_i and receives a slice X_i of the input, computing a partial sum Y_i = X_i W_i. An all-reduce operation then sums the partial results: Y = sum(Y_i).

In a transformer FFN block, these two strategies are combined. The first linear layer is column-parallel (no input communication needed), and the second is row-parallel (output requires all-reduce). This pairing minimizes synchronization points to one per FFN block.

Practical Implementation with vLLM

The most accessible way to apply tensor parallelism in production is through a serving framework like vLLM, which implements TP natively along with PagedAttention for efficient KV cache management. Below is a minimal example of serving a model with tensor parallelism across 4 GPUs.

# Install vLLM (requires CUDA-capable machine with multiple GPUs)
# pip install vllm

from vllm import LLM, SamplingParams

# Initialize the LLM with tensor_parallel_size set to the number of GPUs
llm = LLM(
    model="meta-llama/Llama-2-70b-hf",
    tensor_parallel_size=4,
    dtype="float16",
    enforce_eager=False,
)

# Define sampling parameters
sampling_params = SamplingParams(
    temperature=0.7,
    top_p=0.9,
    max_tokens=256,
)

# Generate completions
prompts = [
    "Explain the difference between TCP and UDP.",
    "Write a Python function to compute the Fibonacci sequence.",
]
outputs = llm.generate(prompts, sampling_params)

for output in outputs:
    prompt = output.prompt
    generated = output.outputs[0].text
    print(f"Prompt: {prompt}")
    print(f"Generated: {generated}\n")

To launch the same model as an OpenAI-compatible API server with tensor parallelism, use the command-line interface:

# Launch a vLLM server with 4-way tensor parallelism
python -m vllm.entrypoints.openai.api_server \
    --model meta-llama/Llama-2-70b-hf \
    --tensor-parallel-size 4 \
    --dtype float16 \
    --port 8000

Once running, you can send requests using the standard OpenAI client library:

from openai import OpenAI

client = OpenAI(base_url="http://localhost:8000/v1", api_key="dummy")

response = client.chat.completions.create(
    model="meta-llama/Llama-2-70b-hf",
    messages=[
        {"role": "system", "content": "You are a helpful assistant."},
        {"role": "user", "content": "What is tensor parallelism?"},
    ],
    max_tokens=200,
)

print(response.choices[0].message.content)

Implementing Tensor Parallelism with Hugging Face Accelerate

For more control, or when you need to integrate TP into a custom inference pipeline, the Hugging Face accelerate library provides lower-level primitives. The following example demonstrates a manual column-parallel linear layer using PyTorch distributed primitives.

import torch
import torch.nn as nn
import torch.distributed as dist


class ColumnParallelLinear(nn.Module):
    """A linear layer whose weight matrix is split column-wise across GPUs."""

    def __init__(self, in_features, out_features, process_group):
        super().__init__()
        self.process_group = process_group
        self.world_size = dist.get_world_size(process_group)
        self.rank = dist.get_rank(process_group)

        assert out_features % self.world_size == 0, (
            "out_features must be divisible by the number of GPUs"
        )
        self.out_features_per_partition = out_features // self.world_size
        self.in_features = in_features

        # Each GPU only stores its shard of the weight matrix
        self.weight = nn.Parameter(
            torch.empty(self.out_features_per_partition, in_features)
        )
        nn.init.kaiming_uniform_(self.weight, a=5 ** 0.5)
        self.bias = nn.Parameter(torch.zeros(self.out_features_per_partition))

    def forward(self, x):
        # Input x is replicated across all GPUs (no communication needed)
        # Each GPU computes its slice of the output
        output_local = torch.nn.functional.linear(x, self.weight, self.bias)
        return output_local


class RowParallelLinear(nn.Module):
    """A linear layer whose weight matrix is split row-wise across GPUs."""

    def __init__(self, in_features, out_features, process_group):
        super().__init__()
        self.process_group = process_group
        self.world_size = dist.get_world_size(process_group)
        self.rank = dist.get_rank(process_group)

        assert in_features % self.world_size == 0, (
            "in_features must be divisible by the number of GPUs"
        )
        self.in_features_per_partition = in_features // self.world_size
        self.out_features = out_features

        self.weight = nn.Parameter(
            torch.empty(out_features, self.in_features_per_partition)
        )
        nn.init.kaiming_uniform_(self.weight, a=5 ** 0.5)
        self.bias = nn.Parameter(torch.zeros(out_features))

    def forward(self, x):
        # x is already sharded along the feature dimension from the previous
        # column-parallel layer, so each GPU takes its slice
        x_local = x[..., self.rank * self.in_features_per_partition:
                       (self.rank + 1) * self.in_features_per_partition]

        # Compute partial output
        output_local = torch.nn.functional.linear(x_local, self.weight)

        # All-reduce to combine partial results from all GPUs
        dist.all_reduce(output_local, op=dist.ReduceOp.SUM, group=self.process_group)
        output_local = output_local + self.bias
        return output_local


def init_distributed():
    """Initialize the default process group for tensor parallelism."""
    dist.init_process_group(backend="nccl")
    torch.cuda.set_device(dist.get_rank())


if __name__ == "__main__":
    init_distributed()
    group = dist.group.WORLD

    # Simulate an FFN block: column-parallel followed by row-parallel
    hidden_size = 4096
    intermediate_size = 11008

    ffn_up = ColumnParallelLinear(hidden_size, intermediate_size, group).cuda()
    ffn_down = RowParallelLinear(intermediate_size, hidden_size, group).cuda()

    # Dummy input replicated across GPUs
    batch_size = 4
    x = torch.randn(batch_size, hidden_size, device="cuda")

    # Forward pass through the parallel FFN
    h = torch.nn.functional.gelu(ffn_up(x))
    y = ffn_down(h)

    if dist.get_rank() == 0:
        print(f"Output shape: {y.shape}")

    dist.destroy_process_group()

To run this script across multiple GPUs, use torchrun:

torchrun --nproc_per_node=4 tensor_parallel_ffn.py

Combining Tensor Parallelism with Pipeline Parallelism

For very large models that span multiple nodes, tensor parallelism is often combined with pipeline parallelism (PP). In this hybrid scheme, TP is applied within a node (where NVLink provides high bandwidth), while PP splits layers across nodes. The general rule is:

world_size = tensor_parallel_size * pipeline_parallel_size * data_parallel_size

With vLLM, you can configure both:

python -m vllm.entrypoints.openai.api_server \
    --model meta-llama/Llama-2-70b-hf \
    --tensor-parallel-size 4 \
    --pipeline-parallel-size 2 \
    --dtype float16

Best Practices

Common Pitfalls

Conclusion

Tensor parallelism is a foundational technique for serving large language models that exceed the capacity of a single GPU. By partitioning weight matrices across accelerators and synchronizing only at carefully chosen points, TP enables both the memory footprint and the compute throughput needed for production-scale inference. When applied within a single NVLink-connected node, paired with an optimized serving framework like vLLM or TensorRT-LLM, and combined with pipeline or data parallelism for multi-node scaling, tensor parallelism forms the backbone of virtually every modern multi-GPU LLM deployment. Understanding its mechanics — from column- and row-parallel linear layers to all-reduce synchronization — equips you to diagnose bottlenecks, choose appropriate parallelism strategies, and ultimately deliver low-latency, high-throughput LLM services to your users.

— Ad —

Google AdSense will appear here after approval

← Back to all articles