Multi-GPU Model Parallelization
Model parallelism partitions a neural network across several GPUs when the complete model or its execution state does not fit efficiently on one device.
OpenNN currently compiles one network for one device and does not expose built-in pipeline or tensor-parallel graph partitioning. The strategies in this article therefore describe architectures that require external orchestration.
Contents:
1. Introduction
Large models consume memory for parameters, activations, gradients and optimizer state. When that working set exceeds one GPU’s capacity, replicating the entire model with data parallelism is not sufficient.
Model parallelism divides the layers or individual tensor operations between devices. The partition must balance memory, computation and inter-device communication.
2. Pipeline parallelism
Pipeline parallelism assigns consecutive groups of layers, called stages, to different GPUs. Each stage owns its parameters and passes activations to the next stage during forward propagation and gradients in the opposite direction during back propagation.
Micro-batches keep several stages busy at the same time. Some idle time remains while the pipeline fills and drains; this is the pipeline bubble. More micro-batches can reduce its relative cost but increase scheduling complexity.
3. Tensor parallelism
Tensor parallelism splits an operation inside a layer. For example, different GPUs can hold partitions of a weight matrix and compute different output columns before the partial results are concatenated or reduced.
This keeps devices active within the same layer, but it requires communication during many layer operations. It therefore benefits from high-bandwidth, low-latency links and carefully selected partition dimensions.
4. Compare the two strategies
- Split granularity: pipeline parallelism splits between layers; tensor parallelism splits inside an operation.
- Communication: pipeline stages exchange activations and gradients; tensor partitions exchange partial results inside layers.
- Idle time: pipeline execution has fill and drain bubbles; tensor execution is synchronized at its collectives.
- Best fit: pipelines suit clear stage boundaries, while tensor partitioning suits large, regular matrix operations.
The two approaches can be combined when neither one alone provides the required memory capacity and throughput.
5. Hybrid 3D parallelism
Very large distributed workloads can combine three independent dimensions:
- Tensor parallelism between tightly connected GPUs.
- Pipeline parallelism between groups of layers or nodes.
- Data parallelism between replicas of the complete partitioned pipeline.
The optimal mapping depends on model shape, device memory and the topology and bandwidth of the available interconnects.
6. Conclusions
- Pipeline parallelism partitions consecutive model stages.
- Tensor parallelism partitions operations within a layer.
- Hybrid strategies combine model and data parallel dimensions.
- Current OpenNN applications must implement device partitioning and communication outside the public network API.