Large deep learning models have achieved impressive performance across a range of applications. However, their large memory requirements, including parameter memory and activation memory, have become a significant challenge for their practical serving. While existing methods mainly address parameter memory, the importance of activation memory has been overlooked. Especially for long input sequences, activation memory is expected to experience a significant exponential growth as the length of sequences increases. In this approach, we propose AutoChunk, an automatic and adaptive compiler system that efficiently reduces activation memory for long sequence inference by chunk strategies. The proposed system generates chunk plans by optimizing through multiple stages. In each stage, the chunk search pass explores all possible chunk candidates and the chunk selection pass identifies the optimal one. At runtime, AutoChunk employs code generation to automatically apply chunk strategies. The experiments demonstrate that AutoChunk can reduce over 80\% of activation memory while maintaining speed loss within 10%, extend max sequence length by 3.2x to 11.7x, and outperform state-of-the-art methods by a large margin.

翻译：大型深度学习模型已在众多应用领域展现出卓越性能。然而，其庞大的内存需求（包括参数内存与激活内存）已成为实际部署中的重大挑战。现有方法主要针对参数内存进行优化，却忽视了激活内存的重要性。尤其对于长输入序列，激活内存预计将随序列长度增加呈现显著的指数级增长。本文提出AutoChunk——一种自动自适应编译系统，通过分块策略高效降低长序列推理过程中的激活内存占用。该系统通过多阶段优化生成分块方案：每个阶段中，分块搜索过程探索所有可能的分块候选方案，分块选择过程则确定最优解。在运行时，AutoChunk通过代码生成自动实施分块策略。实验表明，AutoChunk能在保持速度损失低于10%的前提下，减少超过80%的激活内存，将最大序列长度扩展3.2至11.7倍，并显著超越现有最优方法。