STAlloc: Enhancing Memory Efficiency in Large-Scale Model Training with Spatio-Temporal Planning

TL;DR

STAlloc is a novel GPU memory allocator that significantly reduces memory fragmentation in large-scale model training by combining offline planning with online allocation, leading to improved efficiency and throughput.

Contribution

It introduces a spatio-temporal planning-based memory allocator for deep learning frameworks that minimizes fragmentation and enhances training performance.

Findings

Reduces memory fragmentation ratio by up to 100%.

Improves training throughput by up to 32.5%.

Effective for both dense and Mixture-of-Experts models.

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. 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%.