Peak-Memory Schedules You Can Trust: Exact Realization, Certified Optimality, and the Limits of Recomputation for Inference
Statement-level reordering certifies peak-memory-optimal inference schedules more reliably than regenerating compiler code.
The paper finds that handing a better operator order back to PyTorch Inductor is unreliable: regenerated programs miss predicted peaks in 22 of 54 CPU configurations and 19 of 50 GPU configurations, and GPU runs can change numerical results. Profiling generated statements and permuting them realizes the modeled peak to the byte with bitwise-identical CPU results when allocation is history-independent. Trap lower bounds certify optimality for 94% of held-out wrapper graphs and all 16 microcontroller models; DeepLabV3 reordering saves 10%, and selective recomputation cuts ReDimNet peak memory by 5.2-6.1%.
- Inductor misses predicted peaks in 22 of 54 CPU configurations
- Statement permutation matches modeled peaks with identical CPU results
- Trap bounds certify 94% of wrapper graphs and all 16 tiny models
- DeepLabV3 reordering saves 10% peak memory; training saves up to 7.6%
- Recomputation lowers ReDimNet peaks 5.2-6.1% at under 0.4% extra MACs
- Vendors
- PyTorch
- Products
- PyTorch Inductor
From the original article
1 min read at arxiv.orgReal-time inference on small devices needs low activation memory without extra latency, which makes operator order the most valuable memory lever: it is free in time. We show that the usual way of deploying a better order, handing it back to the compiler and regenerating code, is unreliable. In PyTorch Inductor the regenerated program misses its predicted peak in 22 of 54 CPU configurations.
By Ivan Fomin · arxiv.org
A short excerpt; the article and its images belong to the publisher. Summary, key facts and assessment above are AI-written from the article and can be wrong.