Tensor Parallelism is the model parallelism technique that splits individual weight matrices and tensors across multiple GPUs, with each GPU computing a portion of each layer's output — enabling models with layers too large for single-GPU memory by distributing matrix multiplications column-wise or row-wise and synchronizing results through collective communication operations like all-reduce and all-gather.

Tensor Parallelism Fundamentals:

• Matrix Partitioning: for matrix multiplication Y = XW, split weight matrix W across GPUs; column-wise split: each GPU computes Y_i = X·W_i (partial output); row-wise split: each GPU computes Y = X_i·W (partial input)
• Communication Patterns: column-wise split requires all-gather to combine partial outputs; row-wise split requires all-reduce to sum partial results; communication volume = batch_size × sequence_length × hidden_dim
• Intra-Layer Parallelism: unlike pipeline parallelism (distributes layers), tensor parallelism distributes computation within each layer; all GPUs process same batch simultaneously
• Scaling Characteristics: near-linear scaling within a node (8 GPUs with NVLink); efficiency drops with inter-node communication; typically limited to 8-16 GPUs per tensor parallel group

<svg viewBox="0 0 760 470" xmlns="http://www.w3.org/2000/svg" font-family="-apple-system,Segoe UI,Roboto,sans-serif">
  <rect x="0" y="0" width="760" height="470" fill="#0d1117"/>
  <text x="380" y="28" fill="#e6edf3" font-size="21" font-weight="700" text-anchor="middle">Tensor Parallelism — Split Layers Across GPUs</text>
  <text x="380" y="48" fill="#8b98a5" font-size="12" text-anchor="middle">partition weight matrices column-wise or row-wise — each GPU computes a shard of every layer</text>

  <!-- Column parallel (MLP) -->
  <rect x="30" y="65" width="700" height="170" rx="6" fill="#080d14" stroke="#233043" stroke-width="1"/>
  <text x="380" y="86" fill="#e6edf3" font-size="11" font-weight="600" text-anchor="middle">Column Parallel Linear (MLP first layer)</text>

  <!-- Input X -->
  <rect x="55" y="110" width="60" height="80" rx="3" fill="#0b1220" stroke="#60a5fa" stroke-width="1"/>
  <text x="85" y="145" fill="#93c5fd" font-size="9" text-anchor="middle">X</text>
  <text x="85" y="160" fill="#6b7684" font-size="7" text-anchor="middle">[seq, d]</text>

  <!-- Broadcast arrow -->
  <path d="M117,150 L155,130" fill="none" stroke="#8b98a5" stroke-width="0.8"/>
  <path d="M117,150 L155,150" fill="none" stroke="#8b98a5" stroke-width="0.8"/>
  <path d="M117,150 L155,170" fill="none" stroke="#8b98a5" stroke-width="0.8"/>
  <text x="136" y="118" fill="#6b7684" font-size="7" text-anchor="middle">broadcast</text>

  <!-- Weight shards on GPUs -->
  <rect x="160" y="100" width="70" height="35" rx="3" fill="#14261f" stroke="#34d399" stroke-width="0.8"/>
  <text x="195" y="115" fill="#6ee7b7" font-size="7" text-anchor="middle">W₁ (GPU 0)</text>
  <text x="195" y="129" fill="#6b7684" font-size="6.5" text-anchor="middle">[d, d/N]</text>

  <rect x="160" y="140" width="70" height="35" rx="3" fill="#1c1633" stroke="#a78bfa" stroke-width="0.8"/>
  <text x="195" y="155" fill="#c4b5fd" font-size="7" text-anchor="middle">W₂ (GPU 1)</text>
  <text x="195" y="169" fill="#6b7684" font-size="6.5" text-anchor="middle">[d, d/N]</text>

  <rect x="160" y="180" width="70" height="35" rx="3" fill="#2a1a0a" stroke="#f59e0b" stroke-width="0.8"/>
  <text x="195" y="195" fill="#fbbf24" font-size="7" text-anchor="middle">W₃ (GPU 2)</text>
  <text x="195" y="209" fill="#6b7684" font-size="6.5" text-anchor="middle">[d, d/N]</text>

  <!-- Matmul arrows -->
  <path d="M232,117 L268,117" fill="none" stroke="#3a4453" stroke-width="0.8"/>
  <path d="M232,157 L268,157" fill="none" stroke="#3a4453" stroke-width="0.8"/>
  <path d="M232,197 L268,197" fill="none" stroke="#3a4453" stroke-width="0.8"/>

  <!-- Partial outputs -->
  <rect x="272" y="100" width="55" height="35" rx="3" fill="#14261f" stroke="#34d399" stroke-width="0.6"/>
  <text x="299" y="121" fill="#34d399" font-size="7" text-anchor="middle">Y₁</text>

  <rect x="272" y="140" width="55" height="35" rx="3" fill="#1c1633" stroke="#a78bfa" stroke-width="0.6"/>
  <text x="299" y="161" fill="#a78bfa" font-size="7" text-anchor="middle">Y₂</text>

  <rect x="272" y="180" width="55" height="35" rx="3" fill="#2a1a0a" stroke="#f59e0b" stroke-width="0.6"/>
  <text x="299" y="201" fill="#f59e0b" font-size="7" text-anchor="middle">Y₃</text>

  <!-- All-reduce or concat -->
  <path d="M329,117 L370,150" fill="none" stroke="#8b98a5" stroke-width="0.8"/>
  <path d="M329,157 L370,150" fill="none" stroke="#8b98a5" stroke-width="0.8"/>
  <path d="M329,197 L370,150" fill="none" stroke="#8b98a5" stroke-width="0.8"/>

  <rect x="374" y="130" width="80" height="40" rx="4" fill="#0b1220" stroke="#f87171" stroke-width="1"/>
  <text x="414" y="148" fill="#fca5a5" font-size="8" text-anchor="middle">all-reduce</text>
  <text x="414" y="162" fill="#6b7684" font-size="7" text-anchor="middle">(sum partials)</text>

  <!-- Final output -->
  <path d="M456,150 L488,150" fill="none" stroke="#3a4453" stroke-width="0.8"/>
  <polygon points="486,147 492,150 486,153" fill="#3a4453"/>
  <rect x="496" y="130" width="55" height="40" rx="3" fill="#0b1220" stroke="#60a5fa" stroke-width="0.8"/>
  <text x="523" y="155" fill="#93c5fd" font-size="9" text-anchor="middle">Y</text>

  <!-- Annotation -->
  <text x="620" y="115" fill="#8b98a5" font-size="8">column split: each GPU</text>
  <text x="620" y="131" fill="#8b98a5" font-size="8">gets d/N output columns</text>
  <text x="620" y="150" fill="#34d399" font-size="8">→ concat = full output</text>
  <text x="620" y="170" fill="#6b7684" font-size="7.5">comm: 1 all-reduce per layer</text>
  <text x="620" y="186" fill="#6b7684" font-size="7.5">(MLP: column then row split)</text>

  <!-- Attention parallel -->
  <rect x="30" y="245" width="345" height="130" rx="6" fill="#080d14" stroke="#233043" stroke-width="1"/>
  <text x="202" y="265" fill="#e6edf3" font-size="10" font-weight="600" text-anchor="middle">Attention Head Parallelism</text>
  <text x="50" y="288" fill="#8b98a5" font-size="8.5">32 attention heads, 4 GPUs → 8 heads per GPU</text>
  <text x="50" y="306" fill="#8b98a5" font-size="8.5">each GPU computes full attention for its head subset</text>
  <text x="50" y="324" fill="#8b98a5" font-size="8.5">all-reduce after output projection</text>
  <text x="50" y="346" fill="#6b7684" font-size="7.5">natural split: attention heads are independent (no cross-head interaction)</text>
  <text x="50" y="362" fill="#6b7684" font-size="7.5">GQA: split Q heads evenly, replicate shared KV heads</text>

  <!-- Communication cost -->
  <rect x="390" y="245" width="340" height="130" rx="6" fill="#080d14" stroke="#233043" stroke-width="1"/>
  <text x="560" y="265" fill="#e6edf3" font-size="10" font-weight="600" text-anchor="middle">Communication Cost</text>
  <text x="410" y="288" fill="#f87171" font-size="8.5" font-weight="600">2 all-reduces per transformer layer:</text>
  <text x="410" y="306" fill="#8b98a5" font-size="8">1. after attention output projection</text>
  <text x="410" y="322" fill="#8b98a5" font-size="8">2. after MLP second linear</text>
  <text x="410" y="344" fill="#f59e0b" font-size="8.5" font-weight="600">volume per all-reduce: 2×(N-1)/N × seq×d bytes</text>
  <text x="410" y="362" fill="#6b7684" font-size="7.5">→ must use NVLink (900 GB/s), not InfiniBand (50 GB/s)</text>

  <!-- Bottom -->
  <rect x="30" y="388" width="700" height="42" rx="6" fill="#080d14" stroke="#233043" stroke-width="1"/>
  <text x="380" y="408" fill="#e6edf3" font-size="9" font-weight="600" text-anchor="middle">TP is always within a node (NVLink). PP across nodes (IB). DP across pods.</text>
  <text x="380" y="422" fill="#6b7684" font-size="7.5" text-anchor="middle">Llama 3 405B: TP=8 (one node), PP=16 (across nodes), DP=128 (across pods) = 16,384 GPUs</text>

  <text x="380" y="452" fill="#6b7684" font-size="11" text-anchor="middle">Tensor parallelism cuts per-GPU memory by N× — the only way to fit a single layer when weights exceed one GPU.</text>
</svg>

Megatron-LM Tensor Parallelism:

• Attention Layer Splitting: Q, K, V projections split column-wise across GPUs; each GPU computes attention for subset of heads; output projection split row-wise; requires 2 all-reduce operations per attention layer
• MLP Layer Splitting: first linear layer (hidden → intermediate) split column-wise; activation function applied independently; second linear layer (intermediate → hidden) split row-wise; 2 all-reduce operations per MLP
• Communication Minimization: careful splitting strategy minimizes communication; only 2 all-reduce per Transformer block (attention + MLP); communication overlapped with computation where possible
• Identity Operators: inserts identity operators in forward pass that become all-reduce in backward pass (and vice versa); elegant implementation using autograd

Column-Wise Parallelism:

• Operation: Y = X·W where W is split column-wise; W = [W_1, W_2, ..., W_N] across N GPUs; each GPU computes Y_i = X·W_i
• Output Combination: concatenate partial outputs [Y_1, Y_2, ..., Y_N] to form full output Y; requires all-gather communication
• Use Cases: first layer of MLP, Q/K/V projections in attention; enables independent computation of output dimensions
• Memory Distribution: each GPU stores 1/N of weights; activation memory not reduced (all GPUs process full batch)

Row-Wise Parallelism:

• Operation: Y = X·W where W is split row-wise; W = [W_1; W_2; ...; W_N] (stacked vertically); input X also split; each GPU computes Y_i = X_i·W_i
• Output Combination: sum partial outputs Y = Σ Y_i; requires all-reduce communication
• Use Cases: second layer of MLP, output projection in attention; follows column-wise split to minimize communication
• Input Splitting: requires input X to be split across GPUs; typically X is already split from previous column-wise layer

Communication Optimization:

• All-Reduce Fusion: fuses multiple all-reduce operations into single communication; reduces latency overhead; NCCL automatically fuses small all-reduces
• Communication Overlap: starts all-reduce as soon as partial results are ready; overlaps with computation of next layer; requires careful scheduling
• Gradient All-Reduce: backward pass requires all-reduce for gradients; same communication volume as forward pass; can overlap with backward computation
• High-Bandwidth Interconnect: NVLink (300-600 GB/s within node) essential for efficiency; InfiniBand (200-400 Gb/s across nodes) for multi-node; communication-bound without fast interconnect

Memory Distribution:

• Weight Memory: each GPU stores 1/N of model weights; enables models N× larger than single GPU capacity
• Activation Memory: not reduced by tensor parallelism (all GPUs process full batch); combine with pipeline parallelism or activation checkpointing to reduce activation memory
• Optimizer State Memory: each GPU stores optimizer states for its 1/N of weights; total optimizer memory reduced by N×
• Gradient Memory: each GPU computes gradients for its 1/N of weights; gradient memory reduced by N×

Sequence Parallelism Extension:

• Motivation: LayerNorm and Dropout activations not split by standard tensor parallelism; consume significant memory for long sequences
• Sequence Dimension Splitting: splits sequence length across GPUs for LayerNorm/Dropout; each GPU processes subset of tokens
• Communication: requires all-gather before attention (each token attends to all tokens); all-reduce after attention; additional communication but reduces activation memory
• Memory Savings: reduces activation memory by N× for LayerNorm/Dropout; critical for very long sequences (>8K tokens)

Combining with Other Parallelism:

• Tensor + Data Parallelism: tensor parallelism within groups, data parallelism across groups; example: 64 GPUs = 8 TP × 8 DP
• Tensor + Pipeline Parallelism: each pipeline stage uses tensor parallelism; enables very large models; Megatron-LM uses TP within nodes, PP across nodes
• 3D Parallelism: DP × TP × PP; example: 512 GPUs = 8 DP × 8 TP × 8 PP; matches parallelism to hardware topology
• Optimal Configuration: TP within nodes (high bandwidth), PP across nodes (lower bandwidth), DP for remaining GPUs; automated search or manual tuning

Framework Support:

• Megatron-LM (NVIDIA): reference implementation of tensor parallelism for Transformers; highly optimized; used for training GPT, BERT, T5 at scale
• DeepSpeed: supports tensor parallelism via Megatron integration; combines with ZeRO optimizer; comprehensive parallelism toolkit
• Fairscale: PyTorch-native tensor parallelism; modular design; easier integration than Megatron; used by Meta
• Alpa: automatic parallelization including tensor parallelism; compiler-based approach; supports JAX

Implementation Considerations:

• Collective Communication: uses NCCL (NVIDIA) or MPI for all-reduce/all-gather; requires proper initialization and synchronization
• Determinism: tensor parallelism is deterministic (same results as single GPU); unlike data parallelism which may have non-deterministic reduction order
• Gradient Clipping: must clip gradients after all-reduce; clipping before all-reduce gives incorrect results
• Batch Normalization: requires synchronization across tensor parallel group; typically replaced with LayerNorm in Transformers

Performance Analysis:

• Computation Scaling: each GPU does 1/N of computation; ideal speedup = N×
• Communication Overhead: 2 all-reduce per Transformer block; overhead = communication_time / computation_time; want ratio < 10-20%
• Bandwidth Requirements: all-reduce volume = 2 × batch_size × sequence_length × hidden_dim per block; requires high bandwidth for efficiency
• Scaling Efficiency: 90-95% efficiency within node (NVLink); 70-80% efficiency across nodes (InfiniBand); diminishing returns beyond 16 GPUs

Practical Guidelines:

• When to Use: model layers don't fit on single GPU; have high-bandwidth interconnect (NVLink); need low-latency parallelism
• Tensor Parallel Size: 2-8 GPUs typical; 8 GPUs within node optimal; beyond 8 requires inter-node communication (less efficient)
• Batch Size: larger batches amortize communication overhead; batch_size × sequence_length should be large (>1M tokens total)
• Debugging: start with TP=2 to verify correctness; scale up gradually; use smaller models for initial debugging

Tensor parallelism is the fine-grained parallelism technique that enables training of models with individual layers too large for single-GPU memory — by splitting weight matrices and carefully orchestrating collective communication, it achieves near-linear scaling within high-bandwidth GPU clusters, making it essential for frontier models where even a single attention layer exceeds GPU capacity.