Distributed training: data, tensor and pipeline parallelism
Distributed training spreads the workload of training a neural network across multiple accelerators. Data parallelism replicates the model and splits the data, while tensor and pipeline parallelism split the model itself. These techniques can be combined to train models that are too large for a single device.
The need for distributed training
Modern deep learning models can be larger than the memory of a single GPU, and training them on one device can take too long. Multi-GPU setups accelerate training and allow fitting large models in memory that would not otherwise fit on a single GPU, as explained in the Hugging Face Transformers documentation on parallelism methods Hugging Face's parallelism methods guide.
There are several ways to parallelize the workload. Each method splits either the data or the model differently. The choice depends on the model size, the number of GPUs, and the speed of the connections between them.
A common approach is to use a combination of methods, often called 3D parallelism, to achieve both capacity and efficiency.
Data parallelism
Data parallelism is the most straightforward method. It evenly distributes data across multiple GPUs while each GPU holds a complete copy of the model. Each GPU processes its portion of the data, and the results are synchronized and combined at the end, as described in the same Hugging Face guide Hugging Face's parallelism methods guide.
Two common implementations are DataParallel and DistributedDataParallel. DataParallel works on a single machine with multiple GPUs, while DistributedDataParallel works across multiple machines and is recommended because it reduces communication overhead and scales better.
A more memory-efficient form of data parallelism is the Zero Redundancy Optimizer (ZeRO), developed by Microsoft. ZeRO partitions parameters, gradients, and optimizer states across data parallel processes to reduce memory usage, as stated in the ZeRO paper the ZeRO paper. ZeRO has three stages: stage 1 partitions optimizer states, stage 2 also partitions gradients, and stage 3 also partitions parameters.
Tensor parallelism
Tensor parallelism splits the model itself across accelerators. The Megatron-LM paper from NVIDIA introduced a simple intra-layer model parallel approach that enables training transformer models with billions of parameters without requiring a new compiler or library changes the Megatron-LM paper. It inserts a few communication operations in native PyTorch.
In tensor parallelism, the layers of the model are split, so each GPU holds a portion of every layer's weights and computations. This is useful when the largest layer of a model does not fit into the memory of a single GPU.
Tensor parallelism is orthogonal and complementary to pipeline parallelism, meaning the two can be used together.
Pipeline parallelism
Pipeline parallelism divides a model into segments of layers, and each segment is placed on a separate accelerator. The GPipe library from Google introduced a batch-splitting pipelining algorithm that allows almost linear speedup when a model is partitioned across multiple accelerators the GPipe paper.
During the forward pass, the first group of layers processes a batch and passes the intermediate results to the next group. The backward pass goes in reverse order. Because only one segment is active at a time, GPU utilization can be unbalanced, but the technique enables training networks that are too large for a single device.
Pipeline parallelism is especially useful for models that are expressed as a sequence of layers, such as Transformers and convolutional networks.
Combining strategies and tools
For very large models, data, tensor, and pipeline parallelism are often combined. The DeepSpeed library, also from Microsoft, supports a combination called 3D parallelism and has been used to train models such as Megatron-Turing NLG with 530 billion parameters the DeepSpeed site.
DeepSpeed integrates with popular frameworks like Transformers and Accelerate. The Hugging Face Accelerate library allows the same PyTorch code to run across any distributed configuration by adding a few lines of code, and it supports DeepSpeed and fully sharded data parallelism the Accelerate documentation.
The choice of strategy depends on the model size and the hardware. For a model that fits on one GPU, data parallelism is sufficient. For a model that does not fit, pipeline parallelism or ZeRO may be used. If the largest layer itself does not fit, tensor parallelism or ZeRO is recommended, according to the Hugging Face guide.
| Method | Splits | Memory use on each GPU | Communication | Typical use |
|---|---|---|---|---|
| Data parallelism | Replicates model, splits data | Full model copy | Gradient averaging | Moderate model sizes |
| Tensor parallelism | Splits each layer within a GPU group | Partial layer weights | Frequent small all-reduce | Very large layers |
| Pipeline parallelism | Splits layers across GPUs | One layer group | Send activations forward and gradients back | Models too large for one GPU |
| ZeRO | Partitions parameters, gradients, optimizer states | Partitioned states | Collective communication for updates | Large models with limited memory |
Frequently asked questions
What is the difference between data and model parallelism?
Data parallelism replicates the entire model on each GPU and splits the training data, while model parallelism splits the model itself across GPUs. Both accelerate training but for different reasons: data parallelism increases throughput, model parallelism reduces the memory footprint per GPU.
When should tensor parallelism be used instead of pipeline parallelism?
Tensor parallelism is used when a single layer is too large to fit on one GPU, because it splits the layer across multiple devices. Pipeline parallelism is used when the entire model is too large to fit, by distributing different layers to different GPUs.
What is ZeRO and how does it improve data parallelism?
ZeRO is a memory optimization technique from Microsoft that partitions optimizer states, gradients, and parameters across data parallel processes. It reduces memory redundancy while keeping communication volume low, allowing larger models to be trained.
How do DeepSpeed and Accelerate help with distributed training?
DeepSpeed provides system innovations such as ZeRO and 3D parallelism, making large-scale training more efficient. Accelerate is a library that lets the same PyTorch code run on any distributed setup with minimal code changes, and it can integrate with DeepSpeed.