通过时空规划提升大模型训练内存效率,减少碎片化浪费。
STAlloc: Enhancing Memory Efficiency in Large-Scale Model Training with Spatio-Temporal Planning
- 离线规划结合在线分配,利用训练内存行为的时空规律优化分配。
- 平均降低85.1%内存碎片率(最高达100%),最多节省43%内存。
- 兼容MoE等复杂模型,提升训练吞吐量最高达32.5%,开销极小。
大规模语言模型的快速扩展显著增加了GPU内存压力,而虚拟流水线和梯度重计算等训练优化技术破坏了张量生命周期,导致严重内存碎片。现有框架如PyTorch采用的在线内存分配器忽视张量生命周期,造成高达43%的内存浪费并引发显存溢出。为此,我们提出STAlloc,一种基于时空规律的GPU内存分配器。其结合离线规划与在线分配:离线部分利用训练工作负载的时空规律生成近最优分配方案,线上部分处理Mixture-of-Experts(MoE)等动态复杂模型。作为可插拔的PyTorch分配器,STAlloc在密集和MoE模型上平均降低85.1%(最高100%)的碎片率,实现更高吞吐量训练配置,吞吐性能最高提升32.5%,且开销可忽略不计。
原文摘要 · Abstract (English)
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%.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。