TLDR: STWeaver is a new GPU memory allocator for deep learning that significantly reduces memory fragmentation in large model training. It achieves this by using an “offline planning” phase to identify predictable memory usage patterns and then an “online allocation” phase that reuses idle memory. This approach saves substantial GPU memory (up to 56.3GB), prevents out-of-memory errors, and enables more efficient training configurations, leading to up to 32.5% throughput improvement with negligible overhead.
Training large language models (LLMs) has become a cornerstone of modern artificial intelligence, enabling breakthroughs in various fields from language comprehension to code generation. However, this rapid growth comes with a significant challenge: managing the immense GPU memory demands. Traditional GPU memory allocators, like those in popular deep learning frameworks such as PyTorch, often struggle with these demands, leading to a problem known as memory fragmentation. This fragmentation can waste a substantial portion of memory, sometimes up to 43%, and frequently causes “out-of-memory” errors, which can severely hinder the effectiveness of training optimization techniques.
Memory fragmentation occurs because standard allocators use an “online” strategy, meaning they allocate memory as requests come in, without prior knowledge of how long a piece of memory will be needed. When tensors (the data structures used in deep learning) are deallocated in an unpredictable order, it leaves behind small, scattered gaps in memory. Over time, these gaps become too small or too fragmented to accommodate new, larger requests, even if there’s theoretically enough free memory overall. This inefficiency can force developers to use less optimal training configurations, slowing down the entire process.
To tackle this pervasive issue, researchers have introduced STWeaver, a novel GPU memory allocator designed specifically for deep learning frameworks. STWeaver’s core innovation lies in its ability to exploit the predictable patterns, or “spatio-temporal regularities,” in how memory is allocated during training. Unlike traditional online allocators, STWeaver combines an “offline planning” phase with “online allocation.”
How STWeaver Works
STWeaver operates in three main stages: profiling, plan synthesis, and runtime allocation.
First, an Allocation Profiler observes a few training iterations to understand the memory allocation and deallocation patterns. It captures crucial information like the size of memory requests, their lifespan (how long they are needed), and whether they are static (predictable size and lifespan) or dynamic (unpredictable size, like in Mixture-of-Experts models).
Next, the Plan Synthesizer takes this profiled information and creates a near-optimal memory allocation plan. It does this by grouping memory requests based on their spatial (size) and temporal (lifespan) regularities. For static requests, it generates a fixed plan that minimizes fragmentation. For dynamic requests, it identifies “Dynamic Reusable Space” – idle areas within the static plan that can be safely reused by dynamic allocations at runtime without causing conflicts.
Finally, the Runtime Allocator executes this plan during actual training. It has a “Static Allocator” that follows the pre-planned addresses for static requests, eliminating the need for real-time searches. A “Dynamic Allocator” handles dynamic requests, prioritizing the identified reusable spaces. If a request cannot be accommodated by the plan or reusable space, it falls back to a standard caching allocator, ensuring robustness.
Also Read:
- LaCache: A Smart Memory Solution for Long-Context LLMs
- Unlocking GNN Potential on 100-Billion Edge Graphs
Key Benefits and Performance
STWeaver has demonstrated remarkable improvements in memory efficiency. Across various dense and sparse models, and different training configurations, it reduces memory fragmentation by an average of 79.2%, and in some cases, up to 100%. This translates to significant GPU memory savings, up to 56.3GB in some tests. By freeing up this memory, STWeaver enables the use of more efficient, high-throughput training configurations that would otherwise lead to out-of-memory errors. This can result in performance improvements of up to 32.5%.
Crucially, STWeaver achieves these benefits with negligible impact on end-to-end training throughput. The profiling and plan synthesis, while taking a few minutes, are one-time overheads that are quickly offset by the long-term training efficiencies. The system is also designed as a pluggable allocator for PyTorch, making it compatible with existing training frameworks like Megatron-LM, Megatron-DeepSpeed, and Colossal-AI.
This innovative approach of combining proactive, ahead-of-time planning with reactive runtime decision-making represents a significant step forward in optimizing GPU memory utilization for large-scale model training. For more technical details, you can refer to the full research paper.
