Models outgrew single GPUs — fine-tuning a large language model needs far more memory than any card has — and CUDA Unified Memory can oversubscribe GPU memory transparently but pays for it in page faults so costly that naive UM runs 3× slower than it should; DeepUM keeps the transparency and kills the cost by learning what DNN training hands it for free, a kernel sequence that repeats identically every iteration, recording it in two correlation tables (which kernel follows which, and which 2 MB memory blocks each kernel touches) so page faults trigger chained prefetching that runs ahead across kernel boundaries — cutting faults to under 0.1% of naive UM's, running 3.06× faster, and fitting batch sizes 5–13× larger than tensor-level swapping frameworks, fully automatically.
The argument in one glance
-
Problem
The models stopped fitting
State-of-the-art networks need memory no single GPU has — GPT-3's 175 billion parameters occupy ~326 GB in FP16 against an 80 GB flagship card. Pre-training happens on corporate clusters, but the democratic promise of deep learning is fine-tuning on small systems — and even fine-tuning now exceeds a single GPU. The GPU's fast memory must somehow be backed by the CPU's big memory, with data shuttling across PCIe at the right moments.
-
What people did
Swap tensors, or fault pages
One lineage (vDNN, AutoTM, SwapAdvisor, Capuchin, Sentinel, IBM's LMS) swaps whole tensors between CPU and GPU, planned via profiling, integer programming, or framework integration. The other uses CUDA Unified Memory: a single address space where GPU access to an absent page raises a page fault and the driver migrates it on demand — fully transparent, no planning needed.
-
Why it fell short
Coarse and brittle, or transparent and slow
Tensor-level swapping moves data in indivisible tensor-sized chunks, needs per-model planning or user hints, and — allocating real GPU memory — inherits fragmentation: LMS hits out-of-memory errors at batch sizes DeepUM handles easily, and a kernel whose working set exceeds GPU memory simply cannot launch. Unified Memory dodges all of that but bleeds on every fault: the faulting SM's TLB locks until all its faults resolve, migration is expensive PCIe I/O, and when memory is full, EVICTION sits on the fault-handling critical path. Naive UM training runs about 3× slower than achievable.
-
This paper's idea
Keep the transparency; learn the repetition
DeepUM (Seoul National University) exploits the most predictable workload in computing: DNN training repeats the same kernel sequence with the same memory accesses every iteration. A runtime wrapper hashes each kernel launch into an execution ID; a kernel-space driver records two correlation tables — which kernel follows which, and which 2 MB UM blocks each kernel touches. A page fault then triggers CHAINED prefetching: successors of the faulted block, then, on reaching the current kernel's recorded end, a jump — via predicted next kernel — to that kernel's start block, running ahead N kernels. Two more fixes take eviction off the critical path (correlation-informed pre-eviction) and delete pointless traffic (pages of PyTorch's freed-but-cached blocks are invalidated, not migrated). Results on a V100: faults cut to under 0.1% of naive UM's, 3.06× speedup over UM, 1.11× over LMS, and maximum batch sizes 5–13× larger than every compared framework — with under ten lines changed in PyTorch and zero in user code.
Why this problem is worth solving
The gap this paper attacks is economic as much as technical.
Training frontier models is a datacenter activity, but the
ecosystem’s real breadth comes from fine-tuning pre-trained
models on single machines — and model growth has pushed even
that past a lone GPU’s memory.
a single virtual address space
shared by CPUs and GPUs. Since the Pascal generation, GPUs
have a page-migration engine: touching a page that is not
resident raises a GPU page fault, and the NVIDIA driver
migrates the 4 KB page in and replays the access. Programs can
therefore allocate more than GPU memory and run unmodified —
oversubscription with zero programmer effort — but faults are
brutal: the faulting SM’s TLB locks until its faults resolve,
migration crosses PCIe, and the driver manages pages in UM
blocks of up to 512 contiguous pages (2 MB)
The alternative lineage buys speed with knowledge: frameworks that swap whole tensors according to plans derived from profiling runs, integer-programming solvers, or hand annotations. But whole tensors are coarse currency, the planners bind to specific frameworks and models, and because these systems allocate real (non-virtual) GPU memory, they inherit two hard limits — fragmentation-induced out-of-memory failures, and the inability to launch any kernel whose own working set exceeds physical memory. The prize DeepUM chases is having it both ways: UM’s transparency and 4 KB-granularity robustness, at planned-swapping speed — by noticing that the “plan” is sitting in plain sight, re-executed identically every training iteration.
What people did before, and where it broke
- vDNN, AutoTM, SwapAdvisor, Capuchin, Sentinel (TensorFlow-lineage swappers). Move tensors per a plan: vDNN by layer heuristics, AutoTM by integer programming, SwapAdvisor by genetic search, Capuchin by dynamic profiling, Sentinel by fine-grained profiling with OS cooperation. Effective — Sentinel is DeepUM’s closest performance peer — but none are transparent (profiling passes, annotations, framework surgery), and all move tensor-or-layer-sized units.
- IBM LMS (PyTorch swapper). Automatic tensor-level swapping, directly comparable. It trails DeepUM by 1.11× on speed — and fails outright (out-of-memory from fragmentation) at batch sizes DeepUM runs, unless modified to periodically flush PyTorch’s memory pool, which costs more speed.
- Naive UM, or UM plus hint APIs. Correct and transparent; three times too slow. NVIDIA’s own remedy is manual: scatter cudaMemPrefetchAsync and cudaMemAdvise calls through your code — reintroducing exactly the expert labor UM was meant to remove.
- Classic correlation prefetching. Pair-based tables mapping each missed cache line to its historical successors — proven for CPU caches, never adapted to the GPU page world where the “misses” are page faults observed by a driver and the right granularity is the 2 MB UM block.
How DeepUM works
Two tables, four threads, and a fault that triggers a cascade
-
Runtime: everything becomes UM, every kernel gets an ID new
A user-space wrapper library redirects all GPU memory allocation (PyTorch's allocator included) into Unified Memory space — oversubscription now falls out for free. Wrappers on kernel launches (cuDNN and cuBLAS included) hash each kernel's name and arguments into an EXECUTION ID: the same layer's same kernel gets the same ID every iteration. A CUDA callback ships each upcoming ID to the driver. User code: unchanged. PyTorch: fewer than ten lines.
-
Driver: four kernel threads around the NVIDIA fault path new
A Linux kernel module intercepts GPU page-fault interrupts. The FAULT-HANDLING thread reads the GPU's fault buffer and feeds three others: the CORRELATOR updates the tables from observed faults; the PREFETCHING thread looks up successors and enqueues predictions; the MIGRATION thread moves UM blocks — always draining the fault queue (the GPU is stalled on it) before the prefetch queue.
-
Execution ID correlation table (which kernel comes next) new
One table records kernel-sequence history as (three-predecessor context → successor) records, keeping ALL observed successors per context rather than a fixed few. Kernel-level misprediction is expensive — a wrong jump prefetches a whole wrong working set — so this table spends memory on accuracy. DNN training's fixed, repeated kernel order makes it near-perfect after the first iterations.
-
Per-kernel UM block correlation tables (which memory it touches) new
One table per execution ID records, at UM-block granularity (matching the NVIDIA driver's own management unit — page granularity would drown in entries), which blocks follow which within that kernel — plus START and END block pointers captured at kernel transitions. Best configuration: 2048 rows, 2-way associative, 4 successors; tables live in CPU memory, 13–348 MB per model.
-
Chained prefetching across kernel boundaries new
A fault on block b prefetches b's recorded successors; when the walk reaches the current kernel's END block, the prefetcher consults the execution table, predicts the NEXT kernel, and continues from that kernel's START block — chaining ahead up to N kernels (sweet spot: N=32), pausing until the running kernel finishes, stopping on any new fault or failed prediction. One fault can thus pull in the working sets of many upcoming kernels.
-
Two critical-path repairs new
PRE-EVICTION: eviction inside fault handling doubles the stall, so when free memory runs low the driver evicts early — choosing victims both least-recently-migrated AND not predicted for the next N kernels (the correlation tables again). INVALIDATION: PyTorch's caching allocator keeps freed "PT blocks" in its pool; their page CONTENTS are garbage, yet UM would dutifully migrate them to the CPU on eviction. Ten lines in the allocator tell the driver a block went inactive; its victim pages are simply dropped. Ablation: prefetching alone cuts execution time 45.6%; plus pre-eviction, 63.7%; plus invalidation, 66.7%.
The idea
Training’s repetition is a free oracle. Every swapping planner before DeepUM manufactured its knowledge: profiling passes, ILP solvers, genetic search, user hints. DeepUM observes that the knowledge regenerates itself at runtime — iteration N’s kernel sequence and per-kernel page accesses are iteration N+1’s, exactly — and that UM’s despised fault mechanism is a perfect free sensor: every miss lands in the driver’s fault buffer with its address. So the system that suffers the faults also learns from them, and by the second epoch the tables have turned a reactive fault path into a predictive pipeline. The design would fail on arbitrary software; on DNN training it approaches clairvoyance — GPT-2 XL’s 7.4 million faults per iteration become 687.
Predict at two levels, because errors cost differently. DeepUM splits classic correlation prefetching into a kernel-sequence table and per-kernel block tables, with asymmetric engineering: the execution table hoards unlimited successor records because a wrong KERNEL prediction fetches an entire wrong working set, while block tables stay lean (4 successors) because a wrong BLOCK costs one 2 MB migration. Chaining is what makes the split pay — the prefetcher does not stop at the current kernel’s edge but walks the predicted future, start-pointer to end-pointer, kernel after kernel, staying N=32 kernels ahead of the GPU. The classic single-table, multi-level design collapses into something simpler and further-sighted because the workload’s structure (discrete kernels in a fixed order) is mirrored in the tables’ structure.
The fault path is a critical path — clear it. The two optimizations are exercises in moving work off the moment the GPU is stalled. Eviction inside fault handling means a stalled SM waits for a PCIe round-trip in BOTH directions; pre-eviction does the outbound trip early, guided by the same tables (never evict what the next N kernels will touch). Invalidation goes further and deletes the outbound trip: pages belonging to PyTorch’s freed-but-pooled blocks hold garbage by definition, so migrating them to the CPU is pure waste — cooperation between the allocator’s semantics (“inactive”) and the driver’s mechanics (“drop, don’t move”) that only a cross-layer design can see.
One iteration, watched from the driver
Fine-tuning BERT Large on a 32 GB V100, oversubscribed
Execution ID = hash of kernel name+arguments (stable across iterations). UM block = the NVIDIA driver's 2 MB page-management unit. N = prefetch depth (32 kernels). Numbers from the paper's V100 evaluation.
Step 1 of 6
Iteration 1: faults teach the tables
The first forward pass runs cold. Each kernel launch flows through the DeepUM runtime, which hashes name+arguments into an execution ID and notifies the driver via callback. Every page fault lands in the fault buffer; the correlator records, per execution ID, which UM blocks fault after which — and, in the execution table, which kernel follows which. This iteration pays full UM cost. The next several thousand will not.
the fault mechanism doubles as the training-data collector same-layer kernels hash to the same ID every iteration — the key invariant
Why it works
- DNN training re-executes an identical kernel sequence with near-identical page accesses every iteration → history is prophecy, and the fault handler sees every miss with its address for free → correlation tables converge within an iteration or two and cut faults to under 0.1%.
- Kernel-level and block-level prediction have different error costs → the execution table buys accuracy with unbounded successor records while block tables stay small, and chaining composes them into an N=32-kernel lookahead → one fault stages many kernels’ working sets.
- UM block granularity (2 MB) matches the NVIDIA driver’s own management unit → tables stay tractable for models with terabyte-scale address footprints, and every prefetch moves a unit the driver handles natively → 13–348 MB of CPU-side tables govern the whole run.
- Eviction work is movable in time but fault work is not → pre-eviction (correlation-informed victims) does the outbound transfer before the stall, and invalidation of inactive-PT-block pages deletes the transfer entirely → the fault critical path shrinks to inbound migration alone.
- Virtual-memory management at 4 KB pages sidesteps fragmentation and lets kernels run with working sets beyond physical memory → batch sizes grow until CPU memory, not GPU memory or allocator luck, is the wall — 192 vs 14 on BERT Large.
What it costs
- A kernel module shadowing NVIDIA’s driver. DeepUM intercepts fault interrupts and reuses NVIDIA driver internals — deeply privileged, version-coupled code, well beyond a pip install.
- First-iteration cold start and table memory. The learning iteration pays full UM cost, and tables occupy 13–348 MB of CPU memory per model (plus threads doing lookups — cheap, but present).
- Repetition is load-bearing. Workloads without stable kernel-and-access repetition get nothing: DLRM’s input-dependent embedding lookups defeat DeepUM exactly as they defeat every compared system. Dynamic-shape models, data-dependent control flow, and changing batch compositions all erode the oracle.
- PCIe is still PCIe. Prefetching hides latency; it does not add bandwidth. With heavy oversubscription the ideal bars in the paper’s own figures sit well above what any swapping system achieves — compute-to-transfer ratio decides the ceiling (ResNet’s residual blocks: great; small-oversubscription BERT: little to gain).
- Aggressive lookahead can misfire. Prefetch depth N is a real knob — too deep wastes bandwidth and evicts soon-needed pages; the N=32 sweet spot is empirical, and per-model tuning (like the block-table geometry) is part of the deployment story.
The evaluation
Nine real models spanning language, vision, recommendation, and GANs; direct execution against naive UM and IBM LMS on identical hardware; indirect comparison (same models, same GPU, published numbers) against five TensorFlow-lineage swappers; ablations for each mechanism; wall-socket energy measurement.
Three roads past GPU memory
| Property | Naive CUDA Unified Memory | Tensor-level swapping (LMS; vDNN / AutoTM / SwapAdvisor / Capuchin / Sentinel) | DeepUM |
|---|---|---|---|
| Transparency | |||
| Data movement granularity | |||
| Page faults per iteration (BERT Large, batch 3) | 1,171,717 | Fewer (planned transfers) but coarse-grained stalls remain | |
| Training speed | 1.0× reference (about 3× slower than achievable) | ||
| Maximum batch size | Large (virtual memory) but unusably slow |
Check yourself
Check your understanding
1. DLRM — Facebook's recommendation model — gets essentially zero speedup from DeepUM, from LMS, and from every other swapping framework. Predict the property that defeats them all.
2. DeepUM splits prediction into one kernel-sequence table (unbounded successor records) and many per-kernel block tables (just 4 successors each). Why the asymmetric engineering, and what does chaining add?
3. Pre-eviction and inactive-block invalidation together nearly double prefetching's benefit (45.6% → 66.7% time reduction). What does each remove from the fault critical path?
4. DeepUM runs BERT Large at batch 192 where IBM LMS caps at 14 — a 13× capacity gap on identical hardware. Which differences produce it?
Sources
- ACM DL: DeepUM: Tensor Migration and Prefetching in Unified Memory — ASPLOS 2023 (open access)