Source-linked AI summary
Chimera: Efficiently Training Large-Scale Neural Networks with Bidirectional Pipelines
Shigang Li, Torsten Hoefler
TL;DR
Training increasingly large models requires distributed execution, while pipeline parallelism introduces bubbles and imbalanced activation memory. Chimera combines bidirectional pipelines in a synchronous schedule, achieving improved throughput and a balance among pipeline efficiency, memory cost, and convergence friendliness. Evaluations on language-model training show throughput improvements over counterpart pipeline approaches at up to 2,048 GPU nodes.
Problem
Large-scale model training requires distributed execution, while pipeline parallelism suffers from pipeline bubbles and uneven activation-memory requirements.
Method
Chimera combines down and up pipelines running in opposite directions and merges their schedules into a synchronous bidirectional pipeline.
Results
Chimera significantly improves training throughput over counterpart pipeline approaches in empirical language-model training on up to 2,048 GPU nodes.
Takeaways & Limitations
Chimera achieves a balance among pipeline efficiency, memory cost, and convergence friendliness while scaling large-model training.
Takeaways & Limitations
The evaluation uses the same model partitioning as PipeDream-2BW; alternative weight-balancing partitions are outside the paper’s scope.
Abstract
from arXiv · showhide
Training large deep learning models at scale is very challenging. This paper proposes Chimera, a novel pipeline parallelism scheme which combines bidirectional pipelines for efficiently training large-scale models. Chimera is a synchronous approach and therefore no loss of accuracy, which is more convergence-friendly than asynchronous approaches. Compared with the latest synchronous pipeline approach, Chimera reduces the number of bubbles by up to 50%; benefiting from the sophisticated scheduling of bidirectional pipelines, Chimera has a more balanced activation memory consumption. Evaluations are conducted on Transformer based language models. For a GPT-2 model with 1.3 billion parameters running on 2,048 GPU nodes of the Piz Daint supercomputer, Chimera improves the training throughput by 1.16x-2.34x over the state-of-the-art synchronous and asynchronous pipeline approaches.
1 INTRODUCTION
Large Transformer models require distributed training, but pipeline parallelism faces activation-memory imbalance and pipeline bubbles. Chimera addresses these challenges with synchronous bidirectional pipelines, improving utilization while balancing memory consumption.
- Large models must be distributed across many accelerators because full model replicas do not fit in individual accelerator memory.
- Pipeline parallelism lowers communication cost through point-to-point activation transfers, but its schedules create uneven activation-memory requirements across stages.
- Larger micro-batches improve computational efficiency by increasing matrix-operation reuse and reducing pipeline bubbles.
- Chimera runs fully packed bidirectional pipelines synchronously on the same accelerators.
- Figure 1 compares approaches by bubble ratio, memory cost, and best throughput for GPT-2 on 2,048 GPU nodes.Memory-cost entries indicate when activation recomputation is required because of out-of-memory conditions.
- 1.38x-2.34x per iteration: Chimera improves end-to-end performance over synchronous training for comparable GPT-2 models on 2,048 Piz Daint GPU nodes.The reported savings exceed $1.2m-$5m when training very large models on practical systems.
2 BACKGROUND AND RELATED WORK
Large models exceed the practical limits of pure data parallelism, motivating pipeline-based distribution, but efficient training must balance communication, bubbles, memory, and convergence. Existing synchronous methods avoid weight staleness at the cost of bubbles, whereas asynchronous methods reduce bubbles but introduce stale weights and uncertain convergence quality.
- Scaling large models: Large models may not fit on one accelerator, so data parallelism alone can be inefficient or infeasible.Data parallelism replicates the entire model and synchronizes gradients across workers.
- Scaling large models: Pipeline parallelism partitions consecutive layers across stages and transfers activations point-to-point, commonly reducing communication relative to operator parallelism.The mini-batch is divided into micro-batches that are pipelined across stages to improve resource utilization.
- Pipeline efficiency: Increasing the micro-batch count can reduce bubbles, but practical limits on mini-batch and micro-batch sizes constrain this strategy.Large mini-batches can compromise convergence, while small micro-batches may reduce accelerator efficiency; data parallelism can also reduce N.
- Memory and convergence: Existing approaches trade off activation memory, pipeline efficiency, and weight freshness across different scheduling choices.GPipe stores activations proportional to N, 1F1B-based methods inject up to D micro-batches, and asynchronous methods avoid bubbles with stale weights.
- Memory and convergence: Synchronous approaches retain standard mini-batch SGD accuracy guarantees, while asynchronous approaches may lose ideal accuracy because of weight staleness.Chimera is presented as balancing synchronization, utilization, memory consumption, and configurability across pipeline schemes.
3 THE SCHEME OF CHIMERA
Chimera combines bidirectional pipelines to improve utilization while preserving synchronous training and balancing activation memory. Its performance model guides configurations, communication overlap, and extensions to larger micro-batch counts or more pipelines.
- 3.1 Bidirectional Pipelines: Chimera combines down and up pipelines mapped in opposite directions across the same accelerators.The two pipelines partition micro-batches and each uses a 1F1B schedule.
- 3.1 Bidirectional Pipelines: Chimera supports N< D by partitioning micro-batches between the two pipelines as evenly as possible.When N=1, only one micro-batch runs on a single pipeline.
- 3.2 Communication Scheme: Chimera overlaps gradient synchronization with pipeline computation using asynchronous allreduce launched when stage gradients become available.Point-to-point communication transfers activations and input gradients, while allreduce synchronizes weight gradients across stage replicas.
- 3.2 Communication Scheme: Chimera replicates bidirectional pipelines W times, creating a hybrid pipeline-and-data-parallel configuration with W·D workers.The local gradient size remains unchanged, while W times more processes participate in synchronization.
- 3.4 Configuration Selection Based on Performance Modelling: Larger micro-batches generally improve efficiency, so Chimera greedily chooses the largest size fitting device memory and reduces tuning space.The approach exploits Chimera’s reduced bubble problem compared with synchronous pipeline methods.
- 3.4 Configuration Selection Based on Performance Modelling: The performance model uses latency-bandwidth communication costs and accounts for communication that remains after overlap to select W and D.It models point-to-point communication, allreduce, computation, and the maximum unoverlapped synchronization cost across workers.
- 3.5 Scaling Micro-batches and 3.6 Generalization: For more than D micro-batches, Chimera concatenates basic scheduling units, while forward doubling removes intermediate bubbles.The generalized design trades fewer bubbles and more balanced activation memory for higher gradient-synchronization and weight-memory costs.
4 EXPERIMENTAL EVALUATION
Experiments on Piz Daint and a 32-GPU V100 cluster evaluate Chimera against synchronous and asynchronous pipeline baselines for Transformer models. Chimera consistently improves throughput while reducing bubble and memory-balance constraints across configurations and scales.
- 4.1 Memory Consumption: Chimera’s bidirectional schedule produces more balanced memory consumption than PipeDream-2BW and DAPPLE, whose first workers commonly reach peak memory.PipeDream-2BW and DAPPLE can additionally require recomputation or encounter out-of-memory conditions under coarser pipeline stages.
- 4.2.2 Performance Modelling of Chimera.: The performance model’s error remains within 10% for Bert-48 and GPT-2, and GPT-2’s model-selected configuration is within 1.7% of the best performance.The model accurately selects Bert-48’s best configuration, while GPT-2’s best configuration differs from the selected one.
- 4.2.3 Comparison with the Best Performance.: Chimera outperforms all evaluated baselines at every tested scale for both Bert-48 and GPT-2.The comparison includes GPipe, GEMS, DAPPLE, PipeDream, and PipeDream-2BW.
- 4.2.3 Comparison with the Best Performance.: 2.01x and 1.16x are Chimera’s GPT-2 throughput gains over PipeDream and PipeDream-2BW on 2,048 nodes.Against synchronous GPipe, GEMS, and DAPPLE, the corresponding gains are 1.42x, 2.34x, and 1.38x.
- 4.2.3 Comparison with the Best Performance.: 91.4% parallel efficiency is achieved when scaling GPT-2 from 512 to 2,048 nodes in weak scaling.The result uses 512 nodes as the baseline.
- 4.3 Scale to Large Mini-Batches on a Given Number of Machines: 1.13x, 1.18x, 2.60x, and 1.34x are Chimera’s average speedups over PipeDream-2BW, GPipe, GEMS, and DAPPLE when scaling GPT-2 to large mini-batches.Chimera with forward doubling removes intermediate bubbles while activation recomputation is required in both compared methods.
5 CONCLUSION
Chimera improves large-scale neural network pipelining by balancing pipeline efficiency, memory cost, and convergence friendliness. Experiments on up to 2,048 GPU nodes show higher training throughput than other pipeline approaches.
- Chimera achieves the best balance among pipeline efficiency, memory cost, and convergence friendliness.
- Chimera significantly improves training throughput over counterpart pipeline approaches on large language model workloads.The reported experiments scale to 2,048 GPU nodes.
- The authors identify reducing gradient-synchronization communication through sparsification and quantization as future work.
A.1 SUMMARY OF THE EXPERIMENTS REPORTED
The experiments were conducted on the CSCS Piz Daint supercomputer using GPU acceleration and PyTorch's GLOO distributed backend. The Chimera artifact is available through a persistent GitHub identifier.
- The evaluation used the CSCS Piz Daint supercomputer with GPU acceleration.Each Cray XC50 node contains one NVIDIA Tesla P100 GPU with 16 GB memory.
- PyTorch's GLOO served as the distributed backend for the experiments.
- The Chimera source artifact is identified by a persistent GitHub URL.The listed persistent identifier is https://github.com/Shigangli/Chimera.
A.2 BASELINE EXPERIMENTAL SETUP, AND MODIFICATIONS MADE FOR THE PAPER
The baseline environment used Cray XC50 hardware with Tesla P100 GPUs and a Lustre filesystem. Experiments used SUSE SLES, gcc, PyTorch, BERT, GPT-2, and stochastic gradient descent on Wikipedia and WikiText-2 datasets.
- Each baseline Cray XC50 node included a 12-core Intel Xeon E5-2690 CPU and one NVIDIA Tesla P100 GPU.
- The baseline filesystem was Lustre.
- The reported software setup used SUSE SLES 11.3, gcc 9.3.0, PyTorch 1.6, BERT, GPT-2, and stochastic gradient descent.Input datasets included Wikipedia and WikiText-2.