🤖 machine learning

Memory-Efficient Activation Checkpointing with Sliding Window and Hirschberg's Algorithm for 0/1 Knapsack Solving in PyTorch

本文介绍了一种用于 PyTorch 的内存高效型激活检查点求解器，该求解器结合了滑动窗口算法和 Hirschberg 算法，将峰值内存使用量从 $O(nW)$ 降低至 $O(W)$ ，从而能够解决规模显著更大的 0/1 背包问题，并实现了 25-28% 的运行速度提升，随后被集成到 PyTorch 2.10 中。

原作者： Jędrzej Maczan

发布于 2026-08-11

📖 1 分钟阅读☕ 轻松阅读

CC BY 4.0

原作者： Jędrzej Maczan

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正试图烘焙世界上最美味、最复杂的蛋糕，但你的厨房却非常狭窄拥挤。你的食谱要求你记录下混合的每一种原料、每一次温度变化以及每一次搅拌动作，以便稍后能完美地逆向还原整个过程，看看蛋糕做得如何。问题在于，你的厨房柜台（即你的计算机内存）太小了，装不下所有的笔记。如果你试图把一切都写下来，柜台就会溢出，导致你不得不停止烘焙。这就是科学家们在训练大规模人工智能模型时面临的日常挣扎。他们需要记住很多步骤来教导人工智能，但他们的计算机空间会耗尽。为了解决这个问题，他们使用了一个巧妙的技巧，叫做“激活检查点”（activation checkpointing）。与其写下每一个步骤，不如挑选出最重要的步骤来保存，并约定好稍后重新执行那些不太重要的步骤。这就像是决定哪些照片要放入一本小相册，而哪些照片即使忘了也可以再拍一次。目标是将整个烘焙蛋糕的过程塞进那个小小的厨房里，同时又不丢失食谱的魔力。

长期以来，许多人工智能科学家用来构建模型的计算机程序 PyTorch，对于决定保存哪些步骤有着特定的方式。它将这个决策视为一个经典的谜题，叫做“0/1 背包问题”。想象你是一名徒步旅行者，背着一个只能承载一定重量的背包。你有一份物品清单，每件物品都有一个重量和一个价值（即它对你的帮助程度）。你想挑选出能提供最大价值的物品，同时又不会撑破你的背包。PyTorch 解决此问题的默认方法就像是在一张巨大的纸上写下所有可能的物品组合。虽然这种方法很完美，能找到绝对最优解，但那张纸变得如此巨大，以至于计算机的内存会爆炸，导致程序崩溃。研究人员发现，如果他们只有 100 个项目可供选择，所需的纸张规模大到需要 304 GB 的空间，这远超他们机器上仅有的 64 GB。这是一个完美的解决方案，却仅仅因为无法装进房间而无法使用。

在这篇论文中，作者介绍了一种更聪明的方法来解决这个谜题，称之为 dp_knapsack_sliding_hirschberg。与其试图一次性写下整张巨大的纸，他们使用了一个“滑动窗口”技巧。想象你在读一本长篇小说，但你只有一个小放大镜，每次只能看两页。你让放大镜在书页间滑动，看两页，再看接下来的两页，以此类经过。这样一来，你每次只需要在大脑中保留两页的内容，从而节省了大量的精神空间。然而，仅仅看两页是不够的，因为你还需要记住整个故事；你需要知道具体该挑选哪些物品。为了解决这个问题，他们将滑动窗口与一种古老且聪明的策略——“Hirschberg 算法”结合起来。把这想象成一个“分而治之”的游戏。与其试图一次性解决整个背包问题，不如将物品清单一分为二。他们先解决左半部分，再解决右半部分，然后研究如何将两个最优解结合起来。他们递归地进行这个过程，将问题不断分解为越来越小的部分，直到可以轻松解决，同时始终只使用极少量的内存。

新方法的成果令人印象深刻。作者在一台拥有 64 GB RAM 的计算机上测试了它。虽然旧方法在处理仅 100 个项目的问题时就会崩溃，但新方法成功解决了包含 2,000 个项目的难题，且峰值内存仅为 58.4 GB。这意味着计算机现在可以处理比以前大 20 倍的问题。此外，新方法不仅节省了内存，而且速度更快。在测试中，它比旧方法快了 25% 到 28%。作者通过在特定机器上运行 1,000 次相同的谜题来衡量这一点，并发现新求解器在速度上始终优于旧求解器。至关重要的是，与某些通过猜测答案可能导致轻微偏差的“快速修复”方法不同，这种新方法每次都能找到精确、完美的解。它与旧方法一样准确，但效率更高。

论文证实，这种新方法不仅仅是一个理论；它已成功合并到 PyTorch 软件中，并在版本 2.10 中可用。作者展示了通过结合使用滑动窗口和分而治之的方法，他们可以解决阻碍人工智能模型规模扩大的内存瓶颈。他们并不声称这是解决该问题的唯一方法，也不暗示它适用于每一种类型的计算机谜题，但对于决定保存哪些 AI 步骤这一特定任务，它是一个经过验证、精确且高效的升级。论文排除了旧方法在处理大型模型时仍然适用的可能性，清楚地表明当项目数量过高时，旧方法会失效。相反，他们提供了一个既保留了旧有方式的完美准确性，又消除了内存崩溃风险的解决方案，让科学家们能够在他们狭小的厨房里，烘焙出更大、更复杂的 AI 蛋糕。

技术摘要：结合滑动窗口与 Hirschberg 算法的高效内存激活检查点技术

问题陈述
激活检查点（Activation Checkpointing）是针对在严格内存预算下最小化神经网络运行时间的一项关键技术，其通过选择性地存储中间张量并重新计算其他张量来实现。在 PyTorch 中，torch.compile 框架将“保存哪些操作”与“重新计算哪些操作”的决策建模为一个 0/1 背包问题。默认求解器 dp_knapsack 采用标准的动态规划（DP）方法，需要一个形状为 $(n + 1) \times (W + 1)$ 的完整二维表，其中 $n$ 是操作数量， $W$ 是量化的内存预算。

作者指出一个关键的可扩展性瓶颈：默认求解器的内存复杂度为 $O(nW) $。实证测试表明，这种方法在处理现代深度学习图时会变得不可行；具体而言，当$ n$ 达到仅 100 个项目时，该求解器就会在 64 GB 的机器上触发内存溢出（OOM）崩溃。这一限制阻碍了对具有长序列长度或宽计算图的大规模模型的优化。

方法论
为了解决内存瓶颈同时保留精确的最优解，本文引入了 dp_knapsack_sliding_hirschberg 求解器。该求解器集成了两种算法优化：

滑动窗口技巧（Sliding Window Trick）： 该方法不再维护完整的 DP 表，而是每次仅存储两行。这使得计算最优值的空间复杂度从 $O(nW) $降低到$ O(W)$。
Hirschberg 算法： 由于滑动窗口优化会丢弃用于回溯（项目选择）所需的中间行，作者引入了 Hirschberg 的分治策略。该算法递归地将项目集分为两半，使用滑动窗口计算每一半的 DP 轮廓，并确定能使总价值最大化的最优容量分割点 ( $k^*$ )。随后，它在子问题上进行递归，以恢复要保存的精确项目集。

为了防止在大规模 $n$ 情况下发生栈溢出，实现中使用了显式的 LIFO 栈而非递归函数调用。该算法的空间复杂度为 $O(W + n)$ ，时间复杂度为 $O(nW \log n)$ 。

核心贡献

算法集成： 将滑动窗口技术与 Hirschberg 算法成功结合，形成了一个用于 PyTorch 的生产级背包求解器。
实现： 该求解器已被合并至 PyTorch 代码库，并在版本 2.10 中发布。
实证验证： 通过综合基准测试，将新求解器与默认的 dp_knapsack、基于整数线性规划的求解器 (ilp_knapsack) 以及贪心启发式算法 (greedy_knapsack) 进行了对比。

结果
评估是在模拟真实激活内存规划场景的合成实例上进行的，使用的是一台配备 AMD Ryzen 7 9800X3D 处理器的 64 GB RAM 机器。

内存效率： 新求解器在 $n = 2,000$ 时运行，峰值内存使用量为 58.4 GB。相比之下，默认求解器在 $n = 100$ 时即宣告失败。这代表了计算问题规模实现了 20 倍 的提升。在 $n=100$ 时，DP 表的内存占用从约 304 GB（默认）降至约 6 GB（新求解器）。
运行时性能： 尽管理论时间复杂度有所增加（ $O(nW \log n)$ 对比 $O(nW)$），但新求解器表现出了比默认 dp_knapsack 持续 25–28% 的运行速度提升。作者将其归功于优化的内存访问模式和固定大小缓冲区的使用。
最优性： 与默认求解器一样，dp_knapsack_sliding_hirschberg 能产生精确的最优解。这使其区别于被发现最多存在 7.4% 次优性的 greedy_knapsack。

意义与主张
本文声称 dp_knapsack_sliding_hirschberg 在不牺牲解的最优性前提下，解决了 PyTorch 默认激活检查点求解器固有的内存瓶颈。通过支持优化规模显著更大（操作数量高达 20 倍）的计算图并提供运行速度提升，该方法被视为对 PyTorch 生态系统的实质性改进。

作者建议将该新求解器作为大型或内存受限图中需要精确解且无法使用 SciPy（ilp_knapsack 所需）时的默认选择。这项工作被定位为将标准教科书技术（滑动窗口）和经典算法（Hirschberg）应用于解决现代深度学习编译器中特定且高影响力的工程挑战的实践应用。

技术摘要：结合滑动窗口与 Hirschberg 算法的高效内存激活检查点技术

类似论文