A 70B parameter model in FP16 needs 140 GB just for the weights. An H100 has 80 GB of HBM. The model simply does not fit on one GPU. You need to split it, and tensor parallelism (TP) is the most common way to do it for inference.
I covered the theory of model parallelism earlier in this journey. Today I want to get concrete: simulate the actual matrix splits, the partial computations on each GPU, and the all-reduce communication that stitches the results back together. No multi-GPU hardware required. Just NumPy, clear thinking, and the Megatron-LM paper's partitioning scheme.
The core idea
A linear layer computes Y = XW + b, where X is the input activation and W is the weight matrix. Tensor parallelism splits W across N GPUs. There are two ways to split:
- Column parallel: split W along its columns. Each GPU holds a vertical slice of W and computes a slice of the output. The partial outputs are concatenated.
- Row parallel: split W along its rows. Each GPU holds a horizontal slice of W and computes a partial sum of the output. The partial sums are all-reduced (summed).
The Megatron-LM insight is to pair these: use column parallel for the first linear layer in a transformer block (e.g., the up-projection in the MLP), and row parallel for the second (the down-projection). This way, you only need one all-reduce per transformer block instead of two, because the column-parallel output is already distributed in the right layout for the row-parallel input.
Simulating column parallel linear
Let's start with the simpler case. We have a weight matrix W of shape (in_features, out_features) and we want to split it across N GPUs along the output dimension:
import numpy as np
def column_parallel_linear(x, W, bias, num_gpus):
"""Simulate a column-parallel linear layer.
W is split along columns (axis=1). Each GPU computes its slice
of the output independently. No communication needed until
the results are used.
Args:
x: input activation, shape (batch, in_features)
W: weight matrix, shape (in_features, out_features)
bias: bias vector, shape (out_features,)
num_gpus: number of GPUs to split across
Returns:
List of partial outputs, one per GPU
"""
in_features, out_features = W.shape
assert out_features % num_gpus == 0
slice_size = out_features // num_gpus
gpu_outputs = []
for i in range(num_gpus):
# Each GPU gets columns [i*slice : (i+1)*slice]
W_slice = W[:, i * slice_size : (i + 1) * slice_size]
b_slice = bias[i * slice_size : (i + 1) * slice_size]
# Local matmul: no communication needed
y_slice = x @ W_slice + b_slice
gpu_outputs.append(y_slice)
return gpu_outputs
# Verify: concatenating all slices equals the full computation
np.random.seed(42)
batch, in_f, out_f = 4, 512, 2048
x = np.random.randn(batch, in_f).astype(np.float32)
W = np.random.randn(in_f, out_f).astype(np.float32)
b = np.random.randn(out_f).astype(np.float32)
y_full = x @ W + b
gpu_outputs = column_parallel_linear(x, W, b, num_gpus=4)
y_concat = np.concatenate(gpu_outputs, axis=1)
print(f"Max error: {np.max(np.abs(y_full - y_concat)):.2e}")
# Max error: 0.00e+00Column parallel is clean: each GPU computes independently, and the outputs are simply concatenated. No communication during the forward pass. The activation function (GeLU, SiLU) can be applied locally on each GPU because it operates element-wise on each slice.
Simulating row parallel linear
Row parallel splits W along the input dimension. Each GPU holds rows of W and receives a slice of the input. The catch is that each GPU computes a partial sum of the output, and we need an all-reduce to get the final result:
def row_parallel_linear(x_slices, W, bias, num_gpus):
"""Simulate a row-parallel linear layer.
W is split along rows (axis=0). Each GPU gets a slice of the input
and a horizontal slice of W. The partial outputs are summed via
all-reduce.
Args:
x_slices: list of input slices, one per GPU.
Each has shape (batch, in_features // num_gpus)
W: full weight matrix, shape (in_features, out_features)
bias: bias vector, shape (out_features,)
num_gpus: number of GPUs
Returns:
Final output after all-reduce, shape (batch, out_features)
"""
in_features, out_features = W.shape
assert in_features % num_gpus == 0
slice_size = in_features // num_gpus
partial_sums = []
for i in range(num_gpus):
W_slice = W[i * slice_size : (i + 1) * slice_size, :]
partial = x_slices[i] @ W_slice
partial_sums.append(partial)
# All-reduce: sum partial results across GPUs
# In real hardware, this is an NCCL all-reduce over NVLink
y = sum(partial_sums) + bias
return yThe all-reduce is the communication cost. On NVLink-connected GPUs (like 8x H100 in a DGX node), all-reduce for a typical activation tensor takes on the order of tens of microseconds. On PCIe, it is significantly slower. This is why tensor parallelism is practical within a node (NVLink) but not across nodes (network), where pipeline parallelism is preferred.
Pairing them: the Megatron-LM MLP
A standard transformer MLP has two linear layers with an activation function in between:
y = W_down @ activation(W_up @ x + b_up) + b_downMegatron-LM makes W_up column-parallel and W_down row-parallel. Here is why this is clever:
def megatron_mlp_forward(x, W_up, b_up, W_down, b_down, num_gpus):
"""Simulate a Megatron-LM style tensor-parallel MLP.
W_up: column-parallel (split along output dim)
W_down: row-parallel (split along input dim)
The column-parallel output is already partitioned correctly
to serve as input to the row-parallel layer. No reshape or
communication between the two layers.
"""
# Step 1: Column-parallel up-projection
# Each GPU computes a slice of the hidden activations
gpu_hidden = column_parallel_linear(x, W_up, b_up, num_gpus)
# Step 2: Apply activation locally on each GPU
gpu_activated = [np.maximum(0, h) for h in gpu_hidden] # ReLU for simplicity
# Step 3: Row-parallel down-projection
# Each GPU's activated slice is the correct input slice for row-parallel
y = row_parallel_linear(gpu_activated, W_down, b_down, num_gpus)
# Only ONE all-reduce happened (inside row_parallel_linear)
return yIf both layers were independently parallelized, you would need two all-reduces per MLP block (one after each layer). By pairing column-parallel with row-parallel, the output partitioning of the first layer matches the input partitioning of the second, eliminating one all-reduce. Over 80 transformer layers, this halves the total communication volume.
Attention heads: naturally parallel
Tensor parallelism applies to attention even more naturally than to the MLP. Multi-head attention already splits computation across heads. With TP degree N, each GPU handles (num_heads / N) attention heads. The Q, K, V projection is column-parallel (each GPU computes its heads' projections), attention is computed locally per head, and the output projection is row-parallel.
def tensor_parallel_attention(x, W_qkv, W_out, num_heads, num_gpus):
"""Sketch of tensor-parallel multi-head attention.
Each GPU owns (num_heads // num_gpus) attention heads.
"""
heads_per_gpu = num_heads // num_gpus
for gpu_id in range(num_gpus):
# Column-parallel QKV projection: each GPU gets its heads
# No communication needed
q, k, v = project_qkv(x, W_qkv, gpu_id, heads_per_gpu)
# Local attention computation per head
# No communication needed (each head is independent)
attn_out = scaled_dot_product_attention(q, k, v)
# Row-parallel output projection: all-reduce to combine
# This is the only communication point
output = row_parallel_output_projection(attn_outputs, W_out)
return outputThe constraint is that num_heads must be divisible by the TP degree. A model with 32 attention heads can use TP=1, 2, 4, 8, 16, or 32. This is why model architectures choose head counts that are powers of two or have many factors.
The communication cost model
Each transformer layer requires two all-reduce operations: one in the MLP and one in the attention output projection. An all-reduce of a tensor with B elements across N GPUs transfers 2 * B * (N-1) / N bytes (using the ring all-reduce algorithm). For a 7B model with hidden dimension 4096 in FP16, each all-reduce moves:
# Activation size: batch_size * seq_len * hidden_dim * 2 bytes
# For batch=1, seq_len=1 (decode), hidden=4096, FP16:
allreduce_bytes = 1 * 1 * 4096 * 2 # 8 KB per all-reduce
# Two all-reduces per layer, 32 layers:
total_bytes = 8192 * 2 * 32 # 512 KB total
# NVLink bandwidth: ~900 GB/s per direction on H100
# Time: 512 KB / 900 GB/s = ~0.5 microseconds
# But latency dominates at small sizes: ~5-10 microseconds per all-reduceFor decode (batch size 1, sequence length 1), the all-reduce data is tiny but the latency of initiating the collective operation dominates. With 32 layers and 2 all-reduces each, that is 64 all-reduce calls at 5-10 microseconds each, adding 320-640 microseconds of communication overhead per token. For a model that generates a token in 10ms on a single GPU, this is a 3-6% overhead, which is acceptable.
For prefill with large batches and long sequences, the data volume grows but so does the compute, keeping the communication-to-computation ratio favorable.
Tensor parallelism splits the work, but communication is the tax. The art is minimizing the number of synchronization points (Megatron pairing) and keeping each one small (all-reduce only activations, never weights).
What I took away
- TP is about splitting matrices, not layers. Unlike pipeline parallelism, every GPU participates in every layer. This keeps latency low at the cost of communication.
- Column + row pairing halves communication. The Megatron-LM insight is simple but impactful: match the output partition of one layer to the input partition of the next.
- Attention heads are free parallelism. Heads are independent, so splitting them across GPUs requires no algorithmic changes.
- NVLink makes TP practical. The all-reduce latency at small tensor sizes (decode) is the bottleneck. High-bandwidth interconnects like NVLink keep this under control.
Next up: benchmarking the ops:byte ratio in practice, where I measure actual arithmetic intensity on real workloads and see how closely reality matches the roofline theory from day 6.