Multi-GPU Model Parallelization
When the model itself is too large to fit in a single GPU’s memory — even before considering activations and optimizer state — replicating it is not an option. Instead, the model must be split across multiple GPUs. This is model parallelism.
Model parallelism takes two distinct forms that solve different problems, and large-scale training systems often combine both. This article explains pipeline parallelism, tensor parallelism, and how they fit together with data parallelism into a hybrid scaling strategy.
Contents:
- Introduction.
- Pipeline Parallelism.
- Tensor Parallelism.
- Comparing the Two Strategies.
- Hybrid 3D Parallelism.
- Conclusions.
1. Introduction
A modern transformer model with tens of billions of parameters does not fit on a single GPU, regardless of memory size. The parameters alone consume hundreds of gigabytes, before counting activations, gradients, and optimizer state, which together typically multiply the memory footprint by a factor of four to eight.
Data parallelism, which replicates the model on every device, cannot help here — there is no single device large enough to hold one copy. The model must be partitioned, and the work distributed across devices that each hold only a portion of it.
There are two principled ways to do this, and they split the model along different axes.
2. Pipeline Parallelism
Pipeline parallelism assigns contiguous groups of layers — called stages — to different GPUs. A batch flows through stage one, then stage two, then stage three. Each GPU holds the weights for its own stage and nothing else.
The naïve implementation is sequential — GPU two sits idle while GPU one runs the first stage, then vice versa, wasting most of the available compute. To recover throughput, batches are split into micro-batches that overlap: while GPU two processes micro-batch one, GPU one already starts micro-batch two.
This is the classic instruction-pipeline pattern applied to neural network execution. The leftover idle time at the start and end of each step — while the pipeline fills and drains — is called the pipeline bubble. The bubble shrinks as the number of micro-batches increases, but it cannot be eliminated entirely. Schedulers such as GPipe and PipeDream differ in how they overlap forward and backward passes to minimize this overhead.
3. Tensor Parallelism
Tensor parallelism splits a single layer’s weight matrix across GPUs. A matrix multiplication Y = X · W becomes Y = X · [W₁ | W₂] on two GPUs, each computing half of the output columns, with the result concatenated (or summed, depending on which dimension is split).
Unlike pipeline parallelism, tensor parallelism keeps every GPU busy at every moment — there is no bubble. The trade-off is that communication happens inside every layer rather than once between stages, so high-bandwidth interconnects (NVLink rather than PCIe) become essential.
Tensor parallelism is the dominant strategy for the attention and feed-forward blocks of large transformer models, where individual matrix multiplications already saturate a single GPU’s compute.
4. Comparing the Two Strategies
Pipeline and tensor parallelism solve the same problem in fundamentally different ways:
- Split granularity. Pipeline parallelism splits between layers; tensor parallelism splits inside a single operation.
- Communication frequency. Pipeline parallelism communicates once per stage boundary; tensor parallelism communicates inside every layer.
- Bandwidth requirements. Pipeline parallelism tolerates modest bandwidth; tensor parallelism needs very high bandwidth (typically NVLink).
- Idle time. Pipeline parallelism has the pipeline bubble; tensor parallelism is fully synchronous and has none.
- Best fit. Pipeline parallelism suits long, sequential models with clean layer boundaries; tensor parallelism suits wide layers, especially in transformers.
The two are not mutually exclusive. Production systems training the largest models routinely apply tensor parallelism within a node and pipeline parallelism across nodes.
5. Hybrid 3D Parallelism
Training of very large models combines all three strategies — data, tensor, and pipeline parallelism — into a single hybrid strategy:
- Tensor parallel within a node, across the GPUs connected by high-bandwidth NVLink.
- Pipeline parallel across nodes, over slower inter-node links.
- Data parallel across replicas of the entire pipeline, to scale throughput beyond what a single pipeline can sustain.
Choosing the right partition is a balance between compute utilization, memory pressure, and the available interconnect topology. The same model can be partitioned very differently on a cluster of 8-GPU nodes connected by InfiniBand than on a single node of 8 NVLink-connected GPUs.
This combined strategy is sometimes called 3D parallelism, after the three independent dimensions along which the workload is split.
6. Conclusions
Model parallelism is the answer to the question: what do you do when a model is too large to fit on one GPU?
- Pipeline parallelism splits the model between layers, trading throughput for the pipeline bubble.
- Tensor parallelism splits the model inside layers, trading bandwidth for fully overlapped compute.
- Hybrid 3D parallelism combines both with data parallelism to train the largest models on commodity clusters.
Model parallelism is significantly more complex than data parallelism, both in implementation and in performance tuning. It should be reserved for the cases where a model genuinely cannot fit on one device — but at scale, it is the only way forward.