GMLake: Efficient and Transparent GPU Memory Defragmentation for Large-scale DNN Training with Virtual Memory Stitching

The very techniques invented to save GPU memory during LLM training — recomputation, offloading, LoRA, sharded distributed training — bombard the allocator with irregular requests that shatter its memory pool into unusable fragments, wasting up to 30% of an 80 GB GPU and killing jobs with out-of-memory errors while free memory sits stranded; GMLake fixes this beneath the framework, using CUDA's low-level virtual memory API to STITCH scattered free blocks under one contiguous virtual address (no data moves, no user code changes), turning an NP-hard packing problem into exact-fit block reshaping — reclaiming 9.2 GB on average (up to 25 GB) across eight LLMs, holding fragmentation under 10% at 16-GPU scale, and running batch sizes PyTorch cannot, at identical steady-state throughput.

The argument in one glance

Problem

The GPU has free memory; the job dies of out-of-memory anyway

Training or fine-tuning a large language model is bounded first by GPU memory — 80 GB on an A100 against models whose training state runs to hundreds of gigabytes. Frameworks cannot use the GPU vendor's native allocator directly (cudaMalloc requires device synchronization; training OPT-1.3B through it runs 9.7× slower), so PyTorch and TensorFlow interpose a caching allocator: grab memory once, keep it in a pool, and serve requests by SPLITTING pooled blocks to fit. Splitting is fast — and it fragments. A request larger than any single free block fails with an out-of-memory error even when the free blocks, summed, would hold it several times over.
What people did

Save memory with recomputation, offloading, LoRA, and sharding

The ecosystem's answer to memory pressure is a family of reduction techniques: recomputation (recompute activations in the backward pass instead of storing them), offloading (swap optimizer state and tensors to CPU memory), low-rank adaptation (LoRA — train small adapter matrices instead of full weights), and sharded distributed training (DeepSpeed ZeRO-3, FSDP — partition weights, gradients, and optimizer state across GPUs). Each genuinely shrinks the footprint, and practitioners stack them.
Why it fell short

The memory savers are the fragmenters

Every one of those techniques replaces a few large, regular, long-lived allocations with many small, irregular, short-lived ones — tensors swap in and out, activations appear and vanish mid-backward-pass, shards shrink as GPU count grows. Measured on OPT-1.3B across four A100s: plain PyTorch runs at 97% memory utilization; stacking LoRA, recomputation, and offloading drops it as low as 70%. Scaling from 1 to 16 GPUs drops utilization from over 90% to 76%, because ZeRO-3's per-GPU shards get smaller and more numerous. GPT-NeoX-20B's trace goes from 46 thousand allocations averaging 93 MB to 76 thousand averaging 85 MB once the techniques switch on. The obvious fix — CUDA's low-level virtual memory API, which can present scattered physical chunks as one contiguous virtual range — is 115× slower than cudaMalloc when used naively, because every 2 MB chunk pays mapping and access-permission costs.
This paper's idea

Stitch the fragments under one virtual address — and let training's rhythm hide the cost

GMLake (Shanghai Jiao Tong University and Ant Group) rebuilds the caching allocator's internals on the virtual memory API while keeping its interface exactly — transparent to models, reduction techniques, and platforms alike. Physical memory lives in primitive blocks (pBlocks, uniform 2 MB chunks); when a request exceeds every free block, virtual memory STITCHING fuses several non-contiguous pBlocks under one freshly reserved virtual range (an sBlock) — no data moves, and since one physical chunk may sit under multiple virtual addresses, nothing needs unmapping. Split and stitch together let the allocator manufacture an EXACT-FIT block for any request, sidestepping the NP-hard packing problem. The expensive virtual-memory calls happen only during the first few iterations: DNN training is periodic, so after about four iterations the pool's shape matches the workload and every request hits an instant exact match. Results: 15% average fragmentation reduction (up to 33%), 9.2 GB average memory savings (up to 25 GB) across 76 workloads on eight LLMs, ~90% utilization sustained at 16 GPUs, batch sizes PyTorch cannot reach — at identical steady-state throughput.

Why this problem is worth solving

An A100’s 80 GB is the hard wall every large-model practitioner negotiates with daily — the number that decides the largest batch, the longest sequence, whether fine-tuning fits on the cluster at all. Wasting a quarter of it on bookkeeping is not a rounding error; it is the difference between a job that runs and one that dies.

the memory-management layer PyTorch and TensorFlow interpose between tensors and the GPU, because native cudaMalloc/cudaFree require device synchronization and run an order of magnitude too slow for training loops. The allocator grabs memory from the device once and recycles it through a pool using best-fit with coalescing (BFC): find the smallest free block that fits, SPLIT it if it is larger than the request, mark freed blocks inactive without returning them to the device, and merge adjacent inactive neighbors. In the best case every physical allocation happens exactly once

design served the classic era well: a transformer is a stack of identical layers allocating identical tensor sizes in a steady rhythm, so splits are few and fragments rare.

What broke the design is the modern training stack itself. The paper’s characterization is pointed: memory fragmentation — free memory stranded in blocks too small or badly shaped to serve incoming requests, measured as the gap between reserved and peak-active memory — reaches 30% precisely when practitioners deploy the memory-REDUCTION arsenal, because recomputation, offloading, LoRA, and ZeRO-style sharding all convert the allocation stream from few-large-regular to many-small-irregular. The failure mode is bitter: an out-of-memory crash with gigabytes technically free, or a batch size capped far below what the silicon could hold. And the escape hatch CUDA provides —

the low-level API family (cuMemAddressReserve, cuMemCreate, cuMemMap, cuMemSetAccess) that CUDA introduced in the spirit of mmap: reserve a contiguous VIRTUAL address range, create physical chunks wherever they fit, and map one onto the other, so contiguity becomes a property of addresses rather than of silicon. It removes the physical-contiguity requirement that causes fragmentation — but naively replacing the allocator with it is 115× slower than cudaMalloc for 2 GB in 2 MB chunks, dominated by per-chunk mapping and access-permission setup

— is a trap without an amortization strategy. The problem is real, the primitive exists, and the entire contribution is the engineering that makes the primitive affordable.

What people did before, and where it broke

Native allocation. Correct, synchronous, and 9.7× slower end-to-end — disqualified for training loops from the start.
Best-fit caching allocators (PyTorch, TensorFlow). The workhorse. Splitting maximizes pool utilization for regular workloads and quietly manufactures fragments for irregular ones; freed blocks merge only with ADJACENT inactive neighbors, so scattered holes never recombine.
Naive virtual-memory allocation. Solves fragmentation by construction, costs two orders of magnitude in allocation latency — unusable without the pool-and-reuse layer this paper builds.
Application-level memory paging (vLLM). Pages the attention KV cache inside the tensor — brilliant for serving, but scoped to self-attention specifically; it restructures the application rather than the allocator, and does nothing for the training-time pool.
Defragmentation by copying. Classic compaction moves live data to consolidate holes — prohibitive at tens-of-gigabytes scale in a training loop’s inner timing.

How GMLake works

Same interface as the caching allocator; stitching underneath

Two pools, two block types new

The primitive pool (pPool) holds pBlocks — the sole owners of physical memory, built from uniform 2 MB chunks, sorted by size. The stitched pool (sPool) holds sBlocks — virtual-address ranges mapped onto the physical chunks of MULTIPLE pBlocks, soft-link style. A physical chunk can sit under several virtual addresses simultaneously (CUDA VMM allows many-to-one mapping), so stitching never requires unmapping anything; active-state flags on pBlocks guarantee each chunk serves only one live tensor.
↓ →
BestFit with four outcomes new

Each allocation request is classified: S1 exact match (an sBlock or pBlock of precisely the right size — the only case where sBlocks are handed out, and the instant path); S2 one larger free pBlock (split it, allocate one half, stitch the remainder for later); S3 no single block fits but several sum to enough (STITCH them into a new sBlock under one reserved virtual range); S4 not even the sum suffices (allocate a fresh pBlock from the device, stitching it with existing free blocks). Out-of-memory is reported only when memory is genuinely, physically exhausted.
↓ →
Stitching mechanics new

To fuse a 1 GB and a 2 GB free block: reserve a fresh 3 GB virtual range (cuMemAddressReserve — microseconds), then map it onto both blocks' existing physical chunks (cuMemMap). No cuMemCreate — sBlocks never own new physical memory. No data movement, no copying: the fragments stay where they are; only the address space is rearranged over them.
↓ →
Deallocation without deallocating new

Frees flip active-state flags — physical memory never returns to the device mid-training. When stitched structures accumulate past a capacity threshold, StitchFree releases the least-recently-used sBlocks (structures only, not physical memory) as a robustness fallback; a fragmentation floor (e.g., 128 MB) exempts small blocks from splitting and stitching entirely, preventing death by a thousand slivers under extremely irregular patterns.
↓ →
Warmup, then silence new

The expensive states (S2, S3, S4) fire during the first iterations while the pool's shape adapts to the workload. Because DNN training is PERIODIC — every iteration allocates the same sizes in the same order — the union of pools becomes a tape recording of the allocation pattern, and after about four iterations every request lands in S1 exact match. The 115× VMM penalty is paid a handful of times, then never again; steady-state throughput equals PyTorch's.
↓ →
Drop-in integration

5,000 lines of C++ replacing internals of PyTorch's caching allocator (adapted to PyTorch 1.13 and 2.0), toggled by configuration. Models, reduction techniques (LoRA, recomputation, offload), and platforms (DeepSpeed, FSDP, Colossal-AI) run unmodified — transparency is the design constraint that scopes every other choice. Open-sourced as part of the GLake project.

The metric: fragmentation ratio = 1 − (peak active memory / peak reserved memory) — active being what tensors actually hold, reserved being what the allocator has claimed from the device. Standard page-based fragmentation indices do not apply because blocks have arbitrary sizes.

The idea

Fragmentation is a contiguity problem, and contiguity is negotiable. A request fails not because memory is lacking but because the free memory is not ADJACENT — a constraint that lives in the address space, not the silicon. CUDA’s virtual memory API dissolves it: reserve a contiguous virtual range, map it over scattered physical chunks, and the tensor sees one flat buffer while the physics stays scattered. GMLake’s stitching is defragmentation without compaction — the classic copy-and-consolidate cost (moving gigabytes inside a training step) is replaced by address bookkeeping measured in microseconds. The elegant enabling detail: physical chunks can sit under multiple virtual addresses at once, so old mappings need no teardown; a stitch is purely additive.

Don’t pack requests into holes; reshape the holes to fit requests. Allocating from fragmented free lists optimally is bin-packing — NP-hard, and heuristics like best-fit are exactly what produced the mess. Split and stitch change the game: any request can be given an EXACT-FIT block in linear time, by splitting a too-large block or fusing too-small ones. The search problem disappears because the allocator stopped searching and started manufacturing. This is the same intellectual move as the paged-KV-cache work on the serving side — virtualize away a contiguity constraint the algorithm never truly needed — applied one layer down, at the framework pool where TRAINING’s irregularity lives.

Amortize the primitive over the workload’s rhythm. The raw VMM API is 115× too slow, and GMLake’s answer is not a faster API but a workload observation: training is periodic. Every iteration replays the same allocation script, so a pool whose shape was carved by iteration one serves iterations two through two hundred thousand from the exact-match fast path, with zero VMM calls. The warmup (about four iterations) is the entire price. This is the predictability dividend this site keeps meeting — know the future because it repeats, and spend accordingly — cashed here not for prefetching but for allocator strategy.

GPT-NeoX-20B, fine-tuned without dying

Four A100s, LoRA + recomputation, a batch PyTorch cannot run

From the paper's memory-trace analysis: reserved vs active memory over 600 seconds of fine-tuning. PyTorch's run ends at 200 seconds; GMLake's does not.

Step 1 of 5

The irregular stream arrives

LoRA allocates small adapter tensors beside huge frozen weights; recomputation frees activations in the forward pass and reallocates them mid-backward; ZeRO-3 gathers and scatters shards continuously. 76 thousand allocations, mixed sizes, short lifetimes. Under PyTorch's best-fit splitting, the pool shatters: reserved memory balloons far above active memory — the gap IS the fragmentation — until a large request finds no home and the job dies of out-of-memory at 200 seconds, with free gigabytes stranded in slivers.

plain PyTorch on this model: 46K allocations averaging 93 MB; with the techniques: 76K averaging 85 MB
fragmentation here is a side effect of memory-SAVING techniques — the cure created the disease

Why it works

Fragmentation is an address-adjacency constraint, not a capacity constraint → mapping one virtual range over scattered physical chunks removes the constraint without moving a byte → out-of-memory now means physically full, not badly shaped.
Split plus stitch can manufacture an exact-fit block for any request → the NP-hard bin-packing search collapses into linear-time reshaping → best-fit heuristics and their pathologies exit the picture.
Physical chunks tolerate multiple simultaneous virtual mappings → stitching is additive, never requiring unmap or teardown → the fast path stays fast and failure windows stay closed.
Training replays the same allocation script every iteration → the pool’s shape converges to the workload’s tape in about four iterations, after which every request is an instant exact match → the 115× virtual-memory penalty is a one-time warmup cost, invisible in steady state.
The interface never changes → every model, reduction technique, and platform combination runs unmodified → adoption is a configuration flag, and the technique composes with the very optimizations that caused the problem.

What it costs

Warmup. The first few iterations pay real virtual-memory overhead while the pool adapts; jobs of thousands of iterations amortize it to nothing, but the benefit assumes the periodic-workload premise.
Irregularity beyond the pattern. Truly aperiodic allocation streams keep triggering the expensive states; the LRU StitchFree fallback and the 128 MB fragmentation floor bound the damage rather than eliminate it.
Vendor coupling. The design rides NVIDIA’s CUDA VMM API and its many-to-one mapping semantics; portability to other accelerators depends on equivalent primitives existing there.
Tuning surface. Chunk size, fragmentation floor, capacity thresholds are empirically configured hyper-parameters — sane defaults ship, but the paper is candid that they were chosen by best practice, not derived.
Scope of evidence. Fine-tuning scenarios up to 20B parameters on up to 16 A100s; the fragmentation metric is the authors’ own (arbitrary block sizes rule out standard indices), and from-scratch pretraining at larger scales is extrapolation.

The evaluation

Eight open-source LLMs fine-tuned across DeepSpeed, FSDP, and Colossal-AI with LoRA/recomputation/offload combinations, on one to sixteen A100-80GB GPUs (NVLink within nodes, RDMA across), against PyTorch 2.0’s caching allocator.

Three allocators for an irregular training stream

Property	PyTorch/TensorFlow caching allocator (BFC)	Naive CUDA virtual-memory allocator
Allocation speed	Fast — the 10× win over native cudaMalloc that made it standard
Fragmentation under reduction techniques
Out-of-memory behavior	Fails when no single free block fits — with gigabytes stranded; GPT-NeoX-20B run dies at 200 s	—
Transparency	The interface everyone already uses	Requires rebuilding allocation logic; unaware of the pool layer where workload knowledge lives
Relation to application-level paging (vLLM)	—	—

Testbed: nodes of eight A100-80GB with NVLink, two nodes via RDMA for multi-node runs; CUDA 11.4. Models: OPT-1.3B/13B, GPT-2, GLM-10B, Vicuna-13B, GPT-NeoX-20B. Implementation: 5,000 lines of C++ inside PyTorch's caching allocator, open-sourced in the GLake project.Underlined values explain themselves when clicked.

Background

The address-indirection trick, one layer down

GMLake is the training-side sibling of a page this site covers elsewhere: the serving-side work that uses the same CUDA virtual-memory API to keep KV caches virtually contiguous. Together they form a clean arc — contiguity constraints that algorithms never truly needed, dissolved by address indirection rather than data movement — which even reached quantum compilation this week in the form of qubit reuse (live working set over total allocation). Two more site threads run through this page. The fragmentation diagnosis is a cross-resource coupling of the kind the site keeps pricing: memory-SAVING techniques cause memory WASTE at the allocator layer, another cure-creates-disease loop like the worst-case provisioning stories. And the warmup amortization is the predictability dividend again — the GPU-memory prefetching page exploits the same iteration periodicity for data movement; GMLake spends the identical observation on allocator strategy: when the future is a replay, you only pay for the first take.

If you remember only three things

1. GPU memory fragmentation is caused by the memory-reduction stack itself: recomputation, offloading, LoRA, and ZeRO-style sharding turn few-large-regular allocations into many-small-irregular ones (GPT-NeoX-20B: 46K allocations averaging 93 MB becomes 76K averaging 85 MB), and best-fit splitting strands up to 30% of an 80 GB A100 — so jobs die of out-of-memory with gigabytes technically free, and utilization decays from 97% to 70% as techniques stack and from 90% to 76% as GPUs scale to 16.
2. GMLake keeps the caching allocator's interface and rebuilds its internals on CUDA's low-level virtual memory API: primitive blocks own physical 2 MB chunks, and virtual memory stitching fuses scattered free blocks under one freshly reserved virtual range (no data movement, no unmapping — chunks tolerate multiple mappings), so split-plus-stitch manufactures exact-fit blocks in linear time instead of searching an NP-hard packing space; the naive API's 115× penalty is amortized by training's periodicity — after about four warmup iterations every request is an instant exact match and throughput equals PyTorch's.
3. Across 76 workloads on eight LLMs and three platforms: fragmentation down 15% on average (up to 33%), 9.2 GB reclaimed on average (up to 25 GB), ~90% utilization sustained at 16-GPU scale, and batch sizes PyTorch cannot run — a transparent, open-source, configuration-flag swap whose honest boundaries are the periodicity premise, NVIDIA-specific primitives, and fine-tuning-scale evidence.

Check yourself

Check your understanding

1. Plain PyTorch training runs at 97% memory utilization; adding LoRA, recomputation, and offloading — techniques that REDUCE memory use — drops utilization to 70%. Predict the mechanism of this paradox.
2. Allocating optimally from a fragmented free list is NP-hard bin-packing, yet GMLake claims linear-time allocation with near-zero fragmentation. What dissolves the hard problem?
3. The naive virtual-memory allocator is 115× slower than cudaMalloc, yet GMLake — built on the same API — matches PyTorch's throughput. Where does the penalty go?
4. vLLM already virtualizes GPU memory (paging the KV cache) using similar ideas. Why does GMLake operate at the allocator layer instead — and what does each scope buy?