← AI Terminology
Tensor Parallelism
Tensor parallelism is a distributed training and inference strategy that splits individual weight matrices (tensors) across multiple GPUs — each GPU computes a portion of each matrix multiplication and results are combined via all-reduce communication — enabling single layers that are too large for one GPU to fit across multiple devices.
It is one of the three axes of 3D parallelism alongside data and pipeline parallelism.
It is one of the three axes of 3D parallelism alongside data and pipeline parallelism.
Why It Matters in AI
A single attention layer in a 70B model might have weight matrices of size 8192×8192 — multiple GB each. Tensor parallelism shards these horizontally or vertically across GPUs: GPU 0 computes columns 0–2048; GPU 1 computes columns 2049–4096; etc. This requires high-bandwidth communication after each operation (all-reduce), so tensor parallelism demands fast NVLink (600+ GB/s) and is typically confined to GPUs within the same node. Megatron-LM popularised it for transformer training.
Key Points
| Aspect | Description |
|---|---|
| Degree | TP=8 = 8 GPUs per model replica — typical for frontier training; higher degrees decrease efficiency |
| Megatron-LM | NVIDIA's tensor parallel transformer: --tensor-model-parallel-size 8 — the reference implementation |
| Row parallel | Split weight rows — each GPU holds partial rows; all-reduce output after matrix multiply |
| All-reduce cost | Communication after each tensor-parallel layer — requires NVLink/InfiniBand for low latency |
| Attention split | Split attention heads across GPUs — head 1-8 on GPU 0, heads 9-16 on GPU 1, etc. |
| Column parallel | Split weight columns across GPUs — each GPU holds complete rows, partial columns; all-gather after |
Simple Analogy
A symphony orchestra where four violin sections each play a different part of the violin score simultaneously: the parts must be coordinated (all-reduce) to produce the full musical output, but each section handles only its portion. High-speed communication channels (NVLink) are the conductor's ability to keep all sections in sync.
Common Usage Examples
megatron-lm --tensor-model-parallel-size 8 --pipeline-model-parallel-size 4— 3D parallel configfrom megatron.core.tensor_parallel import ColumnParallelLinear, RowParallelLinear- DeepSpeed: tensor parallelism via
"tensor_parallel": {"tp_size": 4}in ZeRO Stage 3 - vLLM inference:
--tensor-parallel-size 4— split large model across 4 GPUs for inference torch.distributed.all_reduce(tensor)— synchronise partial results across tensor-parallel GPUs
Summary
In short: Tensor parallelism splits individual weight matrices across multiple GPUs that each compute a portion of each operation — enabling layers too large for a single GPU, requiring fast NVLink communication, and forming one axis of the 3D parallelism used to train frontier LLMs.