Tensor Parallelism › Tensor Parallelism

Training a large language model with hundreds of billions of parameters is not something a single GPU can do on its own. A single NVIDIA A100 has 80 GB of memory, and a model like GPT-3 requires roughly 350 GB just to store its weights in 16-bit precision, before accounting for optimizer states, gradients, and activations. The simplest response to this constraint is data parallelism, where each GPU holds a complete copy of the model but processes a different slice of the minibatch. But data parallelism fails precisely when the model itself does not fit on one device. You cannot copy what you cannot hold.

Tensor parallelism solves this differently. Instead of replicating the entire model across GPUs, it splits individual tensors across multiple devices. Each GPU holds only a shard of each weight matrix, performs computation on its shard, and then communicates with peer GPUs to assemble the full result. The model is not replicated; it is partitioned. This lets you train models whose individual layers would overflow any single device's memory.

The idea sounds straightforward, but realizing it efficiently requires care. A naive partitioning strategy might simply distribute layers to different GPUs (inter-layer parallelism), but this creates severe load-balancing problems and requires activations to travel from one device to another between every layer. Tensor parallelism takes a different approach: it operates within a single layer, distributing the matrix multiplications that dominate transformer computation. By splitting weight matrices in mathematically compatible ways, tensor parallelism allows each GPU to do useful work simultaneously on every layer, with communication happening at precisely defined synchronization points rather than at every layer boundary.

This chapter covers the mechanics of tensor parallelism in depth: how weight matrices are split along columns or rows, what communication is required to combine shards into correct outputs, how sequences of column-parallel and row-parallel layers are chained to minimize synchronization, and how the Megatron-LM system from NVIDIA Research formalized these ideas into a practical training framework used for some of the largest models in existence. We will build understanding from first principles, working through the math of why the splits are correct before looking at implementation.

As we discussed in earlier chapters on the transformer architecture, the dominant compute in a transformer layer comes from two sources: the multi-head attention mechanism and the two-layer feed-forward network (FFN). Tensor parallelism in Megatron style targets both of these components, splitting the weight matrices in carefully chosen directions to enable maximum parallelism with minimum communication.