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

arXiv: https://arxiv.org/abs/2507.16274 · HTML: https://mlsanity.com/paper/2507.16274/

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Placed in Efficient MoE training and inference systems, one of 12 directions in Mixture-of-experts.
- Builds on 22 of its 50 references that are in the map.
- First posted on arXiv 22 Jul 2025.

## Placed in

- Mixture-of-experts / [Efficient MoE training and inference systems](https://mlsanity.com/moe/efficient-moe-training-and-inference-systems/)

## Builds on (most cited first, 22 of 22)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Qwen Technical Report](https://mlsanity.com/paper/2309.16609/) (2023)
- [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://mlsanity.com/paper/1909.08053/) (2019)
- [GShard: Scaling Giant Models with Conditional Computation and Automatic Sharding](https://mlsanity.com/paper/2006.16668/) (2020)
- [ZeRO: Memory Optimizations Toward Training Trillion Parameter Models](https://mlsanity.com/paper/1910.02054/) (2019)
- [DeepSeek-V2: A Strong, Economical, and Efficient Mixture-of-Experts Language Model](https://mlsanity.com/paper/2405.04434/) (2024)
- [Training Deep Nets with Sublinear Memory Cost](https://mlsanity.com/paper/1604.06174/) (2016)
- [Efficient Large-Scale Language Model Training on GPU Clusters Using Megatron-LM](https://mlsanity.com/paper/2104.04473/) (2021)
- [ZeRO-Offload: Democratizing Billion-Scale Model Training](https://mlsanity.com/paper/2101.06840/) (2021)
- [Reducing Activation Recomputation in Large Transformer Models](https://mlsanity.com/paper/2205.05198/) (2022)
- [ZeRO-Infinity: Breaking the GPU Memory Wall for Extreme Scale Deep Learning](https://mlsanity.com/paper/2104.07857/) (2021)
- [TensorFlow Lite Micro: Embedded Machine Learning on TinyML Systems](https://mlsanity.com/paper/2010.08678/) (2020)
- [Memory-Efficient Pipeline-Parallel DNN Training](https://mlsanity.com/paper/2006.09503/) (2020)
- [Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training](https://mlsanity.com/paper/2110.14883/) (2021)
- [Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization](https://mlsanity.com/paper/1910.02653/) (2019)
- [Efficient Training of Large Language Models on Distributed Infrastructures: A Survey](https://mlsanity.com/paper/2407.20018/) (2024)
- [GMLake: Efficient and Transparent GPU Memory Defragmentation for Large-scale DNN Training with Virtual Memory Stitching](https://mlsanity.com/paper/2401.08156/) (2024)
- [TFLMS: Large Model Support in TensorFlow by Graph Rewriting](https://mlsanity.com/paper/1807.02037/) (2018)
- [Fast, Multicore-Scalable, Low-Fragmentation Memory Allocation through Large Virtual Memory and Global Data Structures](https://mlsanity.com/paper/1503.09006/) (2015)

## Built on by (most cited first, 2 of 2)

- [RLinf: Flexible and Efficient Large-scale Reinforcement Learning via Macro-to-Micro Flow Transformation](https://mlsanity.com/paper/2509.15965/) (2025)
- [HyperOffload: Graph-Driven Hierarchical Memory Management for Large Language Models on SuperNode Architectures](https://mlsanity.com/paper/2602.00748/) (2026)

## Abstract (arXiv)

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