xMem: A CPU-Based Approach for Accurate Estimation of GPU Memory in Deep Learning Training Workloads

TL;DR

The paper tackles the challenge of accurately predicting peak GPU memory usage for DL training in shared clusters without using GPUs during estimation. It introduces xMem, a CPU-based dynamic analysis framework with three components—Analyzer, Memory Orchestrator, and Simulator—that simulate a two-level allocator on the GPU to forecast peak memory from CPU traces. Across 25 models and thousands of runs, xMem achieves substantially lower Median Relative Error ($MRE$) and lower Probability of Estimation Failure ($PEF$) than static, data-driven, or direct-GPU baselines, while enabling significant memory conservation ($MCP$). The approach demonstrates strong scalability to large Transformer models and high-end GPUs, and is released as open-source, offering practical benefits for proactive scheduling and OOM avoidance in GPU-constrained environments.

Abstract

The global scarcity of GPUs necessitates more sophisticated strategies for Deep Learning jobs in shared cluster environments. Accurate estimation of how much GPU memory a job will require is fundamental to enabling advanced scheduling and GPU sharing, which helps prevent out-of-memory (OOM) errors and resource underutilization. However, existing estimation methods have limitations. Approaches relying on static analysis or historical data with machine learning often fail to accurately capture runtime dynamics. Furthermore, direct GPU analysis consumes scarce resources, and some techniques require intrusive code modifications. Thus, the key challenge lies in precisely estimating dynamic memory requirements, including memory allocator nuances, without consuming GPU resources and non-intrusive code changes. To address this challenge, we propose xMem, a novel framework that leverages CPU-only dynamic analysis to accurately estimate peak GPU memory requirements a priori. We conducted a thorough evaluation of xMem against state-of-the-art solutions using workloads from 25 different models, including architectures like Convolutional Neural Networks and Transformers. The analysis of 5209 runs, which includes ANOVA and Monte Carlo results, highlights xMem's benefits: it decreases the median relative error by 91% and significantly reduces the probability of estimation failure as safe OOM thresholds by 75%, meaning that the estimated value can often be used directly without causing OOM. Ultimately, these improvements lead to a 368% increase in memory conservation potential over current solutions.

Figures (9)

• Figure 1: The impact of optimizer.zero_grad() placement on GPU memory between calling zero_grad() before backward propagation (POS0) versus at the start of iteration (POS1), illustrating runtime sensitivity and allocator impact. While the underlying Tensor memory activity (green/red areas for POS0/POS1) might show similar patterns, the overall Segment memory footprint (solid and dashed lines) differs significantly.
• Figure 2: An example of how memory blocks are managed in PyTorch for a training job.
• Figure 3: The impact of the sequence of memory activities on GPU memory consumption.
• Figure 4: Architecture of the xMem CPU-based memory estimation framework. The pipeline (center) uses CPU profiling data to estimate peak GPU memory via analysis, orchestrator, and simulation.
• Figure 5: A visualization of the time-based hierarchical attribution in xMem, correlating Python function calls, CPU operator execution time frames, and associated memory blocks.