PatrickStar: Parallel Training of Pre-trained Models via Chunk-based Memory Management
arXiv:2108.05818 · doi:10.1109/TPDS.2022.3219819
Abstract
The pre-trained model (PTM) is revolutionizing Artificial Intelligence (AI) technology. However, the hardware requirement of PTM training is prohibitively high, making it a game for a small proportion of people. Therefore, we proposed PatrickStar system to lower the hardware requirements of PTMs and make them accessible to everyone. PatrickStar uses the CPU-GPU heterogeneous memory space to store the model data. Different from existing works, we organize the model data in memory chunks and dynamically distribute them in the heterogeneous memory. Guided by the runtime memory statistics collected in a warm-up iteration, chunks are orchestrated efficiently in heterogeneous memory and generate lower CPU-GPU data transmission volume and higher bandwidth utilization. Symbiosis with the Zero Redundancy Optimizer, PatrickStar scales to multiple GPUs on multiple nodes. % using data parallelism. The system can train tasks on bigger models and larger batch sizes, which cannot be accomplished by existing works. Experimental results show that PatrickStar extends model scales 2.27 and 2.5 times of DeepSpeed, and consistently exhibits significantly higher execution speed. PatricStar also successfully runs the 175B GPT3 training task on a 32 GPU cluster. Our code is publicly available at https://github.com/Tencent/PatrickStar.
arXiv comment13 pages
References in corpus (17)
- Adam: A Method for Stochastic Optimization
- An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale
- PyTorch: An Imperative Style, High-Performance Deep Learning Library
- Language Models are Few-Shot Learners
- Mixed Precision Training
- Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism
- Training Deep Nets with Sublinear Memory Cost
- Evaluating Modern GPU Interconnect: PCIe, NVLink, NV-SLI, NVSwitch and GPUDirect
- On-Device Neural Net Inference with Mobile GPUs
- ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning
- Mesh-TensorFlow: Deep Learning for Supercomputers
- Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization
- TeraPipe: Token-Level Pipeline Parallelism for Training Large-Scale Language Models
- Training Large Neural Networks with Constant Memory using a New Execution Algorithm
- Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training
- Maximizing Parallelism in Distributed Training for Huge Neural Networks
- Optimal checkpointing for heterogeneous chains: how to train deep neural networks with limited memory