An Efficient Checkpointing System for Large Machine Learning Model Training

Conference · Sat Nov 30 04:00:00 EST 2024

DOI:https://doi.org/10.1109/SCW63240.2024.00127· OSTI ID:2526231

Xu, Wubiao ^[1]; Huang, Xin ^[2]; Meng, Shiman ^[1]; Zhang, Weiping ^[1]; Guo, Luanzheng ^[3]; Sato, Kento ^[4]

---

1. Nanchang Hangkong University
2. Kobe University
3. BATTELLE (PACIFIC NW LAB)
4. RIKEN R-CCS

As machine learning models increase in size and complexity rapidly, the cost of checkpointing in ML training became a bottleneck in storage and performance (time). For example, the latest GPT-4 model has massive parameters at the scale of 1.76 trillion. It is highly time and storage consuming to frequently writes the model to checkpoints with more than 1 trillion floating point values to storage. This work aims to understand and attempt to mitigate this problem. First, we characterize the checkpointing interface in a collection of representative large machine learning/language models with respect to storage consumption and performance overhead. Second, we propose the two optimizations: i) A periodic cleaning strategy that periodically cleans up outdated checkpoints to reduce the storage burden; ii) A data staging optimization that coordinates checkpoints between local and shared file systems for performance improvement.

• Research Organization: Pacific Northwest National Laboratory (PNNL), Richland, WA (United States)
• Sponsoring Organization: USDOE
• DOE Contract Number: AC05-76RL01830
• OSTI ID: 2526231
• Report Number(s): PNNL-SA-204755
• Country of Publication: United States
• Language: English

Related Subjects

machine learning
artificial intelligence