一句话：MAGIS 协同图变换与图调度，扩展 DNN 内存-性能权衡空间，大幅降低峰值内存。
深度神经网络的内存消耗因张量生命周期长、形状大而急剧增长。图调度是有效的内存优化手段，可为每个算子/张量决定执行、重算与换入换出的最优时机，但往往显著损害性能，且只能操纵张量生命周期而无法改变形状，限制了优化空间。作者发现，能改变张量形状与图结构的图变换可开辟新的内存-性能权衡空间，然而现有图变换通常单独使用且只关注性能优化。为此，论文提出 MAGIS 框架，将图变换与图调度协同起来：用层次树表示 Fission Transformation（F-Trans）这一可有效缩减子图张量形状的变换；基于图结构分析构建轻量级搜索空间以控制复杂度；将图调度分解为图变换与重排序，并设计增量调度算法以在每个图变换步骤后缓解调度开销，从而高效协同二者。实验表明，在相同延迟约束下，MAGIS 的峰值内存仅为现有最优工作的 15%~85%，并在内存与性能的双目标优化上取得更优的 Pareto 边界。
Recently, memory consumption of Deep Neural Network (DNN) rapidly increases, mainly due to long lifetimes and large shapes of tensors. Graph scheduling has emerged as an effective memory optimization technique, which determines the optimal execution, re-computation, swap-out, and swap-in timings for each operator/tensor. However, it often hurts performance significantly and can only manipulate tensors' lifetimes but not shapes, limiting the optimization space. We find that graph transformation, which can change the tensor shapes and graph structure, creates a new trade-off space between memory and performance. Nevertheless, graph transformation are applied separately so far, with primary focus on optimizing performance and not memory. In this paper, we propose MAGIS, a DNN memory optimization framework that coordinates graph transformation with graph scheduling. MAGIS uses a hierarchical tree to represent Fission Transformation (F-Trans), a type of transformation which can effectively reduce tensor shapes in a sub-graph. To keep the complexity low, we build a light-weight search space based on graph structure analysis. MAGIS decomposes graph scheduling into graph transformation and re-ordering and designs an incremental scheduling algorithm to alleviate the scheduling overhead after each graph transformation step to efficiently coordinate them. Experimental results show that compared to state-of-the-art works, MAGIS only uses 15%~85% of their peak memory usage with the same latency1 constraint and obtains a better Pareto boundary in dual-objective optimization of memory and performance. Our code is now available at https://github.com/pku-liang/MAGIS.
@inproceedings{ChenD0ZLL024,
title = {MAGIS: Memory Optimization via Coordinated Graph Transformation and Scheduling for DNN},
author = {Renze Chen and Zijian Ding and Size Zheng and Chengrui Zhang and Jingwen Leng and Xuanzhe Liu and Yun Liang},
booktitle = {ASPLOS},
year = {2024},
doi = {10.1145/3620666.3651330}
}