STAlloc: Enhancing Memory Efficiency in Large-Scale Model Training with Spatio-Temporal Planning
Zixiao Huanga, Junhao Hub, Hao Lina, Chunyang Zhuc, Yueran Tangb, Quanlu Zhangc, Zhen Guoc, Zhenhua Lia, Shengen Yana, Zhenhua Zhua, Guohao Daid, Yu Wanga
- aTsinghua University, Beijing, China
- bInfinigence-AI, Shanghai, China
- cInfinigence-AI, Beijing, China
- dShanghai Jiao Tong University, Shanghai, China
Papers open inside SciRah from the original publisher or repository; SciRah never hosts them.
Abstract
The rapid scaling of large language models (LLMs) has significantly increased GPU memory pressure, which is further aggravated by training optimization techniques such as virtual pipeline and recomputation that disrupt tensor lifespans and introduce considerable memory fragmentation. Such fragmentation stems from the use of online GPU memory allocators in popular deep learning frameworks like PyTorch, which disregard tensor lifespans. As a result, this inefficiency can waste as much as 43% of memory and trigger out-of-memory errors, undermining the effectiveness of optimization methods.