Source-linked AI summary
Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM
Deepak Narayanan, Mohammad Shoeybi, Jared Casper, Patrick LeGresley, Mostofa Patwary, Vijay Anand Korthikanti, Dmitri Vainbrand, Prethvi Kashinkunti, Julie Bernauer, Bryan Catanzaro, Amar Phanishayee, Matei Zaharia
TL;DR
Training large language models is limited by GPU memory, computational cost, and scaling problems when parallelism is used naively. The paper composes tensor, pipeline, and data parallelism through PTD-P, studies their trade-offs, and introduces an interleaved pipeline schedule. It reports 502 petaFLOP/s for a trillion-parameter model on 3072 GPUs, with 52% of theoretical peak per-GPU throughput.
Problem
Large language models can exceed GPU memory capacity and require unrealistically long training times, while naive parallelism introduces scaling problems at thousands of GPUs.
Method
The paper composes pipeline, tensor, and data parallelism through PTD-P, proposes an interleaved pipeline schedule, and analyzes trade-offs for configuring distributed training.
Results
502 petaFLOP/s was achieved for a trillion-parameter model on 3072 GPUs, with 52% of theoretical peak per-GPU throughput.
Takeaways & Limitations
Combining pipeline parallelism across servers with tensor parallelism within servers and data parallelism enables practical training of trillion-parameter models at scale.
Takeaways & Limitations
The approach uses heuristics rather than automatically exploring the search space of parallelism strategies, and data parallelism alone is insufficient for very large models with limited batch sizes.
Abstract
from arXiv · showhide
Large language models have led to state-of-the-art accuracies across a range of tasks. However, training these models efficiently is challenging for two reasons: a) GPU memory capacity is limited, making it impossible to fit large models on even a multi-GPU server, and b) the number of compute operations required to train these models can result in unrealistically long training times. Consequently, new methods of model parallelism such as tensor and pipeline parallelism have been proposed. Unfortunately, naive usage of these methods leads to fundamental scaling issues at thousands of GPUs, e.g., due to expensive cross-node communication or devices spending significant time waiting on other devices to make progress. In this paper, we show how different types of parallelism methods (tensor, pipeline, and data parallelism) can be composed to scale to thousands of GPUs and models with trillions of parameters. We survey techniques for pipeline parallelism and propose a novel interleaved pipeline parallelism schedule that can improve throughput by 10+% with memory footprint comparable to existing approaches. We quantitatively study the trade-offs between tensor, pipeline, and data parallelism, and provide intuition as to how to configure distributed training of a large model. Our approach allows us to perform training iterations on a model with 1 trillion parameters at 502 petaFLOP/s on 3072 GPUs with achieved per-GPU throughput of 52% of theoretical peak. Our code is open sourced at https://github.com/nvidia/megatron-lm.
1 INTRODUCTION
Training state-of-the-art language models is constrained by GPU memory and enormous computational demands. The paper studies how tensor, pipeline, and data parallelism can be combined to scale training while managing communication, efficiency, and pipeline idle time.
- Motivation: Large language models increasingly exceed the memory capacity of individual GPUs and require impractically large numbers of compute operations to train.These constraints motivate distributed model-parallel training approaches.
- Research question: The paper asks how parallelism techniques can be combined to maximize training throughput for a fixed batch size while retaining strict optimizer semantics.The challenge arises because combining techniques creates non-trivial interactions.
- Approach: PTD-P combines pipeline parallelism across servers, tensor parallelism within servers, and data parallelism to train trillion-parameter models on thousands of GPUs.The approach targets practical scaling with high-bandwidth links and reports 52% of peak device throughput.
- Trade-offs: Parallelism choices affect communication, kernel efficiency, and pipeline bubbles, and sub-optimal tensor-pipeline combinations can produce up to 2× lower throughput.Tensor parallelism is effective within multi-GPU servers, while pipeline parallelism is needed for larger models.
- Pipeline scheduling: The proposed interleaved pipeline schedule can improve throughput by as much as 10% over prior schedules with comparable memory footprint.Microbatch size also affects memory, arithmetic efficiency, and pipeline bubbles, with an optimal value increasing throughput by 15% in experiments.
- Scope: The authors provide practical heuristics for configuring distributed training rather than automatically searching the space of parallelism strategies.The paper explicitly distinguishes its heuristic approach from automated exploration methods.
2 MODES OF PARALLELISM
This section combines pipeline, tensor, and data parallelism to train models that do not fit on one GPU, while balancing communication, memory, and pipeline idle time. It presents flushed and interleaved pipeline schedules, then describes tensor-parallel transformer-layer partitioning.
- Parallelism overview: PTD-P combines pipeline parallelism across multi-GPU servers, tensor parallelism within servers, and data parallelism across model shards.
- Pipeline model parallelism: Pipeline parallelism shards model layers across devices and pipelines microbatches through those stages while preserving consistent weight versions.Strict optimizer semantics require synchronized updates across devices.
- Default schedule: PipeDream-Flush limits in-flight microbatches to the pipeline depth, reducing activation storage from the batch size to p or fewer microbatches.Its bubble time matches the earlier schedule, but it is much more memory-efficient when m≫p.
- Interleaved stages: The interleaved schedule assigns each device multiple model chunks, reducing pipeline bubble time by v while increasing communication by v.The schedule requires the batch microbatch count to be an integer multiple of the pipeline-parallel degree.
- Tensor model parallelism: Tensor parallelism partitions transformer-layer computations across devices, applying GeLU independently and reducing outputs across GPUs before dropout.The partitioning removes synchronization between partitioned GEMMs because GeLU is applied independently.
3 PERFORMANCE ANALYSIS OF PARALLELIZATION CONFIGURATIONS
The section analyzes how pipeline, tensor, and data parallelism trade off memory footprint, device utilization, communication, and pipeline bubbles. It derives configuration heuristics and examines microbatching and activation recomputation.
- Different parallelism dimensions expose trade-offs among memory footprint, device utilization, and communication under fixed GPU and batch-size budgets.
- Tensor and pipeline parallelism have distinct communication patterns: tensor parallelism uses all-reduce, whereas pipeline parallelism uses cheaper point-to-point communication.
- Tensor parallelism should generally remain within a g-GPU server, while pipeline parallelism scales models across servers.
- 1.3× higher per-GPU throughput is observed with a larger microbatch size on a single GPU.
- The optimal microbatch size depends on model throughput and memory characteristics, pipeline depth, data-parallel size, and batch size.
- Activation recomputation trades additional computation for lower activation storage, while checkpoint placement determines memory footprint.
4 IMPLEMENTATION
The implementation combines communication and computation optimizations for PTD-P on GPU clusters. Its scatter-gather strategy reduces redundant cross-node transfers and supports communication-intensive pipeline schedules.
- PTD-P extends Megatron-LM using PyTorch and NCCL, with optimizations targeting both communication and computation.
- Scatter-gather communication splits replicated tensors into chunks before cross-node transfer and reconstructs them at the receiver.
- With tensor-model-parallel size 8, the optimization reduces redundant transfers across consecutive pipeline stages.
- Model-specific kernels change tensor layouts, fuse element-wise operations, and fuse scale, mask, and softmax operations.
5 EVALUATION
The evaluation examines end-to-end training performance, scaling behavior, and throughput across GPT models from billions to one trillion parameters on Selene. It reports throughput, training-time estimates, and configuration details for these experiments.
- The experiments assess PTD-P performance, pipeline scaling, interactions among parallelization dimensions, microbatch size, and communication optimizations.These questions include end-to-end performance, interleaved scheduling, and hardware limits during training.
- The evaluation studies PTD-P and related parallelism configurations across GPT models ranging from 1 billion to 1 trillion parameters.Experiments use mixed precision on Selene, with A100 GPUs connected by NVLink, NVSwitch, and InfiniBand.
- 52% of peak device throughput is achieved for the largest model, compared with 44% for the smallest model, including data loading, optimization, communication, and logging.The results show superlinear scaling to 3072 A100 GPUs as utilization improves for larger models without a significant relative increase in communication time.
- 34 days are estimated for GPT-3 with 175 billion parameters on 1024 A100 GPUs, while 84 days are estimated for a 1 trillion parameter model on 3072 A100 GPUs.The estimates use reported per-GPU throughput and assumed token counts for end-to-end training.
5.2 Comparison to ZeRO-3
The section compares PTD-P with ZeRO-3 without model parallelism for 175- and 530-billion-parameter GPT models. PTD-P scales more gracefully as GPU count increases while maintaining the global batch size.
- 70% higher throughput is achieved by PTD-P than ZeRO-3 for both models when doubling the number of GPUs at the same batch size.The authors attribute this result to reduced cross-node communication.
- With fewer GPUs and microbatch size 4, PTD-P provides 6% higher throughput for the 175-billion-parameter model and 24% for the 530-billion-parameter model.The comparison keeps global batch size fixed as GPU count changes.
5.3 Pipeline Parallelism
The evaluation studies pipeline scaling, interleaved scheduling, and combinations of tensor, pipeline, and data parallelism. Results emphasize trade-offs between communication cost, pipeline bubbles, batch size, and model fit.
- Pipeline Parallelism: Higher batch sizes scale better in weak-scaling pipeline experiments because the pipeline bubble is amortized over more microbatches.Model size increases proportionally with the number of pipeline stages while tensor-parallel size remains 8.
- Pipeline Parallelism: The interleaved schedule with scatter/gather optimization achieves higher computational performance than the default non-interleaved schedule for a 175-billion-parameter GPT model.The performance gap narrows as batch size increases because the default schedule’s bubble decreases and communication is amortized differently.
- Tensor versus Pipeline Parallelism: Tensor and pipeline model parallelism together provide the strongest configuration for a 161-billion-parameter GPT model, with tensor parallelism within nodes and pipeline parallelism across nodes.Tensor parallelism uses expensive all-reduce communication, whereas pipeline parallelism uses cheaper point-to-point communication but incurs pipeline bubbles.
- Tensor versus Pipeline Parallelism: Peak performance occurs when tensor-parallel size equals the 8 GPUs in a DGX A100 node, while pipeline stages remain limited relative to the number of microbatches.Neither tensor nor pipeline model parallelism alone matches their combined performance in this experiment.
- Tensor versus Data Parallelism: Data parallelism alone cannot support very large models with limited batch sizes because of insufficient memory capacity and scaling limitations.GPT-3’s convergence batch size of 1536 would support data parallelism to only 1536 GPUs, despite roughly 10,000 GPUs being used for reasonable training time.
5.5 Microbatch Size
The section evaluates microbatch size and related pipeline optimizations for large GPT models. Microbatch size affects throughput through arithmetic intensity and the number of microbatches in the pipeline.
- The best microbatch size is 2 for the 91-billion-parameter model using tensor-parallel size 8 and pipeline-parallel size 8.The optimal microbatch size differs across models and is model-dependent.
- Increasing microbatch size reduces the number of microbatches in the pipeline for a fixed batch size, changing pipeline behavior and throughput.The passage links microbatch size to both arithmetic intensity and pipeline bubble size.
- Activation recomputation and scatter/gather communication are evaluated as throughput-affecting optimizations for large GPT models.The cited figures compare throughput with and without activation recomputation and with and without scatter/gather optimization.
5.6 Activation Recomputation
The section reports performance effects of activation recomputation, scatter/gather communication optimization, and operator fusion, alongside communication and checkpointing measurements at large scale.
- Up to 33% lower throughput occurs with activation recomputation at small batch sizes for the 145-billion-parameter GPT model.The cost comes from an extra forward pass during backward propagation; recomputation is needed for larger models.
- Up to 11% higher throughput is achieved by scatter/gather communication optimization for communication-intensive GPT-3 schedules.The optimization reduces communication over cross-node links, especially with large batch sizes and interleaving.
- 19% higher throughput is obtained from operator fusion for GPT-3, increasing per-GPU throughput from 113 teraFLOP/s to 135 teraFLOP/s.
- 11% higher throughput is obtained from operator fusion for the 530-billion-parameter GPT model, increasing per-GPU throughput from 133 teraFLOP/s to 148 teraFLOP/s.
- The trillion-parameter model reaches 892 GB/s effective point-to-point bisection bandwidth and 12.9 TB/s effective data-parallel all-reduce bisection bandwidth on 3072 GPUs.
- The trillion-parameter model uses a 13.8-terabyte checkpoint, with initial loading reaching 1TB/s read bandwidth and saves reaching 273 GB/s.
6 RELATED WORK
Related work covers pipeline, tensor, data, sharded, automatic, and high-performance-computing approaches for large-scale model training, emphasizing differing scalability and semantic trade-offs.
- Parallelism for Large Models: Pipeline parallelism variants differ in granularity, elasticity, hardware heterogeneity, and whether they preserve strict or relaxed optimizer semantics.
- Parallelism for Large Models: Relaxed-semantic pipeline methods can improve throughput compared with pipeline-flush techniques, potentially at the cost of convergence rate or final accuracy.
- Parallelism for Large Models: DeepSpeed combines pipeline, tensor, and data parallelism to train models with up to a trillion parameters, but reports 36% rather than 52% of peak throughput.
- Sharded Data Parallelism: Sharded data parallelism divides optimizer computation and memory across data-parallel partitions without adding communication over vanilla data parallelism.
- Automatic Partitioning: Automatic partitioning methods use cost models but do not consider all parallelism dimensions and memory-saving effects for models exceeding accelerator memory capacity.
- HPC for Model Training: HPC training demonstrations differ because their image-classification models fit on one accelerator, use very large batches, and do not require model parallelism.
7 DISCUSSION AND CONCLUSION
The discussion concludes that composing pipeline, tensor, and data parallelism enables high-throughput trillion-parameter training, while careful interaction-aware design and accelerator-agnostic optimizations remain important.
- 502 petaFLOP/s aggregate throughput is achieved while training trillion-parameter models with PTD-P.The approach enables an estimated end-to-end training time of around 3 months.
- The effectiveness of combined parallelism depends on carefully considering the trade-offs and interactions among pipeline, tensor, and data parallelism.
- Smart graph partitioning, memory-bound-kernel reduction, operator fusion, careful data layout, and domain-specific communication optimizations are identified as accelerator-agnostic ideas.
APPENDIX: FLOATING-POINT OPERATIONS
The appendix defines the variables and component operations used to calculate floating-point operations for transformer language models, including matrix multiplication, attention, feed-forward, and logit computations.
- The FLOP calculation uses transformer layers l, hidden size h, sequence length s, vocabulary size V, and training batch size B.
- A m×k by k×n matrix multiplication requires 2m×k×n FLOPs to account for multiplications and additions.
- A transformer layer combines an attention block with a 2-layer feed-forward network whose hidden dimension expands to 4h before returning to h.
- The language-model logit layer contributes 6BshV FLOPs in total, combining 2BshV forward-pass FLOPs and 4BshV backward-pass FLOPs.
- The appendix then presents the total floating-point-operation expression for a transformer model with l transformer layers.