一句话：提出算子分段抽象与分段 ILP，使张量重计算在保持最优性的同时扩展到真实 DNN 模型。
GPU 内存是扩展 DNN 训练的主要瓶颈，激活张量占用常超过显存容量。张量重计算用额外计算降低峰值内存，但现有方法难以兼顾最优性与可扩展性：ILP 求解器可得数学最优解却因搜索空间组合爆炸而不适用于现代模型，启发式方法可扩展却牺牲最优性。本文提出 Bonsai，核心是算子分段抽象：将计算图划为可变大小的段，并构建轻量的分段 ILP，从而压缩搜索空间、剪除冗余解，在保持整模型全局视角的同时将决策变量减少数个数量级。在多种 DNN 架构与模型上评估，Bonsai 可扩展到真实模型，求解器开销比先进 ILP 求解器最多低 10.13×。
GPU memory is increasingly the primary bottleneck in scaling deep neural network (DNN) training, where the activation tensors footprint of a model may exceed the memory capacity. Tensor recomputation is a powerful technique that trades additional computation for reduced peak memory usage. However, existing approaches face a fundamental tension between performance optimality and computational scalability. On the one hand, solvers leverage Integer Linear Programming (ILP) to provide mathematically optimal solutions but suffer from the combinatorial explosion of the search space and thus become intractable for modern DNN models. On the other hand, heuristics-based approaches achieve scalability but sacrifice optimality altogether, resulting in suboptimal execution schedules. The root cause of these inefficiencies in the state of the art is the mismatch in abstraction. This paper introduces Bonsai, a framework that tackles this scalability-granularity tension. At the heart of Bonsai is a novel abstraction of operator segmentation that breaks the computation graph into flexible, variable-sized units to enable a lightweight yet effective segment-based ILP formulation. By having segments, Bonsai collapses the search space and prunes redundant solutions that stall existing solvers. This abstraction enables Bonsai to maintain a holistic view of the entire model, ensuring that no optimization opportunity is lost while reducing the number of decision variables by orders of magnitude. The evaluation across a diverse set of DNN architectures and models demonstrates that Bonsai scales to real-world models, is up to 10.13× lower solver cost than state-of-the-art ILP solvers, and delivers up to 11
@article{paperbot12144,
title = {Bonsai: Efficient and Optimal Automatic Tensor Rematerialization for Memory-Constrained DNN Training},
author = {Dat Nguyen and Vasudha Devarakonda and Anxiao Jiang and Khanh Nguyen},
journal = {Proceedings of the ACM on Programming Languages},
volume = {10},
number = {OOPSLA2},
year = {2026},
doi = {10.1145/3839510}
}