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 viewBox="0 0 443 226" xmlns="http://www.w3.org/2000/svg" style="max-width:100%;height:auto" role="img"><rect x="0" y="0" width="443" height="226" rx="12" fill="#0d1117"/><g font-family="ui-monospace,SFMono-Regular,Menlo,Consolas,&quot;Liberation Mono&quot;,monospace" font-size="14"><text xml:space="preserve" x="20" y="31.7"><tspan fill="#c9d1d9">GPU 0:                    GPU 1:</tspan></text><text xml:space="preserve" x="20" y="50.7"><tspan fill="#c9d1d9">X </tspan><tspan fill="#6e7681">──→</tspan><tspan fill="#c9d1d9"> [A₁] </tspan><tspan fill="#6e7681">──→</tspan><tspan fill="#c9d1d9"> GeLU </tspan><tspan fill="#6e7681">──→</tspan><tspan fill="#c9d1d9">  X </tspan><tspan fill="#6e7681">──→</tspan><tspan fill="#c9d1d9"> [A₂] </tspan><tspan fill="#6e7681">──→</tspan><tspan fill="#c9d1d9"> GeLU </tspan><tspan fill="#6e7681">──→</tspan></text><text xml:space="preserve" x="20" y="69.7"><tspan fill="#c9d1d9">      (col split)               (col split)</tspan></text><text xml:space="preserve" x="20" y="88.7"><tspan fill="#c9d1d9">      </tspan><tspan fill="#6e7681">↓</tspan><tspan fill="#c9d1d9">                         </tspan><tspan fill="#6e7681">↓</tspan></text><text xml:space="preserve" x="20" y="107.7"><tspan fill="#c9d1d9">      [B₁] </tspan><tspan fill="#6e7681">──→</tspan><tspan fill="#c9d1d9"> Y₀              [B₂] </tspan><tspan fill="#6e7681">──→</tspan><tspan fill="#c9d1d9"> Y₁</tspan></text><text xml:space="preserve" x="20" y="126.7"><tspan fill="#c9d1d9">      (row split)               (row split)</tspan></text><text xml:space="preserve" x="20" y="145.7"><tspan fill="#c9d1d9">      </tspan><tspan fill="#6e7681">↓</tspan><tspan fill="#c9d1d9">                         </tspan><tspan fill="#6e7681">↓</tspan></text><text xml:space="preserve" x="20" y="164.7"><tspan fill="#c9d1d9">      </tspan><tspan fill="#6e7681">───────</tspan><tspan fill="#c9d1d9"> AllReduce </tspan><tspan fill="#6e7681">────────</tspan></text><text xml:space="preserve" x="20" y="183.7"><tspan fill="#c9d1d9">              </tspan><tspan fill="#6e7681">↓</tspan></text><text xml:space="preserve" x="20" y="202.7"><tspan fill="#c9d1d9">              Y (complete output)</tspan></text></g></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.