Parallel Training of Pre-Trained Models via Chunk-Based Dynamic Memory Management

By Jiarui Fang, Yang Yu, Shenggui Li, Yang You, Jie Zhou. 2021

48 citations recorded by the source corpus.

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. 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 exhibits significantly higher execution speed. PatricStar also successfully runs the 175B GPT3 training task on a 32 GPU cluster. Our code is available at https://github.com/Tencent/PatrickStar.

Subfields: Computer Networks and Communications, Computer Vision and Pattern Recognition, Hardware and Architecture, Efficiency & Compression

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Computer Networks and Communications

1. Chien-chin Huang
2. Myle Ott
3. Sam Shleifer

In Computer Vision and Pattern Recognition

1. Sam Shleifer
2. Chien-chin Huang
3. Min Xu
4. Alban Desmaison
5. Rohan Varma
6. Myle Ott
7. Liangchen Luo
8. Luke S. Zettlemoyer

In Hardware and Architecture

1. Chien-chin Huang
2. Myle Ott
3. Sam Shleifer

In Efficiency & Compression

1. Chien-chin Huang
2. Myle Ott
3. Sam Shleifer

Citations in the corpus

• PyTorch FSDP: Experiences on Scaling Fully Sharded Data Parallel (2023, background): “Orthogonal memory-saving techniques include gradient compression [2], mixed-precision training [7], tensor rematerialization [12] and CPU-offloading [4], but they could have implications on model accuracy and incur overhead in (un)compression, quantization, recomputation, and host-to-device copies, ”