一句话:提出 T-Control,结合张量依赖图拓扑与实时内存的自适应动态张量重计算系统,显著提升 DNN 训练吞吐。
随着 DNN 模型和 batch size 持续增长,训练面临显存过度消耗问题。张量重计算可在有限内存下支持训练,但动态方法常因贪心驱逐调度和运行时开销而不如静态方法。本文提出 T-Control,一种面向 DNN 训练的动态张量重计算系统。它将追踪到的张量依赖图拓扑与实时内存使用结合,自适应决定张量保留,保护关键张量并减少驱逐引发的重计算;同时扩展 PyTorch 原生内存管理器,采用细粒度内存管理策略提升内存利用率,减少不必要驱逐。实验表明,T-Control 相比最先进的静态与动态张量重计算系统,吞吐量最高分别提升 1.58× 和 1.91×。
With the continuous growth of model and batch sizes, DNN model training increasingly suffers from excessive memory consumption. Tensor rematerialization has emerged as an effective technique to enable training under limited memory constraints. However, dynamic rematerialization methods often underperform static approaches, primarily due to their greedy tensor eviction schedules and runtime overhead. In this work, we develop T-Control, a dynamic tensor rematerialization system for DNN model training. T-Control integrates the topology of the traced tensor dependency graph with real-time memory usage to make informed, adaptive tensor retention decisions, preserving critical tensors and reducing eviction-induced recomputation. Furthermore, by extending PyTorch's native memory manager with our fine-grained memory management strategy, T-control effectively improves memory utilization, thereby reducing unnecessary tensor eviction. Experimental results demonstrate that T-Control boosts throughput by up to 1.58× and 1.91× over state-of-the-art static and dynamic tensor rematerialization systems.
@inproceedings{WangXDWMLPT26,
title = {T-Control: An Efficient Dynamic Tensor Rematerialization System for DNN Training},
author = {Zehua Wang and Junmin Xiao and Xiaochuan Deng and Huibing Wang and Hui Ma and Mingyi Li and Yunfei Pang and Guangming Tan},
booktitle = {ASPLOS},
year = {2026},
doi = {10.1145/3779212.3790230}
}