TapeSAC: Automated, Pareto-Optimal Fine-Grained Memory Management for Large Model Training
Kaihui Ren, Yifan Zhao, Zejian Wang, Yingtong Xiong, Xiaoling Zhou, Daohong Xu, Chenhao Guan, Zhikang Mo 等 14 位作者
Selective activation checkpointing (SAC) decides which activations to save for the backward pass and which to recompute: each choice is a point in the plane of peak memory and recomputation time, and the goal is that plane's Pareto frontier. Current practice falls short of this goal. Hand-tuned policies fix a single point that does not adapt when the model, batch or kernels change; PyTorch's SAC reruns each block's forward pass in program order, recomputing values the backward pass never reads; and the memory budget of torch.compile treats recomputation as free or forbidden and selects saved tensors by bytes alone, an assumption that fails for the custom kernels that produce many of a large model's activations. We present TapeSAC, which automates SAC in eager PyTorch without editing the model's forward code. TapeSAC records each block's forward pass at the dispatcher and, during backward, replays only the producers of the values autograd consumes, so saving a tensor eliminates every producer needed only for it. A planner selects the saved tensors from measured recomputation costs under a single limit on step-peak memory. On Qwen3.8-27B and Qwen3.6-35B-A3B, replay alone reduces step time by 3.0% and 1.7% versus full checkpointing at no greater memory cost; memory-aware planning raises these gains to 6.3% and 13.9% under a 90%-of-device-memory limit. The evaluated hand-tuned policies fall short of this memory-recomputation frontier and, on Qwen3.6-35B-A3B, run out of memory at larger batches. Over 500 steps on real data, TapeSAC reduces training time by 8-11%; on Qwen3.5-397B, it is 2.7% faster than a model-specific hand-tuned policy while using 14% less memory.
尚无 AI 解读，展示原文摘要