Tensor Parallelism

Table of Contents

1. Learning Objectives
2. Notation
3. Core Intuition
4. Column-Parallel Linear Layer
5. Row-Parallel Linear Layer
6. Megatron-LM Attention Splitting
7. Megatron-LM FFN Splitting
8. Communication Analysis
9. Scaling Efficiency
10. Common Pitfalls
11. Summary
12. Exercises

Learning Objectives

1. Derive column-parallel and row-parallel matrix multiplication.
2. Explain how Megatron-LM splits attention heads across GPUs.
3. Compute the communication volume per layer.
4. Analyze why tensor parallelism is limited to within a node.
5. Compare TP scaling efficiency at 2, 4, 8 GPUs.

Notation

• $N$ — tensor parallel degree (number of GPUs)
• $\mathbf{A}_i$ — partition $i$ of matrix $\mathbf{A}$
• AR — all-reduce operation
• AG — all-gather operation

Core Intuition

When a single layer is too large for one GPU (or we want faster per-step latency), we split the weight matrices across GPUs. Each GPU computes part of the matrix multiplication, then they communicate to combine results. The key insight: by carefully choosing how to split (column-wise or row-wise), we can minimize communication to just 2 all-reduces per transformer layer.

Column-Parallel Linear Layer

Split weight $\mathbf{A} \in \mathbb{R}^{k \times n}$ along columns: $\mathbf{A} = [\mathbf{A}_1, \mathbf{A}_2, \ldots, \mathbf{A}_N]$.

Each GPU $i$ computes: $\mathbf{y}_i = \mathbf{x}\mathbf{A}_i \in \mathbb{R}^{b \times n/N}$.

Input: Replicated across all GPUs (each has full $\mathbf{x}$). Output: Partitioned (each GPU has columns $n/N$).

Communication: None in forward (if input is already replicated). All-gather needed if output must be full.

Row-Parallel Linear Layer

Split weight $\mathbf{B} \in \mathbb{R}^{n \times k}$ along rows: $\mathbf{B} = [\mathbf{B}_1; \mathbf{B}_2; \ldots; \mathbf{B}_N]$.

Input must be partitioned: $\mathbf{x} = [\mathbf{x}_1, \mathbf{x}_2, \ldots, \mathbf{x}_N]$.

Each GPU $i$ computes: $\mathbf{z}_i = \mathbf{x}_i\mathbf{B}_i \in \mathbb{R}^{b \times k}$.

Full output: $\mathbf{z} = \sum_i \mathbf{z}_i$ → all-reduce.

Communication: One all-reduce to sum partial results.

Megatron-LM Attention Splitting

Natural partition: Split attention heads across GPUs. Each GPU handles $H/N$ heads.

Per-GPU computation:

• Q, K, V projections for $H/N$ heads: $\mathbf{Q}_i = \mathbf{X}\mathbf{W}^Q_i$ (column-parallel).
• Attention computation: local (no cross-GPU communication).
• Output projection: row-parallel (all-reduce to combine).

Result: Only 1 all-reduce needed for the attention block.

Megatron-LM FFN Splitting

For FFN: $\mathbf{Y} = \text{GeLU}(\mathbf{XA})\mathbf{B}$.

Split:

• First linear $\mathbf{A}$: column-parallel (split output features). GeLU applied locally.
• Second linear $\mathbf{B}$: row-parallel (takes partitioned input, all-reduce output).

Result: 1 all-reduce for the FFN block.

Total per transformer layer: 2 all-reduces in forward, 2 in backward = 4 all-reduces.

Communication Analysis

All-reduce volume per operation: $2(N-1)/N \times M$ bytes, where $M$ is the tensor size.

For one transformer layer with sequence length $T$, hidden dim $d$:

• Attention all-reduce: $2(N-1)/N \times T \times d \times 2$ bytes (FP16).
• FFN all-reduce: $2(N-1)/N \times T \times d \times 2$ bytes.
• Total per layer: $\approx 8Td$ bytes (for large $N$).

Bandwidth requirement: For latency $\Delta t$ per layer: need bandwidth $\geq 8Td/\Delta t$.

Scaling Efficiency

Computation scales linearly: Each GPU does $1/N$ of the work.

Communication overhead: All-reduce time depends on interconnect bandwidth.

Efficiency: $\eta = \frac{\text{compute time}}{\text{compute time} + \text{communication time}}$.

For NVLink (900 GB/s bidirectional within a node):

• TP=2: $\eta \approx 95\%$
• TP=4: $\eta \approx 90\%$
• TP=8: $\eta \approx 80\%$

For cross-node (InfiniBand, 200 GB/s):

• TP=2: $\eta \approx 70\%$ — too slow!

Rule: Tensor parallelism only within a single node (connected by NVLink).

Common Pitfalls

Pitfall 1. Using tensor parallelism across nodes. Cross-node bandwidth is 4-5x lower than NVLink; communication dominates compute.

Pitfall 2. Non-uniform splitting when $H$ is not divisible by $N$. Requires padding or uneven distribution, causing load imbalance.

Pitfall 3. Forgetting the dropout synchronization. Random dropout masks must be synchronized across TP ranks to ensure correctness.

Summary

• Column-parallel: Split output features; input replicated, output partitioned.
• Row-parallel: Split input features; input partitioned, output all-reduced.
• Megatron-LM: 2 all-reduces per layer (1 attention + 1 FFN).
• Best within a node: NVLink provides sufficient bandwidth; cross-node is too slow.
• Typical TP degree: 2, 4, or 8 (matching GPUs per node).

Exercises

Exercise 1. For TP=4, $d=4096$, $T=2048$: compute the all-reduce volume per transformer layer.

Exercise 2. Derive the efficiency $\eta$ for TP=8 with NVLink bandwidth 450 GB/s and per-layer compute time 2ms.

Exercise 3. Show that Megatron-LM's attention splitting produces mathematically identical results to single-GPU computation.

Exercise 4. Compute the maximum TP degree for a model with $H=32$ attention heads.

Exercise 5. Compare the end-to-end training throughput for TP=4+DP=8 vs TP=8+DP=4 on 32 GPUs.