STWeaver: Efficient GPU Memory Allocation

• STWeaver is a GPU memory allocator that uses spatio-temporal planning to minimize memory fragmentation in deep learning training.
• It integrates offline static planning with dynamic runtime allocation to efficiently manage tensor lifetimes and allocation sizes.
• Evaluations demonstrate up to 100% fragmentation reduction, saving up to 56.3GB GPU memory and boosting throughput by 32.5%.

STWeaver is a GPU memory allocator for deep learning frameworks that reduces memory fragmentation during large-scale model training by combining offline allocation planning with online runtime allocation. Implemented as a pluggable PyTorch allocator, it is designed for training regimes in which virtual pipeline parallelism, recomputation, ZeRO, offloading, and Mixture-of-Experts (MoE) disrupt tensor lifespans and make default online caching allocators waste substantial GPU memory. Its central premise is that training iterations exhibit strong regularity in both tensor sizes and tensor lifetimes, making most allocations predictable enough to pre-plan, while dynamic cases can still be handled safely at runtime. Across the evaluated dense and sparse workloads, STWeaver reduces fragmentation ratio by 79.2% on average and by up to 100%, saves up to 56.3 GB of GPU memory, and can improve training throughput by up to 32.5% (Huang et al., 22 Jul 2025).

1. Problem setting and formal objective

STWeaver is motivated by the gap between theoretical tensor footprint and actual reserved GPU memory in large-model training. The paper defines memory efficiency as

$E = \frac{M_a}{M_r},$

where $M_a$ is the size of allocated memory under the current training configuration and $M_r$ is the total memory reserved by the allocator. The fragmentation ratio is therefore $1-E$. In this formulation, fragmentation is the proportion of reserved memory that is not actually utilized by tensors (Huang et al., 22 Jul 2025).

The underlying workloads are characterized by very large request counts, varied tensor lifetimes, and complex interleaving of allocation and deallocation. The number of memory requests can exceed $10^5$ in large-scale training. Persistent tensors such as weights and optimizer states coexist with activations, operator intermediates, and tensors whose lifetimes are reshaped by recomputation, offloading, and virtual pipeline parallelism. The paper’s GPT-2 example on 8 A800 GPUs illustrates the effect quantitatively: under baseline 1F1B pipeline parallelism, PyTorch memory efficiency is around 90% with 52.1 GB reserved; enabling Virtual Pipeline Parallelism raises allocated memory to 51.8 GB, but memory efficiency drops to 80%, increasing reserved memory to 59.9 GB; recomputation reduces allocated memory but can further reduce efficiency to around 60%.

The critique of default framework allocators is specifically temporal. PyTorch’s allocator is described as an online caching allocator using large reserved blocks and a best-fit policy. Because it allocates without knowledge of future lifespans, it can make locally sensible decisions that produce globally poor reuse patterns. The paper therefore recasts allocation planning as an optimization problem over both address layout and tensor lifetimes. Each memory request event is modeled as

$m := (s, t_s, t_e, p_s, p_e, d),$

where $s$ is request size, $t_s$ and $t_e$ are allocation and free timestamps, $p_s$ and $M_a$0 are allocation and free computation phases, and $M_a$1 records whether the request is dynamic. For dynamic requests, the profiler additionally records $M_a$2 and $M_a$3, the module names of allocation and free. Since the allocated memory $M_a$4 is fixed for a given training configuration, maximizing $M_a$5 reduces to minimizing $M_a$6. The paper explicitly notes that optimal planning is NP-hard as an instance of the Dynamic Storage Allocation problem.

2. Spatio-temporal planning and allocation synthesis

The core idea of STWeaver is “spatio-temporal planning”: planning jointly over space, meaning tensor sizes and address layout, and time, meaning allocation and free times. The method relies on the observation that large-model training is “mostly repetitive and highly regular” across iterations. Spatial regularity is strong because repeated Transformer blocks produce repeated allocation sizes; the paper’s Llama2-7B example reports that among more than 50,000 allocations larger than 512 bytes in one iteration, there are only about 32 distinct tensor sizes. Temporal regularity is expressed through three recurring classes of objects: persistent tensors, scoped tensors, and transient tensors (Huang et al., 22 Jul 2025).

STWeaver exploits this structure through three components: Allocation Profiler, Plan Synthesizer, and Runtime Allocator. The profiler traces torch-level allocation and free requests and records timestamp, address, size, dynamicity, current computation phase, micro-batch ID, and module name. During profiling it uses native GPU allocation APIs such as cudaMalloc and cudaFree rather than PyTorch’s caching allocator, because native APIs avoid framework-level fragmentation and permit profiling even when the same configuration fails under PyTorch due to fragmentation.

The offline planner first partitions requests into static and dynamic subsets,

$M_a$7

Static requests are then decomposed by temporal and spatial regularity. A HomoPhase Group contains requests sharing the same start and end computation phases,

$M_a$8

These groups are fused when doing so improves the time-memory product,

$M_a$9

which measures how efficiently reserved memory is occupied over time. The paper interprets low $M_r$0 as evidence of “spatio-temporal bubbles.”

After temporal grouping and fusion, STWeaver performs HomoSize Group planning. A HomoSize Group contains requests with identical size $M_r$1, and the planner interprets each reusable address region as a memory layer. The objective is to minimize the number of layers, so that equal-size requests whose lifetimes do not overlap can share one region. The resulting global plan is assembled by partitioning into HomoPhase Groups, fusing them into local temporal plans, treating those plans as larger requests, regrouping by size, constructing memory layers for each size group, processing groups in descending order of request size, and finally assigning addresses for all static requests. The paper characterizes the outcome as “near-optimal”: it does not solve the NP-hard global problem exactly, but combines local optima through grouping and greedy heuristics.

A second planning stage addresses dynamic requests. Static and dynamic peak memory usage often do not occur simultaneously, so STWeaver searches the static plan for safe idle intervals that dynamic requests can reuse. For a dynamic group $M_r$2 over temporal range $M_r$3, the occupied addresses in the static plan are

$M_r$4

and the dynamic reusable intervals are

$M_r$5

where $M_r$6 is the full static-pool address space. This mechanism is the paper’s main answer to MoE and other runtime-variable layers: exact addresses are not pre-assigned, but reusable regions are pre-vetted.

3. Runtime allocator, PyTorch integration, and deployment model

At runtime, STWeaver consists of a Static Allocator, a Dynamic Allocator, and a caching allocator fallback, with a Request Matcher that routes requests according to whether they are static or dynamic and according to current model-layer context (Huang et al., 22 Jul 2025).

Before training, STWeaver reserves one contiguous static memory pool whose size is determined by the Static Allocation Plan. Static requests are then served by returning the pre-planned addresses in sequence, eliminating online search for the dominant fraction of requests. Dynamic requests are handled inside the same pool but only within precomputed safe intervals. If $M_r$7 denotes the currently free intervals in the static pool and $M_r$8 the reusable space for the request’s HomoLayer Group, then the valid candidates for a dynamic request $M_r$9 are

$1-E$0

Among these candidates, STWeaver uses a best-fit policy. If no candidate interval can satisfy the request, it falls back to the caching allocator. Correctness and robustness therefore do not depend on a perfect static prediction.

The implementation is intentionally lightweight from the framework perspective. STWeaver is implemented using about 5700 lines of Python and 3700 lines of C++. The plan synthesizer is a standalone tool, while the profiler and runtime allocator use PyTorch’s PluggableAllocator interface. The allocator can be loaded before training to intercept malloc and free calls, does not require recompiling PyTorch, is compatible with any PyTorch version that supports PluggableAllocator, and works on both NVIDIA and AMD GPU platforms as long as the interface is supported. To capture phase and layer information, STWeaver uses monkey patching and PyTorch hook APIs, requiring no more than five lines of code in the original training framework.

The deployment model assumes a short profiling phase before production training. The paper states that profiling typically requires three iterations, and elsewhere notes that at most five iterations are sufficient because patterns are repetitive. Runtime overhead is negligible: for identical training configurations, STWeaver differs from vanilla PyTorch by <0.05% in all reported throughput comparisons. The scope is also explicit: as a PyTorch pluggable allocator, it manages PyTorch-related GPU allocations but not allocations originating directly from kernels or compilers, although the managed memory typically accounts for over 95% of total GPU memory usage during training.

4. Evaluation across dense, sparse, and multi-platform workloads

The evaluation spans over 48 training configurations on 3 testbeds, using Megatron-LM, Megatron-DeepSpeed, and Colossal-AI, with dense models including GPT-2 (345M), Llama2-7B, and Qwen2.5 at 7B, 14B, 32B, and 72B, and sparse models including Qwen1.5-MoE-A2.7B (16B total parameters). Hardware includes 1 node, 8× NVIDIA A800-80GB, up to 16 nodes, each with 8× NVIDIA H200-141GB, and 8 nodes, each with 8× AMD MI210-64GB. Baselines are the PyTorch caching allocator, PyTorch expandable_segments, and GMLake (Huang et al., 22 Jul 2025).

The headline result is a 79.2% average reduction in fragmentation ratio, with up to 100% reduction, up to 56.3 GB memory savings, and up to 32.5% throughput improvement. For dense models, STWeaver achieves >95% memory efficiency, and in some cases 100%, across all tested configurations. The reported baseline ranges are 57.2% to 93.8% for PyTorch, 53.6% to 96.5% for GMLake, and 62.4% to 94.9% for PyTorch ES. Relative to these baselines, STWeaver reduces fragmentation memory on average by 88.1% versus PyTorch, 77.1% versus GMLake, and 76.0% versus PyTorch ES, with reserved-memory reduction up to 11 GB on A800 experiments. The largest dense-model gain is reported for GPT-2 with ZeRO-3 and recomputation.

For sparse MoE workloads, the allocator still achieves 93.8% to 97.8% memory efficiency, corresponding to an average fragmentation ratio of 4.3%. The baseline fragmentation ratios on MoE are 17.7% for PyTorch, 20.3% for GMLake, and 6.9% for PyTorch ES. STWeaver reduces fragmentation memory by 77.1% versus PyTorch, 78.3% versus GMLake, and 39.4% versus PyTorch ES. The paper further decomposes the gain: STWeaver with only the Static Allocation Plan already reduces fragmentation memory by 70.2% versus PyTorch, accounting for 91% of the total fragmentation reduction achieved by the full system; static allocations comprise 73.4% to 99.3% of total memory depending on configuration; dynamic reuse adds an additional 22.9% reduction compared with STWeaver without dynamic reuse.

Cross-platform and scaling results are similarly strong. On AMD, using recomputation, STWeaver achieves over 90% memory efficiency and up to 99.7%, while PyTorch stays below 60% memory efficiency across all scales for Llama2-7B; STWeaver reduces fragmentation memory by 22.8 GB, which is 35.6% of GPU memory. On H200 with larger Qwen2.5 models under recomputation, STWeaver achieves 99.1% memory efficiency, with fragmentation reduction exceeding 98.5% versus PyTorch 2.6 and 98.4% versus PyTorch ES, average memory savings of 37.9 GB, and maximum savings of 56.3 GB. Under virtual pipeline, it achieves >99% memory efficiency in all cases, with average savings of 15.7 GB, while baseline efficiency degrades as model and cluster scale increase and STWeaver remains stable within 0.7%.

The most operationally consequential result is OOM avoidance. In the Qwen2.5-14B on 16 GPUs case with VPP enabled and $1-E$1, PyTorch and PyTorch ES both encounter OOM, whereas STWeaver succeeds with 464.3 TFLOPS. Alternative baseline-safe configurations are slower: disabling VPP gives 440.6 TFLOPS, recomputation gives 350.4 TFLOPS, and increasing $1-E$2 to 4 gives 431.5 TFLOPS. In this case, defragmentation enables a configuration whose throughput is 5.4% to 32.5% better than the fallback alternatives.

5. Operating regime, assumptions, and limitations

STWeaver works best when allocation patterns repeat across iterations, many tensors share a small set of sizes, tensor lifetimes are regular enough to be grouped, the workload has large static memory structure, and dynamic and static peaks are temporally misaligned. The paper states that these conditions are common in LLM training, especially Transformer-based dense models and many sparse MoE setups. Dense models benefit most because nearly everything can be planned statically, which is why STWeaver often reaches >99% efficiency there (Huang et al., 22 Jul 2025).

The planner is heuristic rather than globally optimal. This follows directly from the NP-hardness of the underlying Dynamic Storage Allocation problem. The design therefore balances planning quality, generality, and runtime efficiency. A purely offline allocator would be effective for dense deterministic training loops but would fail on dynamic behaviors such as MoE; a purely online allocator would preserve generality but would miss the global reuse opportunities created by repetitive iteration structure. STWeaver’s hybrid design keeps online allocation and fallback precisely to preserve robustness under dynamic request sizes and runtime deviations.

Several limitations are explicit. STWeaver depends on a profiling phase and assumes future iterations resemble the profiled ones. If the workload changes meaningfully, a new plan would be needed. Dynamic behavior cannot be fully pre-planned, so MoE and other conditional execution patterns still require online allocation and caching-allocator fallback. Gains from dynamic reuse are smaller when dynamic and static requests strongly overlap in time, as suggested by the weaker improvement in non-recomputation MoE settings. The scope of managed memory is limited to PyTorch-visible allocations, excluding some kernel- or compiler-originated allocations. Finally, profiling and plan synthesis take minutes: Table 5 reports, for example, 78.82 s of profiling and 24.36 s of plan synthesis for GPT-2-N, 204.73 s and 104.96 s for Llama2-7B-N, and 273.74 s and 374.18 s for Qwen1.5-MoE-N. The paper presents these as acceptable one-time pre-training costs rather than steady-state overheads.

6. Terminological disambiguation: STWeaver, Weave, and WEAVE

The term “STWeaver” in the provided literature refers specifically to the GPU memory allocator introduced in “Reducing GPU Memory Fragmentation via Spatio-Temporal Planning for Efficient Large-Scale Model Training” (Huang et al., 22 Jul 2025). Two other names in the corpus are close orthographically but denote unrelated systems.

“Weave,” in continuous gravitational-wave data analysis, is a semicoherent all-sky search implementation in LALSuite for isolated continuous gravitational waves. Its defining features are coherent per-segment searches using the $1-E$3-statistic, semicoherent combination across segments, template-bank generation using the parameter-space metric, placement on optimal lattices, and use of the reduced supersky metric so that the metric is effectively constant across the all-sky parameter space. That paper explicitly states that it uses the name Weave and does not define “STWeaver”; any connection is described only as a plausible colloquial or variant label evoking “StackSlide + Weave” rather than as a formal algorithmic identity (Wette et al., 2018).

“WEAVE,” by contrast, is the wide-field spectroscopy facility at the prime focus of the William Herschel Telescope, together with its fibre positioner and acquisition-and-guiding software. The fibre-positioner calibration paper describes orientation-dependent metrology, plate maps, park-height prediction, camera calibration, and configuration-time selection of the closest map to the telescope position; the acquisition-and-guiding paper presents automated LIFU acquisition by pattern recognition of stellar asterisms matched against Gaia predictions and MOS/mIFU guiding using up to eight coherent image guide fibre bundles, with full astrometric recalculation and closed-loop correction through the WHT TCS (Hughes et al., 2022); (Gafton et al., 31 Mar 2026).

A plausible implication is that “STWeaver” can be misunderstood in cross-domain searches because “Weave” and “WEAVE” already denote a gravitational-wave search engine and an astronomical instrument software ecosystem. Within the present corpus, however, only the 2025 systems paper assigns the exact name STWeaver, and it does so in the context of GPU memory allocation for large-scale model training.