Papers

arXiv 论文提出精确内存调度方法,可在 1 秒内证明算子排序最优性

Peak-Memory Schedules You Can Trust: Exact Realization, Certified Optimality, and the Limits of Recomputation for Inference

Why we picked it

这篇 arXiv 论文发现 PyTorch Inductor 优化内存排序经常失效,作者给出能逐位复现、可证明最优的调度方法,DeepLabV3 直接省 10% 内存,做边缘部署的值得看看。

论文研究推理时通过重排算子顺序降低激活内存峰值。PyTorch Inductor 重新生成代码后在 54 个 CPU 配置中的 22 个未达到预测峰值,9 个反而比原程序更差。作者改为对生成程序建模并直接置换语句,证明在分配不依赖历史时可精确实现峰值,CPU 上所有运行结果逐位一致。他们给出正反两个下界,每个用 O(n²) 最小割计算,在 1 秒内证明 94% 的 wrapper 图和全部 16 个微控制器级模型已是最优排序。分支网络可再优化,如 DeepLabV3 节省 10.0%,编译训练最多省 7.6%;选择性重计算在 ReDimNet 上降低峰值 5.2-6.1%,额外 MACs 不到 0.4%。

Original · arXiv cs.LG

Peak-Memory Schedules You Can Trust: Exact Realization, Certified Optimality, and the Limits of Recomputation for Inference

Real-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. It is worse than the stock program in 9 of them and in 19 of 50 on a GPU, and on the GPU it also changes numerical results. We instead profile the generated program, model its memory exactly from its own statements and their read/write effects, and permute those statements. We prove that this realizes the modeled peak whenever allocation does not depend on history. It does so to the byte, with bitwise-identical results, in all our CPU runs. On the GPU we characterize why the caching allocator breaks the condition. To know when to stop, we give forward and reverse "trap" lower bounds, each computable with $O(n^2)$ minimum cuts. They certify the found order optimal for 94% of held-out wrapper graphs and all 16 microcontroller-class models we study, in under a second. So most inference graphs are already scheduled optimally. Branching networks are the exception, where reordering is free and pays, e.g. 10.0% on DeepLabV3. In compiled training we save up to 7.6%. Beyond the certified optimum we study selective recomputation with per-operator caps. It is strongly NP-hard for every fixed cap even at treewidth two, encodes knapsack exactly on series-parallel graphs, and is solvable exactly on bounded treewidth. It is useless on public tiny models whose peak is a single layer, but on architectures that hold a cheap value for a late reuse it pays: on the public ReDimNet speaker networks, whose order is already optimal, it lowers the peak by 5.2-6.1% at under 0.4% extra MACs, with bitwise-identical outputs.