Tensor Parallelism

9 min read

Learn

Quick Reading

Estimated 9 mins

Prereq

Intermediate

Basic ML concepts helpful

Interactive

Static Playbook

Static guide & reference tables

Splitting a single matrix multiply across devices

Tensor parallelism addresses a different bottleneck than pipeline parallelism: what happens when an individual weight matrix — say, the feed-forward layer's weight matrix in a transformer, which can be tens of thousands by tens of thousands of elements in large models — is itself too large to comfortably fit or compute on a single device. Rather than keeping each layer whole and splitting across layers (pipeline parallelism's approach), tensor parallelism splits individual matrices and distributes the computation of a single operation across multiple devices simultaneously.

Column-parallel and row-parallel splits

The Megatron-LM approach to tensor parallelism (the reference design most implementations follow) splits transformer feed-forward blocks in a specific, carefully chosen way. The first linear layer's weight matrix is split by columns across devices ('column-parallel'): each device computes its own partial output using its column-slice of the weights, and because of how the subsequent nonlinearity (e.g., GELU) can be applied independently to each column-slice's output, no communication is needed yet. The second linear layer's weight matrix is split by rows ('row-parallel'): each device computes a partial sum using its row-slice, and these partial sums must then be summed across devices with an all-reduce to produce the final correct output. This column-then-row pairing is deliberately chosen to minimize the number of communication steps — just one all-reduce per feed-forward block, not one per matrix.

tensor_parallel_ffn.py

python

# Simplified illustration of Megatron-style tensor parallelism
# for a 2-layer feed-forward block, split across 2 devices.
import torch
import torch.distributed as dist

def column_parallel_linear(x, W1_shard, rank, world_size):
    # Each device holds a column-slice of W1: shape [d_in, d_hidden / world_size]
    return x @ W1_shard   # [batch, d_hidden / world_size] -- no communication needed yet

def row_parallel_linear(h_shard, W2_shard):
    # Each device holds a row-slice of W2: shape [d_hidden / world_size, d_out]
    partial_out = h_shard @ W2_shard   # [batch, d_out], a partial sum
    dist.all_reduce(partial_out, op=dist.ReduceOp.SUM)  # sum across devices
    return partial_out

def tensor_parallel_ffn(x, W1_shard, W2_shard, rank, world_size):
    h_shard = torch.nn.functional.gelu(
        column_parallel_linear(x, W1_shard, rank, world_size)
    )
    return row_parallel_linear(h_shard, W2_shard)   # one all-reduce total per FFN block

Tensor parallelism in attention layers

Attention layers are split along a natural seam: individual attention heads are independent of each other, so different heads (or groups of heads) can be assigned to different devices with no communication needed during the per-head attention computation itself. Only the final output projection, which combines results across all heads, requires an all-reduce — again, just once per attention block. This head-wise split is why tensor-parallel degree is often chosen to evenly divide the number of attention heads.

Tensor parallelism vs. pipeline parallelism

Property	Tensor parallelism	Pipeline parallelism
What's split	Individual weight matrices within a layer	Whole layers, across devices
Communication frequency	Every layer, every forward/backward pass	Only at stage boundaries
Interconnect requirement	Very high bandwidth, low latency (e.g. NVLink within a node)	Tolerates higher latency (e.g. across nodes)
Idle time pattern	None inherently, but blocked on communication	Pipeline bubbles unless micro-batched carefully

Tensor parallelism doesn't scale well across slow networks

Because tensor parallelism requires an all-reduce on nearly every layer, its communication volume is far higher-frequency than pipeline or data parallelism. Splitting tensor-parallel groups across nodes connected by standard datacenter networking (rather than intra-node interconnects like NVLink) typically causes severe slowdowns — this is why tensor-parallel degree is usually capped at the number of GPUs available within a single node.

Correctness, not just speed, is the hard part

The specific column/row split pattern isn't an arbitrary engineering choice — it's the minimal communication pattern that still produces a mathematically identical result to running the full, unsplit computation on one device. Getting this wrong (e.g., splitting in a way that requires communication mid-nonlinearity) would either break correctness or force far more frequent synchronization than necessary.

What's next?

Tensor and pipeline parallelism address compute distribution, but memory redundancy from ordinary data parallelism — every replica storing its own full copy of optimizer states and gradients — is a separate inefficiency entirely, and that's precisely the problem the ZeRO optimizer family was designed to eliminate.

I build these systems professionally.

Whether it's a RAG pipeline, analytics migration, or AI workflow — let's talk.

Work With Me →