Tensor Parallelism Fundamentals

When an individual Large Language Model exceeds the physical memory capacity of a single GPU, the workload must be partitioned across multiple devices. A 70-billion parameter model loaded in 16-bit precision requires roughly 140 GB of VRAM just for the static weights. This completely saturates the 80 GB capacity of a single enterprise H100 accelerator and far exceeds the 24 GB limit of consumer hardware like the RTX 3090 or RTX 4090. To serve or fine-tune models of this scale, engineering teams distribute the linear algebra operations simultaneously across multiple GPUs. The most prevalent intra-node distribution technique is tensor parallelism. Unlike distributing entire layers sequentially, tensor parallelism slices individual matrix multiplication operations apart, processes them concurrently on separate GPUs, and merges the results over high-speed interconnects.

Mechanics of Matrix Slicing

At the core of transformer architectures are massive matrix multiplications inside the multi-head attention and feed-forward network blocks. Tensor parallelism divides these massive matrices into smaller blocks that can be computed independently. The standard implementation follows a column-parallel and row-parallel sequence.

In a column-parallel operation, the weight matrix is sliced vertically. The input tensor is duplicated across all participating GPUs. Each GPU computes its local matrix multiplication, resulting in a partial output tensor. Next, in a row-parallel operation, the subsequent weight matrix is sliced horizontally. The partial output from the previous step is fed directly into this layer. After the row-parallel computation finishes, the system must perform an all-reduce operation over the network to sum the partial results from every GPU into the final output tensor.

Megatron-style tensor parallelism execution flow through a standard feed-forward layer partitioned across two GPUs.

Memory Overhead and Communication Bottlenecks

Sharding a model with a tensor parallel degree of $t$ divides the static memory footprint of the weights by $t$. For a 70B dense model at FP16 precision, the base weight footprint is roughly 140 GB. Running with $t=4$ on four 80 GB A100 GPUs reduces the per-GPU weight footprint to 35 GB.

The dynamic memory footprint required for the KV cache also scales down proportionally. You calculate the KV cache memory per GPU under tensor parallelism using the following formula:

$Memory_{KV\_per\_GPU} = \frac{2 \times num\_layers \times num\_kv\_heads \times head\_dim \times seq\_length \times bytes\_per\_element \times batch\_size}{t}$

As a concrete example, meta-llama/Meta-Llama-3.1-70B-Instruct has 80 layers, 8 KV heads, and a head dimension of 128. If you run a sequence length of 8192 with a batch size of 16 in FP16 precision (2 bytes per element) across 4 GPUs, you first determine the total elements per sequence:

$2 \times 80 \times 8 \times 128 \times 8192 = 1,342,177,280 \text{ elements}$

At 2 bytes per element, this is roughly 2.5 GB of KV cache per sequence. Multiplying by a batch size of 16 yields a total KV cache footprint of 40 GB. By dividing this footprint across the tensor parallel group ($t=4$), the actual VRAM allocation drops to exactly 10 GB per GPU.

Every transformer layer requires an all-reduce operation across the GPUs in the tensor parallel group. This means the system is highly sensitive to interconnect bandwidth. NVLink 4.0 on Hopper H100 hardware provides 900 GB/s of bidirectional bandwidth, which absorbs this communication overhead easily. Consumer boards relying on PCIe 4.0 x16 are limited to roughly 32 GB/s bidirectional. If you attempt tensor parallelism on dual RTX 4090s, the PCIe bottleneck causes massive latency spikes during the all-reduce phases, severely degrading the Time-Per-Output-Token (TPOT). For Apple Silicon architectures utilizing Unified Memory, the CPU and GPU share the same physical memory pool, bypassing standard PCIe limitations, though memory bandwidth caps (e.g., 800 GB/s on the M2 Ultra) still dictate maximum generation speeds.

Implementing Tensor Parallelism

When deploying infrastructure, modern inference servers like vLLM abstract the matrix partitioning logic. You configure the topology by specifying the tensor parallel size parameter during initialization.

from vllm import LLM, SamplingParams

# Initialize a 70B model across four 80GB VRAM GPUs
# tensor_parallel_size defines the value of 't' in the cluster topology
llm = LLM(
    model="meta-llama/Meta-Llama-3.1-70B-Instruct",
    tensor_parallel_size=4,
    dtype="float16",
    gpu_memory_utilization=0.90
)

prompts = ["Explain the physics of semiconductor scaling."]
sampling_params = SamplingParams(temperature=0.7, max_tokens=256)

outputs = llm.generate(prompts, sampling_params)
for output in outputs:
    print(output.outputs[0].text)

For smaller setups using CLI tools, you can configure the same parameters. If you are operating a dual-GPU 48 GB setup consisting of two RTX 3090s and want to run a 32B model like Qwen/Qwen2.5-32B quantized to 4-bit AWQ, you pass the tensor parallel argument directly to the engine:

vllm serve Qwen/Qwen2.5-32B \
    --tensor-parallel-size 2 \
    --quantization awq \
    --dtype half \
    --max-model-len 8192

Sizing Heuristics and Limits

Tensor parallelism is strictly an intra-node topology. It should never span across network boundaries between physical servers. The latency of Ethernet or InfiniBand networks is too high to support the frequent, per-layer all-reduce operations without causing severe GPU pipeline stalling. A standard sizing heuristic is to keep $t$ less than or equal to the number of GPUs in a single physical node, which is typically capped at 8 on enterprise chassis. If a model requires more than 8 GPUs to fit its weights and KV cache, operators layer tensor parallelism on top of pipeline parallelism.

Mixture of Experts architectures introduce specific memory characteristics. When deploying mistralai/Mixtral-8x7B-Instruct-v0.1, tensor parallelism shards both the base attention layers and the expert feed-forward layers. MoE models possess a massive static parameter footprint (47B parameters) but a small active parameter footprint per token (13B parameters). Spreading this model across two 24 GB VRAM GPUs using $t=2$ effectively halves the static memory pressure. This allows the 47B weights to fit entirely in fast VRAM without relying on slow CPU offloading, keeping generation speeds bounded strictly by the high-speed HBM bandwidth of the GPUs.