值得信赖的峰值内存调度：精确实现、经认证的最优性，以及推理中重计算的局限

arXiv:2610.23328v1 公告类型：new 摘要：小型设备上的实时推理需要较低的激活内存，同时不能增加额外延迟，因此算子顺序成为最有价值的内存优化手段：它不会带来时间开销。我们表明，部署更优顺序的常见方式——将其交还给编译器并重新生成代码——并不可靠。在 PyTorch Inductor 中，重新生成的程序在 54 个 CPU 配置中的 22 个配置里未达到其预测的峰值；其中 9 个配置的表现甚至差于原始程序。在 GPU 上，50 个配置中有 19 个配置的表现更差，而且 GPU 上还会改变数值结果。对此，我们改为对生成的程序进行分析，从其自身的语句及读写效果出发精确建模内存，并对这些语句进行置换。我们证明，只要内存分配不依赖历史状态，这种方法就能实现所建模的峰值。在我们所有的 CPU 运行中，它都能精确到字节，并产生按位完全一致的结果。对于 GPU，我们刻画了缓存分配器破坏这一条件的原因。为了确定何时停止优化，我们给出了正向和反向的“陷阱”下界；二者都可以通过 $O(n^2)$ 次最小割计算得出。对于 94% 的留出包装器图，以及我们研究的全部 16 个微控制器级模型，它们都能证明所找到的顺序是最优的，且耗时不到一秒。因此，大多数推理图已经实现了最优调度。分支网络是例外，此时重排序无需额外开销且能够带来收益，例如在 DeepLabV3 上可降低 10.0%。在编译训练中，我们最多可节省 7.6%。在经过认证的最优点之外，我们进一步研究了带有逐算子上限的选择性重计算。即使树宽为 2，对于任意固定上限，该问题仍是强 NP-hard；在串并联图上，它可以精确编码背包问题；而在有界树宽图上则可以精确求解。对于公开的小型模型，如果其峰值只来自单层，重计算毫无用处；但对于会为较晚的重复使用保留廉价值的架构，重计算能够带来收益：在公开的 ReDimNet 说话人网络上，尽管其顺序已经最优，仍可在额外 MACs 少于 0.4% 的情况下将峰值降低 5.2%–6.1%，同时保持按位完全一致的输出。