Tensor Parallelism Explained: Splitting Attention and FFN Weights Across GPUs

In 2019, NVIDIA researchers published Megatron-LM (Shoeybi et al.), introducing Tensor Parallelism (TP)—the gold standard technique for sharding individual Large Language Model layers across multiple GPUs within a high-speed NVLink domain.

Rather than placing entire layers on different GPUs (which introduces pipeline latency bubbles), Tensor Parallelism splits the internal weight matrices of every Multi-Head Attention and Feed-Forward Network block across $N$ GPUs, computing matrix multiplications in parallel with only two All-Reduce communication operations per transformer layer.

---

1. Matrix Partitioning: Column-Parallel vs Row-Parallel Linear Layers

The mathematical genius of Megatron-LM is pairing a Column-Parallel linear layer directly with a Row-Parallel linear layer. This eliminates intermediate communication steps and requires an All-Reduce only at the very end of the sub-block.

A. Column-Parallel Linear Layer

In a column-parallel layer, the weight matrix $\mathbf{W}\in {\mathbb{R}}^{D_{\text{in}}\times D_{\text{out}}}$ is split along its columns across $N$ GPUs:

$\mathbf{W}=\left(\begin{array}{cccc}{\mathbf{W}}_1& {\mathbf{W}}_2& \dots & {\mathbf{W}}_N\end{array}\right),{\mathbf{W}}_i\in {\mathbb{R}}^{D_{\text{in}}\times \frac{D_{\text{out}}}{N}}$

Each GPU $i$ holds a duplicate of input $\mathbf{X}\in {\mathbb{R}}^{B\times D_{\text{in}}}$ and computes:

${\mathbf{Y}}_i=\mathbf{X}{\mathbf{W}}_i\in {\mathbb{R}}^{B\times \frac{D_{\text{out}}}{N}}$

Zero inter-GPU communication is required. The output is naturally split along its feature columns.

B. Row-Parallel Linear Layer

In a row-parallel layer, the weight matrix ${\mathbf{W}}^{\prime }\in {\mathbb{R}}^{D_{\text{out}}\times D_{\text{final}}}$ is split along its rows:

${\mathbf{W}}^{\prime }=\left(\begin{array}{c}{\mathbf{W}}_1^{\prime }\\ {\mathbf{W}}_2^{\prime }\\ ⋮\\ {\mathbf{W}}_N^{\prime }\end{array}\right),{\mathbf{W}}_i^{\prime }\in {\mathbb{R}}^{\frac{D_{\text{out}}}{N}\times D_{\text{final}}}$

Each GPU takes its local input slice ${\mathbf{Y}}_i$ from the preceding column-parallel layer and computes:

${\mathbf{Z}}_i={\mathbf{Y}}_i{\mathbf{W}}_i^{\prime }=(\mathbf{X}{\mathbf{W}}_i){\mathbf{W}}_i^{\prime }\in {\mathbb{R}}^{B\times D_{\text{final}}}$

To reconstruct the final mathematical product $\mathbf{Z}=\mathbf{XW}{\mathbf{W}}^{\prime }$, the GPUs execute an All-Reduce (Sum) collective operation:

$\mathbf{Z}={\sum }_{i=1}^N{\mathbf{Z}}_i={\sum }_{i=1}^N(\mathbf{X}{\mathbf{W}}_i{\mathbf{W}}_i^{\prime })=\mathbf{X}\left({\sum }_{i=1}^N{\mathbf{W}}_i{\mathbf{W}}_i^{\prime }\right)=\mathbf{X}(\mathbf{W}{\mathbf{W}}^{\prime })$

---

2. Applying Tensor Parallelism to Transformer Blocks

A. The Attention Block

1. Column-Parallel Projections: ${\mathbf{W}}_q,{\mathbf{W}}_k,{\mathbf{W}}_v$ are column-partitioned.
   • If the model has 64 attention heads and $TP=4$, each GPU computes attention for 16 heads independently in local VRAM.
2. Local Attention Computation: Each GPU computes scaled dot-product attention on its assigned subset of heads using local KV cache blocks.
3. Row-Parallel Output Projection: ${\mathbf{W}}_o$ is row-partitioned.
4. Synchronization: An All-Reduce sums the partial output vectors across the 4 GPUs.

B. The Feed-Forward (MLP / SwiGLU) Block

1. Column-Parallel Gate & Up Projections: ${\mathbf{W}}_{\text{gate}}$ and ${\mathbf{W}}_{\text{up}}$ are column-partitioned.
2. Local Activation Function: Elementwise multiplication: $\text{Swish}(\mathbf{X}{\mathbf{W}}_{\text{gate},i})\odot (\mathbf{X}{\mathbf{W}}_{\text{up},i})$ happens completely locally.
3. Row-Parallel Down Projection: ${\mathbf{W}}_{\text{down}}$ is row-partitioned.
4. Synchronization: A second All-Reduce sums the FFN output vectors.

---

3. Communication Cost per Transformer Layer

For a model with hidden dimension $d_{\text{model}}$, batch size $B$, and sequence length $S$:

$\text{Bytes Transferred per Layer}=2\times \text{All-Reduce}=2\times \left(2\times \frac{N-1}{N}\times B\times S\times d_{\text{model}}\times b_{\text{dtype}}\right)$

For $TP=8$ with FP16 ($b_{\text{dtype}}=2$) and $d_{\text{model}}=8192$:

• Each All-Reduce transfers $\approx 31.5\text{ KB}$ per token across NVLink.
• Over NVLink 4 ($900\text{ GB/s}$), this takes less than 5 microseconds per layer, introducing practically zero latency overhead!

---

4. Production Failure Modes: Attention Head Divisibility

Failure Mode: Crash on Startup with Assertion Error on Head Counts

• Symptom: RuntimeError: Number of attention heads (64) must be divisible by tensor_parallel_size (5).
• Root Cause: Tensor Parallelism shards attention heads evenly across GPUs. If $n_{\text{heads}}$ or $n_{\text{kv\_heads}}$ is not evenly divisible by $TP$, the model cannot partition its heads.
• Resolution: $TP$ must always be a power of two ($1,2,4,8$) that evenly divides both $n_{\text{heads}}$ and $n_{\text{kv\_heads}}$. For Llama 3 70B ($n_{\text{kv\_heads}}=8$), the valid $TP$ values are $1,2,4,$ and $8$.

---

5. Summary & Key Takeaways

1. Tensor Parallelism Shards Layers Internally: By partitioning matrices along column and row dimensions, TP distributes memory and compute across GPUs within the same node.
2. Only 2 All-Reduces per Layer: The Column-to-Row pairing eliminates intermediate communications, requiring synchronization only after the Attention and FFN blocks.
3. NVLink is Mandatory: Due to 160 All-Reduce operations per generated token, TP requires the sub-microsecond latency of NVLink interconnects.
4. Even Head Partitioning: The number of attention heads and KV heads must always be divisible by the Tensor Parallel size.