Reordering generated programs makes memory savings verifiable
In plain words
This paper introduces a way to lower the memory peak of a compiled neural-network program by changing the order of statements in its generated wrapper.
It profiles one run, predicts how much memory each step keeps live, then reorders those existing steps while leaving the kernels and buffer reuse alone.
On CPU, the method matched its predicted peak exactly in 78 training runs and 22 inference runs, while preserving outputs and gradients. It also gives certificates that prove a found order is optimal for many tested graphs.
The research
When a compiler finds an operator order that should reduce memory, will the generated program actually use less, and how can its users know whether a better order exists?
These questions matter when a neural-network program must fit a fixed memory budget without relying on slower recomputation or data movement.
The method reorders statements in a generated program and uses lower-bound certificates to show when its memory schedule is already optimal.
The method profiles a compiled program and reorders its existing wrapper statements, keeping kernels and buffer reuse intact. Under six contract conditions, a theorem says this realizes the modeled peak; the pipeline also measures each emitted wrapper.
In 78 CPU training and 22 inference runs, measured peaks matched predictions to the byte, and outputs and gradients matched stock. Separate certificates matched the found peak on 115 of 122 held-out wrapper graphs, certifying those schedules as optimal.
The exactness guarantee depends on six contract conditions. The CUDA caching allocator violates one of them, and the microcontroller results model operator graphs and an arena planner rather than running firmware.
Analysis
Deep
- Original paper authors
- Ivan Fomin, Ivan Koryabkin, Boris Sheludko
- Submitted to arXiv
- License
- arXiv non-exclusive