Tensor Parallelism

Chapter 10: Matrix Block Partitioning. In the chapter mind map this icon labels LLM Distrib Computing: Row-wise & Column-wise Tensor Parallelism. The discussion below is excerpted and lightly edited from § Tensor Parallelism and Distributed Matrix Operations in Mathematics for AI and Machine Learning.

In tensor parallel training[^tensor-parallel], large weight matrices are partitioned across multiple devices (GPUs) to enable distributed computation.

[^tensor-parallel]: Fuheng Wu, Distributed AI Systems: A Practical Guide to Building Scalable Training, Inference, and Serving Systems for Production AI. \url{https://www.amazon.com/dp/B0GN8NRVQX/} [^megatron-lm]: NVIDIA Megatron-LM: \url{https://github.com/NVIDIA/Megatron-LM} [^pytorch-tp]: PyTorch tensor parallel (torch.distributed.tensor): \url{https://pytorch.org/docs/stable/distributed.tensor.html} [^all-reduce]: All-reduce: each rank holds a tensor; all ranks receive the reduction (e.g., sum) across ranks. NCCL collective operations: \url{https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/colls.html} [^all-gather]: All-gather: each rank holds a shard; all ranks receive the concatenation of all shards. NCCL collective operations: \url{https://docs.nvidia.com/deeplearning/nccl/user-guide/docs/api/colls.html}

Throughout this subsection the layer is $\mathbf y = W\mathbf x$ with $W \in \mathbb{R}^{m \times n}$.

• Row-wise Partitioning: Partition $W$ row-wise: $W = [W_0; \ldots; W_{k-1}]$ with $W_i \in \mathbb{R}^{(m/k) \times n}$. Each device computes its slice $W_i \mathbf x$ of the output, concatenated via all-gather:
• Column-wise Partitioning: Partition $W$ and $\mathbf x$ column-wise: $W = [W_0, \ldots, W_{k-1}]$, $\mathbf x = [\mathbf x_0; \ldots; \mathbf x_{k-1}]$ with $W_i \in \mathbb{R}^{m \times (n/k)}$ and $\mathbf x_i \in \mathbb{R}^{n/k}$. Each device computes a partial product $W_i \mathbf x_i$, summed via all-reduce:

Megatron-LM uses the row-vector convention $Y = XW$, transposing these names (column-parallel concatenates outputs, row-parallel sums partial outputs).