AI Next Labs / Learning
Tensor parallelism
Split every weight matrix across GPUs: column and row splits with real numbers, a TP transformer block, and what it costs in communication.
Published
Written by Jaydeep Raijada.
Instead of giving every GPU a whole layer, tensor parallelism (TP) cuts each weight matrix into pieces and keeps one piece on each GPU permanently. This works because matrix multiplication splits cleanly by rows and columns. The cost is a few all-reduces in every transformer block, which is why TP lives on the fast links inside a single node.
Part 2 of 8 in the 5D Parallelism series. Previous: Data parallelism. Next: Sequence parallelism.
Work through the interactive lesson below: step through each diagram, try the controls, and check yourself with the questions at the end.
References
- Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism (Paper). Shoeybi et al., 2019. The column-then-row split for the MLP and attention, and the f / g conjugate operators.
- Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM (Paper). Narayanan et al., 2021. How TP combines with pipeline and data parallelism, and why TP stays inside a node.
- Reducing Activation Recomputation in Large Transformer Models (Paper). Korthikanti et al., 2022. Measures the activations TP leaves replicated and adds sequence parallelism.
- The Ultra-Scale Playbook (Guide). Hugging Face. The tensor parallelism chapter covers the same splits, with measured throughput and memory as TP grows.
- How to Scale Your Model (Guide). Google DeepMind. Sharded matrix multiplication and collective costs from first principles.