STAlloc: Enhancing Memory Efficiency in Large-Scale Model Training with Spatio-Temporal Planning
Fuente:
arXiv
Saved in:
| Main Authors: | , , , , , , , , , , , |
|---|---|
| Format: | Preprint |
| Published: |
2025
|
| Subjects: | |
| Online Access: | |
| Tags: |
Add Tag
No Tags, Be the first to tag this record!
|
| _version_ | 1866911284189986816 |
|---|---|
| author | Huang, Zixiao Hu, Junhao Lin, Hao Zhu, Chunyang Tang, Yueran Zhang, Quanlu Guo, Zhen Li, Zhenhua Yan, Shengen Zhu, Zhenhua Dai, Guohao Wang, Yu |
| author_facet | Huang, Zixiao Hu, Junhao Lin, Hao Zhu, Chunyang Tang, Yueran Zhang, Quanlu Guo, Zhen Li, Zhenhua Yan, Shengen Zhu, Zhenhua Dai, Guohao Wang, Yu |
| contents | 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. To address this, we introduce STAlloc, a GPU memory allocator for deep learning frameworks that reduces fragmentation by exploiting the spatial and temporal regularity in memory allocation behaviors of training workloads. STAlloc introduces a novel paradigm that combines offline planning with online allocation. The offline planning leverages spatio-temporal regularities to generate a near-optimal allocation plan, while the online allocation handles complex and dynamic models such as Mixture-of-Experts (MoE). Built as a pluggable PyTorch memory allocator, STAlloc reduces fragmentation ratio on average by 85.1% (up to 100%) across both dense and MoE models, with negligible overhead. This enables more efficient, high-throughput training configurations and improves throughput performance by up to 32.5%. |
| format | Preprint |
| id |
arxiv_https___arxiv_org_abs_2507_16274 |
| institution | arXiv |
| publishDate | 2025 |
| record_format | arxiv |
| spellingShingle | STAlloc: Enhancing Memory Efficiency in Large-Scale Model Training with Spatio-Temporal Planning Huang, Zixiao Hu, Junhao Lin, Hao Zhu, Chunyang Tang, Yueran Zhang, Quanlu Guo, Zhen Li, Zhenhua Yan, Shengen Zhu, Zhenhua Dai, Guohao Wang, Yu Machine Learning Artificial Intelligence Distributed, Parallel, and Cluster Computing Performance 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. To address this, we introduce STAlloc, a GPU memory allocator for deep learning frameworks that reduces fragmentation by exploiting the spatial and temporal regularity in memory allocation behaviors of training workloads. STAlloc introduces a novel paradigm that combines offline planning with online allocation. The offline planning leverages spatio-temporal regularities to generate a near-optimal allocation plan, while the online allocation handles complex and dynamic models such as Mixture-of-Experts (MoE). Built as a pluggable PyTorch memory allocator, STAlloc reduces fragmentation ratio on average by 85.1% (up to 100%) across both dense and MoE models, with negligible overhead. This enables more efficient, high-throughput training configurations and improves throughput performance by up to 32.5%. |
| title | STAlloc: Enhancing Memory Efficiency in Large-Scale Model Training with Spatio-Temporal Planning |
| topic | Machine Learning Artificial Intelligence Distributed, Parallel, and Cluster Computing Performance |
| url | https://arxiv.org/abs/2507.16274 |