Reordering generated programs makes memory savings verifiable
Questa pagina non ha ancora una traduzione verificata. È mostrata la versione inglese. Leggi la versione inglese
In parole semplici
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.
La ricerca
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.
Analisi
Dettagli
- Autori del paper originale
- Ivan Fomin, Ivan Koryabkin, Boris Sheludko
- Inviato ad arXiv
- Licenza
- arXiv non-exclusive