Tensor Parallelism — infra.bea.ai

When a model’s individual layers are too large for one GPU’s memory — or when you need lower latency than data parallelism provides — tensor parallelism splits each layer’s weight matrices across GPUs. Every GPU computes a portion of every layer, requiring AllReduce synchronization at each layer boundary.

1. The Core Idea

TP Column Row

Instead of replicating the full weight matrix on every GPU, split it:

Column Parallel (split output dimension): Each GPU holds columns of $W$ , computes a slice of the output.

W = [W₀ | W₁ | W₂ | W₃]    (split across 4 GPUs)
Y₀ = X @ W₀  (GPU 0)        Y = [Y₀, Y₁, Y₂, Y₃]  (concatenate)
Y₁ = X @ W₁  (GPU 1)

Row Parallel (split input dimension): Each GPU holds rows of $W$ , computes with its slice of the input.

W = [W₀]     X = [X₀, X₁, X₂, X₃]    (input already split)
    [W₁]
    [W₂]     Y = X₀@W₀ + X₁@W₁ + X₂@W₂ + X₃@W₃  (AllReduce sum)
    [W₃]

2. TP for Transformer Layers

Attention (Column → Row)

Q, K, V projections:  Column parallel (each GPU gets N_h/P heads)
  → Each GPU computes attention for its own heads (no communication)
Output projection:     Row parallel
  → AllReduce to combine results

Communication: 1 AllReduce per attention layer

FFN (Column → Row)

Gate + Up projection: Column parallel (split d_ff across GPUs)
  → Each GPU computes its slice of the intermediate
Down projection:      Row parallel
  → AllReduce to combine

Communication: 1 AllReduce per FFN layer

Total per transformer block: 2 AllReduces (one for attention, one for FFN).

3. Communication Cost

Per block: 2 AllReduces of size $(B, S, d)$ in BF16:

T_{\text{comm}} = 2 \times \frac{2 \times B \times S \times d \times 2 \times (P-1)}{P \times B_{\text{link}}}

For $B=1, S=2048, d=4096$ , BF16 on 8 GPUs with NVLink (450 GB/s):

T = 2 \times \frac{2 \times 2048 \times 4096 \times 2 \times 7}{8 \times 450 \times 10^9} \approx 0.13 \text{ ms per block}

32 blocks: ~4.2 ms total communication. If computation takes ~50 ms, that’s 8% overhead — acceptable on NVLink.

On InfiniBand (50 GB/s): same calculation gives ~38 ms → 43% overhead. TP across nodes is impractical — always keep TP intra-node.

TP always pairs with Sequence Parallelism to save activation memory. For AllReduce cost analysis, see Communication Cost Analysis.

4. Memory Savings

With TP degree $P$ :

Weight memory per GPU: $\sim N/P$ (each GPU holds 1/P of the weights)
Activation memory: $\sim$ same (input/output activations are full-size)
KV cache: $N_{kv}/P$ per head (if heads are split)

For a 70B model with TP=8: $140 \text{ GB} / 8 = 17.5$ GB weights per GPU. Fits easily on one H100.

5. TP in Practice

# Megatron-LM style TP
from megatron.core import tensor_parallel

# Column parallel linear
self.qkv_proj = tensor_parallel.ColumnParallelLinear(
    hidden_size, 3 * hidden_size, bias=False,
    gather_output=False  # keep split for attention computation
)

# Row parallel linear
self.out_proj = tensor_parallel.RowParallelLinear(
    hidden_size, hidden_size, bias=False,
    input_is_parallel=True  # input is already split across GPUs
)

When to Use TP

Condition	Use TP?
Model fits on one GPU	No — use DP only
Model too large, have NVLink	Yes — TP within the node
Need low per-request latency	Yes — TP reduces per-GPU compute per layer
Only InfiniBand (no NVLink)	No — TP communication is too expensive

Self-Check

Explain column parallel vs row parallel and which layers use which
Calculate the number of AllReduces per transformer block with TP
Calculate TP communication time on NVLink vs InfiniBand
Explain why TP must stay intra-node (NVLink requirement)
Explain memory savings from TP and what is NOT reduced (activations)