峰值内存调度新突破：精确实现与可证明最优性

在边缘设备上运行实时推理，激活内存是稀缺资源。算子顺序作为一种不增加延迟的内存优化手段，价值巨大。但传统做法——将优化后的顺序交回编译器重新生成代码——并不可靠。本文基于最新研究，解析如何通过精确建模和置换语句来实现峰值内存的可靠降低，并给出可证明的最优性证书，同时探讨选择性重计算的理论边界。

编译器重生成代码的不可靠性

研究团队在 PyTorch Inductor 上进行了系统测试，发现重新生成代码的方式存在严重偏差。在 54 个 CPU 配置中，有 22 个配置的实际峰值内存偏离了预测值；其中 9 个配置甚至比原始程序更差。在 GPU 上，50 个配置中有 19 个出现同样问题，并且数值结果也发生了变化。这表明，依赖编译器重新生成代码来落实算子顺序优化，风险很高。

究其原因，编译器在代码生成过程中可能引入额外的临时缓冲区或改变内存分配策略，导致实际内存使用与理论模型不符。这种不可预测性使得开发者难以信任自动化的内存优化流程。

峰值内存（peak memory） 指程序运行过程中某一时刻所有活跃张量（activation）所占内存的最大值。在推理阶段，每个算子的输出张量在其被后续算子消费后即可释放；不同算子顺序会改变张量的活跃区间，从而影响峰值内存。与总内存分配量不同，峰值内存决定了设备上能否容纳整个计算图，是边缘部署的硬约束。

算子顺序（operator ordering） 指在不改变数据依赖的前提下，对计算图中无依赖关系的算子重新排列执行次序。例如两个分支的卷积可以交换先后，使其中一个分支的中间结果尽早释放。理想情况下，重排序不增加任何计算量，因此被视为“免费”的内存优化手段。

精确实现：基于语句置换的方法

为了避免编译器重生成带来的不确定性，研究者提出了一种直接对生成程序进行语句置换的方法。他们首先对生成后的程序进行性能分析，然后根据语句自身的读写效果精确建模内存使用，最后对语句进行置换。

理论证明表明，只要内存分配不依赖于历史状态，该方法就能逐字节地实现建模的峰值内存，并且输出结果与原始程序逐位相同。在所有的 CPU 运行中，这一目标都得到了实现。对于 GPU，研究者进一步刻画了缓存分配器破坏该条件的原因，为未来改进提供了方向。

语句置换（statement permutation） 是在编译器生成的底层程序（如循环嵌套或线性语句序列）层面交换语句位置，而非在高层计算图上重排算子。由于底层程序已明确每个语句的读写集合，交换语句时只需检查数据依赖是否被破坏，因此可以精确控制每个张量的生命周期。与图级重排序相比，语句置换绕过了编译器的再次调度和缓冲区分配，避免了代码重生成引入的额外临时变量。

逐位相同（bitwise identical） 指优化后的程序与原始程序在相同输入下产生完全相同的输出位模式。这一性质对数值敏感的边缘推理至关重要，因为重排序若改变浮点累加顺序，可能引入微小误差并在后续层中放大。语句置换只改变语句执行次序而不改变任何算术表达式的求值顺序，因此能保证逐位一致。

可证明的最优性证书

知道何时停止优化同样重要。研究者提出了前向和反向“陷阱”下界，每个下界都可以通过 O(n²) 的最小割计算得到。这些下界能够为找到的算子顺序提供最优性证书。

实验结果显示，在 94% 的留出包装图（wrapper graphs）和全部 16 个微控制器级模型上，该证书在不到一秒内即可确认顺序的最优性。这意味着大多数推理图已经处于最优调度状态，无需进一步优化。分支网络是例外——在这类网络中，重排序是免费且有益的，例如在 DeepLabV3 上实现了 10.0% 的峰值内存降低。在编译训练中，该方法最多可节省 7.6% 的内存。

最小割（minimum cut） 是图论中的经典问题：在带权有向图中寻找一组边，使得删除这些边后源点与汇点不连通，且被删边的权重之和最小。在本文中，研究者将峰值内存下界建模为图上的割问题，每个割对应一种“必须同时活跃”的张量集合，其权重和即为该调度下峰值内存的下界。由于最小割可在多项式时间内求解，因此能以 O(n²) 的代价为任意调度计算一个可证明的下界。

最优性证书（optimality certificate） 指当某个调度的峰值内存等于该下界时，即可证明它已是最优，无需继续搜索。这为优化器提供了可靠的停止准则，避免在已最优的图上浪费计算资源。

选择性重计算的理论边界

在达到认证最优之后，研究者进一步探索了带逐算子上限的选择性重计算。他们证明，对于任何固定上限，该问题都是强 NP 难的，即使在树宽为 2 的图上也是如此；在串并联图上，它精确编码了背包问题；而在有界树宽图上，则可以精确求解。

这一理论结果揭示了重计算优化的计算复杂性。在实际应用中，对于峰值内存仅由单层决定的公共微小模型，选择性重计算并无用处。但对于那些为后期重用而保留廉价值的架构，它却能带来收益。在公开的 ReDimNet 说话人网络上，尽管算子顺序已经最优，选择性重计算仍在仅增加不到 0.4% MACs 的情况下，将峰值内存降低了 5.2-6.1%，且输出逐位相同。

选择性重计算（selective recomputation） 指在内存紧张时丢弃部分中间激活，待后续需要时重新计算，以计算换内存。与训练中对整个反向传播做重计算不同，选择性重计算为每个算子设定一个内存上限，只对超过上限的激活进行重算。

强 NP 难（strongly NP-hard） 意味着即使所有数值参数以一元形式编码，问题也不存在多项式时间算法（除非 P=NP）。这一结论表明，对一般计算图寻找最优的选择性重计算策略在计算上不可行，必须依赖近似算法或针对特定图结构（如树宽有界图）设计精确算法。树宽（treewidth） 衡量图与树的接近程度，树宽越小，图上动态规划的状态空间越小，问题越容易精确求解。

实践启示与未来方向

这项研究传递出几个关键信息。首先，直接置换生成程序的语句比依赖编译器重新生成更可靠，能够精确实现预测的峰值内存。其次，大多数推理图已经处于最优调度，盲目优化可能收效甚微。最后，选择性重计算在特定架构上具有价值，但其理论复杂性意味着需要针对具体问题设计算法。

对于边缘 AI 开发者而言，这意味着在选择内存优化策略时，应优先考虑算子顺序调整，并利用最优性证书避免过度优化。同时，对于分支网络和特定架构，重排序和选择性重计算可以带来额外收益。未来，随着对 GPU 缓存分配器行为的进一步理解，精确实现方法有望在更多硬件上得到应用。