Every technique for fitting a large neural network into a small memory works by moving tensors around in time — recompute this one later, send that one to the host and fetch it back — which buys memory with latency and never touches the one thing that actually makes a tensor large, its shape; this paper (Peking University) adds the missing move, splitting a whole region of the network so it runs in pieces and its intermediate tensors are a fraction of the size, then folds recomputation and swapping into the same rewrite space so the two kinds of decision are made together, reaching the same latency as the best prior systems while using 15 to 85 percent of their peak memory.
The argument in one glance
-
Problem
Memory decides what can run, and models outgrew the card
A training step fails outright when the tensors alive at its peak moment exceed the memory available. A consumer graphics card of this generation offers a few tens of gigabytes, while a large-batch training run can want hundreds. Two things drive the growth. Many tensors live a long time — weights, and the activations saved during the forward pass so the backward pass can use them. And many tensors are simply large, because batches are big, sequences are long, and images are high resolution.
-
What people did
Schedule the graph so fewer tensors are alive at once
The established family of answers manipulates time. Re-ordering picks a different valid execution order so fewer tensors overlap. Rematerialization throws a result away and recomputes it when it is needed again. Swapping sends a tensor to host memory and brings it back before its next use. All three shorten the window during which a tensor occupies memory, and they are the backbone of every memory-limited training system.
-
Why it fell short
Time is only half the problem, and buying it costs latency
Scheduling changes when tensors exist. It cannot change how large they are, so an entire dimension of the problem is outside its reach. Worse, two of its three tools are not free — recomputation spends arithmetic and swapping spends bandwidth, so memory is purchased directly with latency. Meanwhile graph transformation, the technique that can change shapes, has been developed almost entirely for speed, where it is used to merge small operators into large ones, which makes memory worse rather than better.
-
This paper's idea
Add the opposite transformation and search both spaces at once
This paper (Renze Chen, Zijian Ding, Size Zheng, Chengrui Zhang, Jingwen Leng, Xuanzhe Liu and Yun Liang) names the missing move — fission, the dual of merging, which splits a region of the graph along a dimension so the pieces run one after another and every intermediate tensor inside shrinks. It then rewrites rematerialization and swapping as graph transformations too, so shape decisions and time decisions live in one search space and are traded against each other rather than chosen separately.
Why this problem is worth solving
The quantity that decides whether a model runs is a maximum, not a total.
A
neural network for training or inference is described as a computation graph — operators
as nodes, and an edge wherever one operator’s output tensor is an input to another. Pick
an execution order and every tensor acquires a lifetime, beginning when its operator
finishes and ending after the last operator that reads it. At any instant the memory in
use is the total size of the tensors currently alive, and peak memory is the largest that
total ever gets. Only the peak matters, because that is the moment the allocation either
fits or fails. This paper gives a name to the tensors alive at that moment — the memory
hot-spots. It is a useful notion because it says where effort should go. Shrinking a
tensor that is not alive at the peak changes the peak by nothing at all, however large
that tensor happens to be.
The second thing to be clear about is what the existing toolkit can and cannot reach.
RE-ORDERING
chooses a different valid execution order. Since the graph only constrains an order
partially, some orders keep fewer tensors alive at once, and this costs essentially
nothing — the same work happens, in a different sequence. REMATERIALIZATION frees a
tensor and recomputes it later from its inputs, trading arithmetic for space. SWAPPING
copies a tensor to host memory, frees it, and copies it back before its next use, trading
bandwidth for space. The last two are what make aggressive memory reduction possible and
also what make it slow, since both add work that the original program did not do. What
unites all three is that they only move existing tensors around in time. Not one of them
can change the shape of a tensor, so if a single activation is too large to fit, the whole
family has nothing to offer.
What people did before, and where it broke
-
Rematerialization systems. Drop intermediate results and recompute them on demand, chosen by dynamic programming, integer programming, or run-time heuristics. This is the workhorse of memory-limited training and it works on arbitrary networks.
-
The gap. Recomputation is arithmetic the original program did not do, so the latency cost grows sharply as the memory limit tightens. Under a strict limit, recomputing one operator can require first recomputing another.
-
Swapping systems. Offload tensors that will not be needed soon to host memory and prefetch them back. Some systems combine swapping with rematerialization and choose between them optimally for networks with a simple chain structure.
-
The gap. Transfers consume bandwidth, and hiding them behind computation requires free memory to hold the tensor in flight — which is exactly what is scarce. The strongest combined system in this comparison could barely optimise the two networks with complicated skip-connection structure.
-
Graph transformation for speed. A mature line of work rewrites a computation graph into an equivalent one with better performance, typically by merging small operators into larger ones so the hardware is better utilised.
-
The gap. Merging trades memory away to buy speed, which is the wrong direction here. The dual move — splitting — was not part of the vocabulary, and the search machinery built for these systems does not handle it.
-
Micro-batching by hand. Split the whole graph along the batch dimension and run the pieces in sequence. This is the crude version of the idea this paper formalises, and it genuinely helps.
-
The gap. It is one fixed choice applied to the entire graph. The experiment here shows that the best split factor differs with the memory limit, and that a uniform split is worse than choosing different splits for different regions.
How one framework searches shapes and schedules together
One optimisation loop over a combined space of rewrites
-
-
new
-
new
-
new
-
new
-
new
The idea
Fission is the missing dual of a move everyone already uses. Graph transformation for speed merges operators — the three separate projections at the start of an attention block become one larger matrix multiply, which uses the hardware better and temporarily costs memory. Read that backwards and you get fission. Take a region of the graph, choose a dimension that runs through all of it, split the region into several parts along that dimension, and run the parts one after another. Each part computes on a slice, so every intermediate tensor inside the region is a fraction of its former size, and the memory for one part can be released before the next begins. The price is exactly the mirror image of the benefit — smaller operators use the hardware less efficiently, with worse parallelism and worse locality. Defining a rewrite and its inverse so that a search can move in both directions along one axis is the same structure as the equality-saturation optimiser for hardware synthesis covered elsewhere on this site, which also refuses to commit to a direction and lets a cost model decide at the end.
Scheduling is decomposed until only the free part remains. The awkward fact about combining transformation with scheduling is that both trade memory against latency, and after every transformation the schedule has to be recomputed to know what the transformation was worth. The paper’s move is to split scheduling in two. Recomputation and swapping are re-expressed as graph rewrites — recomputation becomes a rule that duplicates an operator so one consumer uses a fresh copy, and swapping becomes a rule that inserts an explicit store and a matching load. What remains in the scheduler is re-ordering alone, which changes no work and therefore costs no latency. Every memory-versus-latency trade now lives in the transformation space, and the scheduler becomes a pure memory minimiser. Making a costly mechanism cheap enough to use freely is the same instinct as the unified-memory prefetcher covered elsewhere on this site, though that system hides the cost of swapping at run time while this one moves the decision to compile time.
Aim only at the tensors that are alive at the peak. The space of possible splits is enormous, since almost any region of the graph is a candidate, so the paper narrows it analytically before searching. A split helps only if the region contains tensors that are alive at the peak, and it hurts to the extent that the region’s inputs must stay resident while all the parts run. Those two quantities give a score for each candidate region — how much hot-spot memory it holds, minus the input memory it will pin — and candidates are drawn from several bands of that score, so the search has both aggressive and gentle options. Regions are taken to be the part of the graph dominated by a single node, which guarantees one entry point and therefore little pinned input. Defining the target set as the tensors alive at the peak is exactly the notion the GPU memory profiler covered elsewhere on this site uses to decide which objects are worth reporting, arrived at independently on the diagnostic side of the same problem.
Four ways to fit a network into a memory of 100
Why splitting and swapping are worth more together than apart
A simplified training graph with a long skip connection, on a device that can hold 100 units.
Step 1 of 5
Do nothing and the peak is ten times the limit
The structure holds 33 tensors of size 32 alive when the 33rd operator runs, for a peak of 1,056. The limit is 100, so the network does not run.
This shape is common in training, and in networks with long skip connections where an early activation must survive until a late layer consumes it. Nothing about the order can fix it, because the tensors genuinely are all needed.
Why it works
- A tree stands in for the split graph. Recording each split as a node saying which region, which dimension, and how many parts → the graph itself is never expanded, since the parts are identical and one copy describes them all → later transformation and scheduling still see a small graph instead of one that doubled in size.
- Candidate regions are scored before they are searched. Estimating the peak reduction as hot-spot memory inside the region minus input memory pinned outside it → only regions that can actually move the peak become candidates → the ablation shows that choosing regions at random instead costs up to 70 percent more peak memory at the same time budget.
- Dominated regions have a single entry. Taking the candidate to be everything dominated by one node → only that node’s output must stay resident while the parts run → the pinned-input term in the score stays small, which is what makes the split profitable.
- Moving trade-offs out of the scheduler simplifies it. Re-expressing recomputation and swapping as rewrites → the scheduler is left with re-ordering, which changes no work and costs no latency → after each rewrite the scheduling question is pure memory minimisation rather than another trade-off to balance.
- Rescheduling is local because dependence is local. Cutting the order at nodes with few independent neighbours → only the interval a rewrite disturbed is rescheduled → 4 to 30 times faster than scheduling from scratch, with identical quality in 94 of 100 tests.
What it costs (trade-offs and scope)
- Splitting is a real performance loss. Smaller operators mean worse parallelism and worse on-chip locality. Under strict memory limits the curve turns sharply upward, because the only remaining moves are fine splits whose utilisation cost is severe.
- The tree limits the search space on purpose. Representing splits hierarchically keeps complexity manageable but excludes split structures the tree cannot express, and the paper says so plainly. The maximum tree depth is a hyper-parameter — too shallow limits the space, too deep slows the search until results get worse.
- Sliding-window axes are excluded. Dimensions where neighbouring outputs overlap, such as the spatial axes of a convolution with a kernel wider than one, are not considered splittable in this work.
- It is a graph-level optimiser, not a compiler. Operator fusion and the other code-generation optimisations that mature compilers perform are absent, so with no memory pressure at all this system is slower than those compilers even though it beats the unoptimised baseline.
- Results come from a search under a time budget. Every reported number is what a three-minute greedy search found, guided by a simulator with cached operator timings. It is a best-effort frontier rather than an optimum.
The evaluation
| Property | What was measured | The result | What it tells you |
|---|---|---|---|
| Memory under a latency budget | |||
| Where it helps most and least | |||
| Latency under a memory budget | |||
| Cost of the extra coordination | |||
| Which heuristics carry the result |
Check yourself
Check your understanding
1. Why can rematerialization and swapping not solve the problem that fission addresses?
2. The system rewrites rematerialization and swapping as graph transformations. What does that buy?
3. In the motivating example, swapping alone reaches a peak of 96 within a limit of 100, yet the paper calls this a poor outcome. Why?
4. Predict what happens if candidate regions for splitting are chosen at random instead of by the hot-spot score.
Sources
MAGIS: Memory Optimization via Coordinated Graph Transformation and Scheduling for DNN. Renze Chen, Zijian Ding, Size Zheng, Chengrui Zhang, Jingwen Leng, Xuanzhe Liu, Yun Liang. ASPLOS 2024, Volume 3, pages 607 to 621. DOI 10.1145/3620666.3651330. Full text read from the ACM Digital Library PDF.