tensor parallelism attention
**Tensor Parallelism for Attention and MLPs** is the **technique of partitioning individual transformer layer computations (attention heads and MLP matrices) across multiple GPUs so that each GPU computes a portion of every layer** — enabling models too large for a single GPU's memory to be trained and served with minimal communication overhead, as pioneered by Megatron-LM for large-scale transformer training.
**Why Tensor Parallelism?**
For models with billions of parameters, a single transformer layer may require more memory than one GPU has. Unlike data parallelism (which replicates the model) or pipeline parallelism (which assigns different layers to different GPUs), tensor parallelism splits individual matrix multiplications across GPUs.
**MLP Tensor Parallelism (Megatron-LM)**
A transformer MLP block: Y = GeLU(XA) · B
Split into column-parallel and row-parallel:
```svg
```
- **Column Parallel**: Matrix A is split column-wise → each GPU gets A₁, A₂. Input X is replicated. Each GPU computes XA_i independently. GeLU is applied locally (no communication needed because column split preserves independent neurons).
- **Row Parallel**: Matrix B is split row-wise → each GPU gets B₁, B₂. Each GPU computes partial results. An **AllReduce** sums the partial outputs to get the final Y.
Result: Only ONE AllReduce per MLP block (not per matrix multiply).
**Attention Tensor Parallelism**
Multi-head attention is naturally parallelizable — split attention heads across GPUs:
```
Input X (replicated on all GPUs)
GPU 0: Heads 0-15 → Q₀,K₀,V₀ → Attn₀ → O₀ (partial)
GPU 1: Heads 16-31 → Q₁,K₁,V₁ → Attn₁ → O₁ (partial)
AllReduce(O₀ + O₁) → Output
```
Each GPU computes Q, K, V projections for its assigned heads, performs attention, and projects output. A single AllReduce at the end combines results. This is remarkably efficient because attention heads are independent.
**Sequence Parallelism**
Megatron-LM v3 added sequence parallelism for the non-tensor-parallel regions (LayerNorm, dropout, residual connections). These ops operate on the full hidden dimension but can be split along the sequence dimension:
```
Tensor Parallel regions: split on hidden dimension (TP)
Non-TP regions: split on sequence dimension (SP)
Transitions: AllGather / ReduceScatter
```
This reduces the memory footprint of activations in non-TP regions by the TP degree.
**Communication Analysis**
Per transformer layer with TP degree = T:
- 2 AllReduce operations in forward pass (1 for attention, 1 for MLP) → 4 in forward+backward
- Each AllReduce communicates: batch_size × seq_len × hidden_dim elements
- Volume: 4 × 2(T-1)/T × B×S×H bytes per layer per training step
Efficiency requires high-bandwidth interconnect (NVLink: 900 GB/s) — tensor parallelism is typically limited to within a single node (TP=4 or TP=8) with NVLink, while data/pipeline parallelism spans nodes over InfiniBand.
**Tensor parallelism is the foundational distributed strategy for training and serving the largest transformer models** — by splitting every layer's computation across GPUs connected by high-bandwidth links, it enables models with hundreds of billions of parameters to fit in memory and compute efficiently within a single server node.