MAGIS: Memory Optimization via Coordinated Graph Transformation and Scheduling for DNN

文章标题与作者/机构

文章标题：MAGIS: 通过协同图变换和调度的深度神经网络内存优化
作者与机构：
- Renze Chen (北京大学)
- Zijian Ding (加州大学洛杉矶分校)
- Size Zheng (北京大学)
- Chengrui Zhang (北京大学)
- Jingwen Leng (上海交通大学)
- Xuanzhe Liu (北京大学)
- Yun Liang (北京大学)

A1 主要贡献

深度神经网络（DNN）的内存消耗随着其拓扑结构和规模的复杂化而持续增长，这主要归因于两个因素：一是大量张量（如模型参数、训练前向传播中的激活值、复杂网络中的中间张量）具有较长的生命周期；二是许多张量（如大批量大小、长序列长度、高分辨率图像）具有较大的尺寸。这给服务器和移动设备的计算带来了巨大挑战。

核心问题: 现有的内存优化技术主要分为两类，各有其局限性。
1. 图调度 (Graph Scheduling)：包括重计算（rematerialization）、交换（swapping）和重排序（re-ordering）。这类技术通过操纵张量的生命周期来减少峰值内存，但通常会显著损害性能，并且无法改变张量的形状，从而限制了优化空间。
2. 图变换 (Graph Transformation)：通过等价变换改变图的结构来优化性能，主要分为聚合变换（Aggregation Transformation, A-Trans）和中间变换（Interim Transformation, I-Trans）。然而，现有工作主要关注性能优化，而未充分利用图变换进行内存优化。特别是聚合变换的对偶操作——分裂变换 (Fission Transformation, F-Trans)，虽然能通过拆分大算子有效减少内存占用，但会引入两大挑战：
* 复杂性：F-Trans 会导致计算图规模急剧增长，增加了后续优化的难度，且其自身的搜索空间巨大。
* 协同优化：图变换和图调度之间存在复杂的权衡关系，需要进行高效的协同优化，但这两种优化本身都非常复杂。

研究目标与创新点:
为了解决上述挑战，本文提出了 MAGIS，一个通过协同图变换和图调度来优化 DNN 内存的框架。
- 设计并实现了 MAGIS：一个基于协同图变换和图调度的内存优化框架。
- 形式化并优化了图分裂变换：本文形式化了图分裂变换（F-Trans），并提出使用分裂层次树 (Fission Hierarchy Tree, F-Tree) 来表示它，通过图结构分析来减少其巨大的搜索空间，从而在不实际增加图复杂度的前提下探索 F-Trans 的优化潜力。
- 提出了高效的协同优化算法：
1. 将重计算和交换等调度技术分解为图变换和重排序，从而将内存与性能的权衡完全移至变换阶段，简化了调度过程。
2. 设计了一种增量图调度算法，在每次图变换后，能够利用先前的调度信息高效地生成新的调度方案，大幅降低了调度开销。

实验结果表明，与现有技术相比，MAGIS 在相同的延迟约束下，峰值内存使用量仅为它们的 15%~85%，并在内存和延迟的双目标优化中获得了更优的帕累托边界。


图 1. 图变换示例。(a) 和 (b) 是从 TASO 【25，Zhihao Jia et al. TASO: optimizing deep learning computation with automatic generation of graph substitutions. SOSP 2019】中借鉴的变换，用于优化性能。(c) 是聚合变换的对偶变换，可以有效地用性能换取内存。

A3 背景知识与动机

表 1. 符号表

2.1 计算图

  • 图结构：DNN的训练或推理过程通常表示为“计算图” G。G = (V(G), E(G))，其中 V(G) 是算子集合，每个算子有若干输入张量和一个输出张量；E(G) ⊆ V(G) × V(G) 是算子间的数据依赖关系集合。(v1, v2) ∈ E(G) 表示 v1 的输出张量是 v2 的一个输入张量。本文使用的相关符号如表1所示。在没有歧义的情况下，我们使用 xxx(o) 作为 http://o.xxx(o) 的缩写。一些符号可以从其他符号派生，例如 G.inps(S) = (∪v∈S G.pre(v))\S，以及 G.outs(S) = (outs(S) ∪ ∪v∈V(G)\S G.pre(v)) ∩ G。如果从入口节点到节点 u 的每条路径都必须经过节点 v，则节点 v 支配节点 u；此时 v 是 u 的支配者。节点 u 的直接支配者是指 u 的支配者中，除了 u 本身之外，被 u 的所有其他支配者所支配的那个。支配树【4，V Aho Alfred et al. Compilers Principles, Techniques & Tools. 2007】是一棵树，其中每个节点的父节点是它在图中的直接支配者。计算图通常有许多输入节点（例如，输入张量、标签张量和权重张量），因此我们这里使用的支配树通常以输入张量作为入口。注意，对于 G = T(G')，G 本身也是一个图，表1中的操作也适用于它。例如，节点 v 在 T 中的子节点集合是 T.suc(v)。T 的节点也属于 G'，即 V(T(G')) ⊆ V(G')。

  • 执行延迟：在单机情况（例如，单卡GPU）下，图中的算子通常按顺序执行，该顺序 π = (v1, v2, ..., vn) 必须满足算子之间的数据依赖关系。图的执行延迟可以估计为算子延迟的总和：cost(G) ≈ Σv∈V(G) cost(v)。

  • 内存使用：给定一个拓扑序 π = (v1, v2, ..., vn)，假设 ti 是第 i 个算子完成时的时间戳，我们可以计算每个算子 vi 的输出张量的生命周期：开始时间戳是 ts_i = ti-1，释放时间戳是 tf_i = max_vj∈suc(vi) tj。基于每个张量的生命周期，在执行 vi 期间活跃的张量集合是 Ai = {oj | ts_j ≤ ti ≤ tf_j}。那么在 vi 执行期间的活跃内存使用量是 Mi = Σo∈Ai |o|，图 G 执行期间的峰值内存使用量为：M_peak = max_i Mi。我们将内存热点定义为对峰值内存使用量有贡献的张量集合，即当达到峰值内存使用量时活跃的张量：H = ∪{Ai | i ∈ {1, 2, ..., n} ∧ Mi = M_peak}。

2.2 图调度与变换

  • 图调度：是一类广泛使用的DNN内存优化技术，它通过操纵张量的生命周期来安排何时执行（重排序【3，Byung Hoon Ahn et al. Ordering Chaos: Memory-Aware Scheduling of Irregularly Wired Neural Networks for Edge Devices. MLSys 2020】【58，Zihan Wang et al. Hierarchical memory-constrained operator scheduling of neural architecture search networks. DAC 2022】）、驱逐与重计算（重物质化【5，Olivier Beaumont et al. Efficient Combination of Rematerialization and Offloading for Training DNNs. NIPS 2021】【10，Tianqi Chen et al. Training deep nets with sublinear memory cost. 2016】【18，Audrunas Gruslys et al. Memory-efficient backpropagation through time. NIPS 2016】【24，Paras Jain et al. Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization. MLSys 2020】【27，Marisa Kirisame et al. Dynamic Tensor Rematerialization. ICLR 2021】【37，Shishir G. Patil et al. POET: Training Neural Networks on Tiny Devices with Integrated Rematerialization and Paging. ICML 2022】【38，Xuan Peng et al. Capuchin: Tensor-based gpu memory management for deep learning. ASPLOS 2020】）以及卸载与重载（交换【5，Olivier Beaumont et al. Efficient Combination of Rematerialization and Offloading for Training DNNs. NIPS 2021】【20，Mark Hildebrand et al. Autotm: Automatic tensor movement in heterogeneous memory systems using integer linear programming. ASPLOS 2020】【22，Chien-Chin Huang et al. Swapadvisor: Pushing deep learning beyond the gpu memory limit via smart swapping. ASPLOS 2020】【37，Shishir G. Patil et al. POET: Training Neural Networks on Tiny Devices with Integrated Rematerialization and Paging. ICML 2022】【38，Xuan Peng et al. Capuchin: Tensor-based gpu memory management for deep learning. ASPLOS 2020】【41，Jie Ren et al. ZeRO-Offload: Democratizing Billion-Scale Model Training. ATC 2021】【42，Minsoo Rhu et al. vDNN: Virtualized deep neural networks for scalable, memory-efficient neural network design. MICRO 2016】【57，Linnan Wang et al. Superneurons: dynamic GPU memory management for training deep neural networks. PPoPP 2018】）每个算子/张量，而不影响张量的形状。

  • 图变换：是一类通过改变计算图结构同时保持其语义来优化计算图的技术。现有工作【25，Zhihao Jia et al. TASO: optimizing deep learning computation with automatic generation of graph substitutions. SOSP 2019】【26，Zhihao Jia et al. Optimizing DNN Computation with Relaxed Graph Substitutions. MLSys 2019】【56，Haojie Wang et al. PET: Optimizing Tensor Programs with Partially Equivalent Transformations and Automated Corrections. OSDI 2021】【62，Yichen Yang et al. Equality Saturation for Tensor Graph Superoptimization. MLSys 2021】主要通过基于规则的子图替换来优化延迟，可分为两类：聚合变换（A-Trans），将小算子聚合成大算子以用内存换取延迟；中间变换（I-Trans），主要基于代数等价性为其他变换提供机会。

2.3 动机

  • 图变换的内存优化潜力：我们发现适当的图变换也可以改善图的内存使用。例如，如图2(c)所示，拆分算子以牺牲更多算子调用和降低硬件利用率为代价，减少了峰值内存使用。借助图变换，DNN的内存优化可以得到极大的增强。例如，在图2(a)中，这是一个在DNN训练或一些具有长跳跃连接的DNN【23，Gao Huang et al. Densely Connected Convolutional Networks. CVPR 2018】【44，Robin Rombach et al. High-Resolution Image Synthesis With Latent Diffusion Models. CVPR 2022】【73，Zongwei Zhou et al. Unet++: A nested u-net architecture for medical image segmentation. DLMIA 2018】【75，Barret Zoph et al. Learning Transferable Architectures for Scalable Image Recognition. CVPR 2018】中常见的简化图结构。其峰值内存使用量为1056，因为在计算第33个算子时，有33个大小为32的张量处于活动状态，这超过了100的内存限制。在图2(b)中，虽然仅使用图调度可以通过将暂时不用的张量交换到外部存储来将内存使用限制在100以内，但由于数据传输，这会导致较长的延迟。然而，如图2(d)所示，结合图变换可以节省更多内存，并可以利用异步交换来隐藏数据传输延迟。虽然硬件利用率下降，但在这种情况下，异步交换带来的效率提升可以补偿延迟损失。

  • 分裂变换 (F-Trans) 与挑战：我们将图1(c)和图2(c)(d)(e)中使用的变换命名为分裂变换（F-Trans），它是A-Trans的对偶操作，可以通过拆分算子有效优化内存使用。然而，现有的基于规则的子图替换的图变换技术【25，Zhihao Jia et al. TASO: optimizing deep learning computation with automatic generation of graph substitutions. SOSP 2019】【26，Zhihao Jia et al. Optimizing DNN Computation with Relaxed Graph Substitutions. MLSys 2019】【56，Haojie Wang et al. PET: Optimizing Tensor Programs with Partially Equivalent Transformations and Automated Corrections. OSDI 2021】【62，Yichen Yang et al. Equality Saturation for Tensor Graph Superoptimization. MLSys 2021】不能用于F-Trans。首先，F-Trans通常会大大增加图的复杂性，阻碍后续优化。其次，F-Trans涉及巨大的搜索空间，因为它可以应用于几乎任何子图。例如，图2(e)使用了两种不同的F-Trans，即使对于这个例子中如此简单的网络，可行的F-Trans的搜索空间也是巨大的。寻找有效表示和搜索F-Trans的方法是一个挑战。

  • 协同优化的必要性：此外，将图变换与图调度进行协调对于利用图变换优化内存使用至关重要。图2(c)表明，单独应用图变换可以优化内存使用，但过细的算子拆分可能会导致高昂的性能成本。相反，如图2(d)所示，结合图变换和图调度可以通过联合平衡变换和调度的内存与性能权衡，显著减少内存使用并实现更短的延迟。


图 2. 内存限制为100的动机示例。(a) 无任何优化。(b) 使用交换。(c) 使用分裂变换。(d)(e) 使用分裂变换和交换。

A2 方法细节

3 设计概览

  • MAGIS 总体设计：图3展示了MAGIS的总体设计。它接受一个DNN图作为输入，并输出优化后的图和调度方案。MAGIS包含四个主要组件：M-State、M-Analyzer、M-Rules和M-Optimizer。

  • 组件功能：

    • M-State：代表优化状态，包括计算图、分裂层次树（F-Tree）、最佳调度以及仿真与性能分析结果。F-Tree表示分裂变换（F-Trans）的层次化搜索空间，其中n=1的节点代表一个潜在的F-Trans子图和维度候选，而n>1的节点代表一个已经通过F-Trans沿某个维度分裂成n个部分的子图。
    • M-Analyzer：通过构建分裂层次树（F-Tree）来生成F-Trans的搜索空间。
    • M-Optimizer：协调图变换（包括F-Trans）和调度，以优化延迟和内存。
    • M-Rules：为M-Optimizer提供变换规则，包括先前工作中使用的“TASO规则”【25，Zhihao Jia et al. TASO: optimizing deep learning computation with automatic generation of graph substitutions. SOSP 2019】【26，Zhihao Jia et al. Optimizing DNN Computation with Relaxed Graph Substitutions. MLSys 2019】【56，Haojie Wang et al. PET: Optimizing Tensor Programs with Partially Equivalent Transformations and Automated Corrections. OSDI 2021】【62，Yichen Yang et al. Equality Saturation for Tensor Graph Superoptimization. MLSys 2021】、用于操作F-Tree以反映F-Trans在图上应用的F-Tree突变规则（§5.1），以及从图调度分解出的基于调度的规则。请注意，F-Trans被解耦为F-Tree和应用于F-Tree的突变规则。这些规则与其他规则（如TASO规则、基于调度的规则）集成在M-Rules中，形成一个由M-Optimizer探索的统一优化空间。

  • 工作流程：MAGIS以计算图为输入。M-Analyzer分析该图及其初始调度，构建F-Tree，输出初始的M-State并将其发送给M-Optimizer。M-Optimizer应用M-Rules通过改变某些子图或子F-Tree来产生新的M-States。需要注意的是，规则不会选择跨越受F-Trans影响的子图（属于F-Tree中n>1节点的子图）边界的子图进行变换。这是因为对于一个已经受F-Trans影响的区域S，规则不会变换与S部分相交的子图S'，因为S'的某些节点在执行时会被分割，而其他节点则不会。然后，它利用变换的图区域和先前的调度，对这些新图执行快速的增量调度，以迅速推导出接近最优的调度和相关的性能分析结果。有效的M-States会迭代地反馈给M-Optimizer。此外，如果一个图变换应用于一个未受F-Trans影响的子图，M-optimizer将查询M-Analyzer以更新新M-States中的F-Tree。

  • 章节结构：本文的其余部分结构如下：§4介绍MAGIS的M-Analyzer，§5讨论M-Rules，§6详细介绍M-Optimizer。


图 3. MAGIS 概览。

4 M-Analyzer

  • 引言：本节将首先介绍维度图（D-Graph），并用它来定义F-Trans。然后，我们提出F-Tree作为F-Trans优化空间/状态的抽象，并提供一种算法来构建一个轻量级的F-Tree，该算法仅考虑在基于支配树和内存热点选择的某些子图上进行F-Trans。
4.1 维度图

  • 维度图（D-Graph）定义：直观上，F-Trans沿着贯穿子图的“维度”进行分裂。因此，我们提出维度图（D-Graph）来识别图级别的维度。给定一个图G，其中v ∈ V(G)的输出张量有$d_v$个维度，其计算中有$r_v$个规约轴（reduce-axes），我们定义D-Graph $D = D(G)$，其中对于每个v ∈ V(G)和$d = -r_v, ..., -2, -1, 1, 2, ..., d_v$，都有一个节点⟨v, d⟩ ∈ V(D)。对于每个(u, v) ∈ E(G)，如果u的第$d_u$个维度和v的第$d_v$个维度对应于同一个空间轴，则存在(⟨u, $d_u$⟩, ⟨v, $d_v$⟩) ∈ E(D)；如果u的第$d_u$个维度对应于v计算的第$r_v$个规约轴，则存在(⟨u, $d_u$⟩, ⟨v, $-r_v$⟩) ∈ E(D)。例如，一个MatMul算子o（表示为$C[i, j] = Σ_k A[i, k] × B[k, j]$，其中i, j是C的维度，k是规约轴），其输入为A, B ∈ pre(o)，会提供连接(⟨A, 1⟩, ⟨o, 1⟩)、(⟨A, 2⟩, ⟨o, -1⟩)、(⟨B, 1⟩, ⟨o, -1⟩)、(⟨B, 2⟩, ⟨o, 2⟩) ∈ E(D)。

  • 维度图示例：图4(a)展示了从一个Transformer块【55，Ashish Vaswani et al. Attention is all you need. NIPS 2017】中提取的图G，其形状在(b)部分详述。 (c)部分描绘了D(G)的一些子图，例如一个包含除v1, v2, v3, v10之外所有张量的批处理维度（batch-dimensions）的子图，一个包含除v0, v12之外所有张量的头维度（head-dimensions）的子图，以及另一个包含除v1, v2, v3, v10之外所有张量的序列维度（sequence-dimensions）的子图。


(a) 自注意力机制的图G


(b) G中每个节点的形状


批处理维度


头维度


序列维度
(c) G的D-Graph的一些子图

图 4. D-Graph示例。N, T, C, H, h分别代表批大小、序列长度、隐藏维度、头数量、头维度。

4.2 分裂变换

  • F-Trans 定义与约束：借助D-Graph，我们可以将图G的一个F-Trans定义为$f = (S, D, n)$，其中$S ⊆ V(G)$，D是用于沿其分裂子图G[S]的D-Graph，n是分裂数量。它有以下约束：(1) G[S]是弱连通的。(2) G[S]是凸的：$G.inps(S) ∩ ∪_{v∈G.outs(S)} G.des(v) = ∅$。(3) 分裂后的图没有冗余计算，要求对于所有$v ∈ S$，存在唯一的整数$d$使得⟨v, d⟩ ∈ V(D)，并且对于所有$(u, v) ∈ E(G[S])$，存在整数$d_u, d_v$使得(⟨u, $d_u$⟩, ⟨v, $d_v$⟩) ∈ E(D)。

  • F-Trans 变换结果：给定G的一个F-Trans $f = (S, D, n)$，F-Trans之后的结果图是一个包含G[S]的n个分裂部分的图。对于所有$v ∈ G.inps(S)$，如果存在$d > 0$使得⟨v, d⟩ ∈ V(D)，那么v将被切片以供每个分裂部分使用，否则它将被所有分裂部分共享。对于所有$v ∈ G.outs(S)$，如果存在$d > 0$使得⟨v, d⟩ ∈ V(D)，那么v将通过合并分裂部分的相关输出来计算，否则通过规约它们来计算。注意，分裂部分是顺序执行的，以便通过及时释放每个部分的中间张量来节省内存，但这会因为算子形状变小而导致硬件利用率（例如，并行性、局部性）降低。

  • F-Trans 示例：图5展示了一个F-Trans $f = (S, D, n)$的例子，其中n = 2。v1是一个权重张量，因此没有⟨v1, d⟩ ∈ V(D)；所以在结果图中，v1被每个分裂部分共享。其他输入v0和v2被切片给每个部分。v8是v1的梯度，通过沿批处理维度相加计算，所以⟨v8, -1⟩ ∈ V(D)；因此在结果图中，v8是通过将每个分裂部分的输出相加来计算的。其他输出v6和v7是通过连接每个部分的输出来计算的。


图 5. 图G中的F-Trans f = (S, D, n) (n = 2)，该图是从一个MLP的训练图中简化的。(a) 子图 S = {v3, v4, v5, v6, v7, v8}。(b) D-Graph D，表示S的激活值的批处理维度。(c) F-Trans之后的结果图。

4.3 分裂层次树

  • F-Tree 概念：直接对图应用F-Trans将显著增加复杂性，特别是当分裂数很大时。由于每个F-Trans将一个图分成几个同构的子图，我们可以只保存其中一个。我们不直接变换原始图，而是构建一个分裂层次树（F-Tree）。F-Tree中的每个树节点记录一个F-Trans $f = (S, D, n)$。对于任何树节点$f = (S, D, n)$及其父节点$f' = (S', D', n')$，我们有$S ⊆ S'$。图3展示了一个F-Tree的例子，其中每个节点代表左侧图中被虚线框包围的子图，节点旁边的n是分裂数。当n = 1时，表示该节点是一个分裂候选；当n > 1时，表示该节点的子图已被F-Trans分裂成n个部分。这种抽象显著降低了后续图变换和调度的复杂性。

  • 搜索空间缩减：然而，图G上F-Trans的搜索空间仍然很大，可达$O(2^{|V(G)|^2})$，因为几乎任何凸子图都可以成为分裂候选。实际上，任意应用F-Trans并不能保证峰值内存的减少。只有当F-Trans作用于包含内存热点（§2.1）的子图时，才能实现有效的内存节省。

  • 分析：对于图G的一个F-Trans $f = (S, D, n)$，内存热点为H，$I = G.inps(S)$。$M_0$和$M_f$分别表示F-Trans前后的峰值内存使用量，如公式（1）所示。由于输入I在执行分裂子图时驻留在内存中，因此$M_f$应将其大小$Σ_{v∈I} |v|$与$Σ_{v∈H\S} |v|$（S之外的内存热点大小）合并为$Σ_{v∈(H\S)∪I} |v|$。F-Trans后的峰值内存减少量，即$M_0 - M_f$，如公式（2）所示。

$$ M_0 = \sum_{v \in U} |v| \quad M_f \approx \sum_{v \in (H \setminus S) \cup U} |v| + \sum_{v \in H \cap S} \frac{|v|}{n} $$
$$M_0 - M_f = \sum_{v \in H \cap S} \left(1 - \frac{1}{n}\right)|v| - \sum_{v \in F \setminus H} |v|$$

  • 度量标准：我们可以观察到，为了使$M_0 - M_f$更大，我们需要确保S包含更多的内存热点，而I消耗更少的内存。为了最小化F-Trans的输入内存使用，我们选择一个节点，并将其支配的子图作为分裂候选，以确保该子图只有一个入口节点。我们定义一个名为“内存热度”的度量，表示一个节点支配的子图中热点的总大小。给定图G、其支配树$T = T(G)$和内存热点H，我们用公式（3）计算v的内存热度，其中$H ∩ T.des(v)$是被v支配的内存热点。然后，我们为每个节点v分配一个分数，如公式（4）所示，用于估计在v支配的子图上进行F-Trans后潜在的峰值内存减少量。其中第一项是内存热点大小的减少，第二项是输入节点的大小，这些节点在F-Trans后执行每个分裂部分时应驻留在内存中。我们通常设置n = 2，以确保仅将子图分裂成两部分也能带来好处。

$heat(v) = \sum_{w \in H \cap T.des(v)} |w|$
$$ \text{score}(v) = \left(1 - \frac{1}{n} \text{heat}(v) - \sum_{u \in G, \text{mps}(T, \text{des}(v)), H} |u| \right) (4) $$

  • F-Tree 构建算法：基于上述度量，我们提出了算法1来构建F-Tree。其主要思想是识别分数（公式（4））分布在不同区间的节点，因为更高的分数意味着F-Trans能带来更多的峰值内存减少，但也可能意味着更大的延迟开销。超参数L控制区间的数量和F-Tree的最大层级。该算法输入图G和最大层级L，遍历D(G)的连通分量D（第3行），提取子图G'及其支配树T（第4-5行），然后基于公式（3）（4）计算分数（第6行）。在获得最大分数$S_{max}$（第7行）后，它将[0, 1]分割成L个区间，根据归一化分数$S[v]/S_{max}$在不同区间选择节点（第10-11行），并从这些节点支配的子图生成分裂候选（第12-15行）。F-Tree由这些子图构建而成。
算法 1: M-Analyzer: F-Tree 构建
输入: 图: G; 最大层级: L
输出: 分裂层次树: F
1 F := ∅;
2 H := MemoryHotspots(G);
3 for D ∈ D(G) 的连通分量 do
4     G' := 从 D 诱导的 G 的子图;
5     T := T(G');
6     S := GetScores(G', T, H);
7     S_max = max_{v∈V(G')} S[v];
8     if S_max ≤ 0 then continue;
9     for l ∈ {1, 2, ..., L} do
10        C := {v ∈ V(G&#39;) | l/L ≤ S[v]/S_max < (l + 1)/L};
11        for v_dom ∈ {v ∈ C | T.des(v) ∩ C = ∅} do
12            S := T.des(v_dom) \ {v_dom};
13            S&#39; := 从 S 诱导的 G&#39; 的子图;
14            f := (S&#39;, D, 1);
15            if f is valid then F := F ∪ {f};
16 return F;
  • F-Tree 构建示例：图6给出了一个F-Tree构建的例子，该计算图是从各种模型的训练图中简化的。为便于演示，这里我们只显示D(G)的一个连通分量D。部分(a)是算法1第4行的G'。部分(b)显示了支配树$T = T(G')$。部分(c)显示了基于公式（3）（4）计算的热度和分数结果。这里L=5，所以有5个归一化分数区间[0.2, 0.4), [0.4,0.6), [0.6,0.8), [0.8,1), [1,1]，虚线框中的节点是被选中的。部分(d)显示了被选为分裂候选的子图节点。部分(e)显示了最终构建的F-Tree。


图 6. 基于算法1（L=5）的F-Tree构建示例。每个张量的大小为1。(a) 算法1第4行的G'。(b) 支配树 Dom G' = T(G')。(c) 基于公式(3)(4)计算的分数，橙色框中的节点是选定的支配者（算法1第11行的v_dom）。(d) 选定的子图（算法1第12行的S）。(e) 构建的F-Tree。

5 M-Rules
  • 规则集合：MAGIS中的M-Rules借鉴了先前工作如TASO【25，Zhihao Jia et al. TASO: optimizing deep learning computation with automatic generation of graph substitutions. SOSP 2019】中的聚合变换（A-Trans）和中间变换（I-Trans）规则，如图1(a)(b)所示。我们将这些称为TASO规则，可用于优化延迟。除此之外，本节将介绍F-Tree突变规则和基于调度的规则，以进一步优化内存和延迟。
5.1 F-Tree 突变规则

  • 节点状态：由算法1构建的初始F-Tree的所有树节点$f = (S, D, n)$都具有n = 1。我们称它们为禁用节点，其子图尚未执行F-Trans。n > 1的节点称为启用节点，表示其子图已经执行了F-Trans并被分裂成n个部分。F-Tree突变规则主要改变F-Tree节点的n值，以将F-Trans应用于图。它们包括：

  • 规则列表：

    • 启用规则 (Enabling Rule)：它启用F-Tree的一个禁用的叶节点，或一个没有已启用祖先的已启用节点的父节点，如图7(a)所示。
    • 提升规则 (Lifting Rule)：它禁用一个没有已启用祖先的已启用节点，并启用其父节点，如图7(b)所示。
    • 禁用规则 (Disabling Rule)：它禁用一个没有已启用后代节点的已启用节点，如图7(c)所示。
    • 突变规则 (Mutating Rule)：它将一个已启用节点的Fission数n增加到下一个可以整除维度长度的数，如图7(d)所示。

  • 规则应用策略：借助M-Analyzer和上述规则，我们将F-Trans解耦为优化阶段之前的F-Tree构建和优化阶段期间的F-Tree突变。可以观察到，我们实际上首先启用叶节点，然后逐渐向更靠近根的节点移动。由于在靠近根的节点上应用分裂对内存和延迟的影响更大，我们从叶子开始进行更小的突变步骤和更平滑的搜索。


图 7. F-Tree突变规则图示。(a) 启用一个F-Tree节点。(b) 提升一个F-Tree节点。(c) 禁用一个F-Tree节点。(d) 增加分裂数n（维度长度d=12）。

5.2 基于调度的规则

  • 新增算子：我们引入了两个额外的算子，Store和Load，来表示图调度中的交换行为。在此基础上，我们添加了四条规则如下：

  • 规则列表：

    • 重物质化规则 (Re-materialization Rule)：它将算子A的一个用户B从多个用户中分离出来，让它使用一个重新计算的算子A'，如图8(a)(b)所示。
    • 反重物质化规则 (De-re-materialization Rule)：它是重物质化规则的对偶，将两个相同类型、相同输入的算子A和A'合并成一个单一的算子，如图8(c)(d)所示。
    • 交换规则 (Swapping Rule)：它在一个算子A和它的一个用户B之间插入Store和Load，表示A将被交换到外部存储，然后在B需要使用它时再交换回来，如图8(e)所示。
    • 反交换规则 (De-swapping Rule)：它是交换规则的对偶，移除两个算子之间的Store和Load，如图8(f)所示。

  • 调度分解：借助上述规则，我们可以将图调度分解为图变换和重排序，其中变换阶段决定哪些算子需要被重计算/交换，而重排序阶段决定何时进行重计算/交换。这样，内存和延迟之间的权衡可以转移到图变换阶段，图调度阶段只需考虑通常不影响总执行延迟的重排序。这种分解使得每次图变换后的调度变得简单得多。

  • 启发式策略：考虑到重物质化规则和交换规则可以应用于几乎任何算子，导致搜索空间巨大从而减慢优化速度，在实际的子图模式匹配过程中，可以有选择地应用这两条规则，过滤掉不包含内存热点的子图。


图 8. 基于调度的规则，表示从图调度分解出的变换。标有星号(*)的边代表零条或多条边。

6 M-Optimizer

  • 引言：在本节中，我们首先介绍增量调度，以利用突变的子图和先前的调度信息，高效地为变换后的图生成最优调度。然后，我们介绍顶层搜索算法，该算法根据内存和延迟对M-State进行优先级排序，并使用M-Rules变换当前最佳的M-State以生成新的M-State。
6.1 增量调度

  • 增量调度的必要性：为了获取一个图的内存使用和性能，我们需要执行图调度。在每次图变换后执行完整的图调度是昂贵的。为了解决这个问题，我们设计了一种增量调度算法，它根据先前的调度和上一次图变换影响的子图范围，确定需要重新调度的图的子集。这种方法使我们能够只对必要的子图进行调度，从而减少了调度的开销。

  • 算法描述：算法2展示了细节。它首先通过使用GetRescheduleInterval（第9行）获取原始图中需要重新调度的算子序列。接下来，在新的图中为该序列获取相应的子图$S_{new}$（第10行），然后使用GraphPartition将其划分为几个可以独立调度的子图（第11行）。每个子图的调度使用先前工作【3，Byung Hoon Ahn et al. Ordering Chaos: Memory-Aware Scheduling of Irregularly Wired Neural Networks for Edge Devices. MLSys 2020】中的基于动态规划的算法（第11行）进行，最后，将得到的调度结果组合起来形成新图的调度，并与原始图的调度集成（第12行）。

算法 2: M-Optimizer: 增量调度
输入: 旧图、新图: G_old, G_new;
      旧的突变子图节点: S_old;
      旧图的调度: ψ_old
输出: 新图的调度: ψ_new
1 function GetRescheduleInterval(G, S, ψ):
2   function ExtendBound(i, d):
3     ñ := 0; l := 0; v := ψ[i];
4     while l < 20 ∧ (ñ > 10 ∨ nw(v) < 4) ∧ nw(v) < ñ do
5       ñ := nw(v); i := i + d; v := ψ[i]; l := l + 1;
6     return i;
7   I_S := {i | i = 1, ..., |ψ| if ψ[i] ∈ S};
8   return ExtendBound(min I_S, -1), ExtendBound(max I_S, 1);
9 beg, end = GetRescheduleInterval(G_old, S_old, ψ_old);
10 S_new := V(G_new) \ (ψ_old[:beg] ∪ ψ_old[end:]);
11 Ψ := {DpSchedule(S) | S ∈ GraphPartition(S_new)};
12 return Merge(ψ_old[:beg], MergeSubSched(Ψ), ψ_old[end:]);

  • 关键函数 GetRescheduleInterval：GetRescheduleInterval是算法2中的一个关键过程，旨在找到原始调度中需要重新调度的区间。该区间不应太小，否则重新调度的结果将是次优的甚至不正确。同时，区间也不应太大，否则重新调度过程将消耗太多时间。在优化质量和时间成本之间进行权衡非常重要。

  • 窄腰（Narrow Waist, NW）值：我们引入一个节点的窄腰（NW）值nw(v)来解决这个问题。对于图G和节点v ∈ V(G)，nw(v)定义为$|V(G)| - |G.anc(v)| - |G.des(v)| - 1$，即$|V(G) \setminus G.anc(v) \setminus G.des(v)| - 1$。NW值可用于衡量与给定节点无关的节点数量。较低的nw(v)意味着更多的节点依赖于v，并且v依赖于更多的节点，这使得v成为拓扑排序问题的合适划分点。具体来说，所有v依赖的节点都应在v之前调度，所有依赖于v的节点都应在v之后调度，这为调度问题提供了自然的划分。此外，在我们分别为G.anc(v)和G.des(v)找到最优调度后，峰值内存消耗保证小于$M_{opt} + Σ_{u∈V(G)\G.anc(v)\G.des(v)} |u|$，其中$M_{opt}$表示在v的最优调度下达到的峰值内存。如果nw(v) = 0，则图的调度问题可以在v处被划分为两个完全独立的子问题。我们设计了一种基于NW值的启发式算法来选择区间，使其边界的NW值尽可能小（第2-6行），其中常数20、10、4是实践中表现良好的经验性超参数。GraphPartition背后的思想是使用nw(v) ≤ 1的节点作为划分点，将给定图的每个连通分量划分为多个子图。

6.2 顶层搜索算法

  • 搜索策略：MAGIS采用贪心搜索算法来优化图。MAGIS支持两种优化模式：在给定内存限制下优化延迟，或在给定延迟限制下优化内存。算法3显示了前一种模式的搜索算法。

  • 算法描述：算法3的输入包括图G、给定的内存限制M和F-Tree最大层级L。我们首先对给定图进行调度和分析，以获得一个初始的M-State（第9行）。然后，我们构建一个用于存储M-State的优先队列（第10行），其中优先级由BetterThan函数（第1-4行）确定，该函数在两个M-State都满足内存限制M时首先比较延迟；否则，比较内存（注意我们按字典序比较(μ, δ) < (μ', δ')）。然后我们迭代地弹出一个M-State μ（第12行），并应用M-Rules生成一系列新的M-State（第17行）。Analyze函数（第16行）如果μ先前突变的子图未受F-Trans影响，则会更新M-State μ中的F-Tree。我们对新生成的M-State μ'执行增量调度。然后，如果μ'在宽松条件下（由一个小的系数δ控制，经验上设为1.1）不比μ_best差，我们将其推入队列。为了防止冗余搜索，我们借鉴了Weisfeiler-Lehman测试【48，Nino Shervashidze et al. Weisfeiler-lehman graph kernels. JMLR 2011】的思想来哈希一个给定的图（第5-8行，第12-14行），其中⊕表示字节连接操作。

  • 性能评估：为了降低性能测量的开销，我们实现了一个带有算子性能缓存的模拟器。它保存算子的实际执行延迟，并使用模拟方法来获取整个图在某个调度下的整体性能和内存使用情况。当考虑异步交换时，涉及Store/Load算子的重排序也可能轻微影响延迟。为了解决这个问题，我们的重排序策略是尽可能早地放置Store，并尽可能晚地放置Load，以恰好隐藏数据传输延迟。

算法 3: M-Optimizer: 搜索算法
输入: 输入图 G; 内存约束 M; F-Tree 最大层级 L
输出: 优化的 M-State μ_best
1 function BetterThan(μ_1, μ_2, δ = 1):
2   return (max(μ_1.mem, M), μ_1.lat) < (max(δ × μ_2.mem, M), δ × μ_2.lat);
3 function GraphHash(G):
4   for v ∈ topo-order(G) do
5     x_v := hash(v) ⊕ (⊕_{u∈G.pre(v)} x_u);
6   return hash(Σ_{v∈G} x_v);
7 μ_best := InitState(G); X := ∅;
8 Q := PriorityQueue({μ_best}, BetterThan);
9 while Q ≠ ∅ do
10  μ := Q.pop(); x := GraphHash(μ.G);
11  if x ∈ X then continue;
12  X := X ∪ {x};
13  if μ&#39;s F-Tree needs update then
14    μ := Analyze(μ, l); # 算法 1
15  for μ&#39; ∈ ApplyTransformRules(μ) do
16    μ&#39; := ApplyIncrementalSchedule(μ&#39;); # 算法 2
17    if BetterThan(μ&#39;, μ_best) then μ_best := μ&#39;;
18    if BetterThan(μ&#39;, μ_best, 1.1) then Q.push(μ&#39;);
19 return μ_best;

A4 实验

实验环境 (总结)

  • 实现细节：MAGIS的图数据结构使用rustworkx【52，Matthew Treinish et al. rustworkx: A high-performance graph library for python. JOSS 2022】实现。我们实现了一个代码生成后端，根据图和调度生成调用PyTorch API的Python代码。使用PyTorch的CUDA Stream API实现异步的Store和Load操作，数据在GPU内存和CPU内存之间交换。
  • 对比基线 (Baselines)：
    1. PyTorch【36，Adam Paszke et al. PyTorch: An Imperative Style, High-Performance Deep Learning Library. NIPS 2019】：未优化的基线，仅进行简单的拓扑排序和即时张量回收。
    2. POFO【5，Olivier Beaumont et al. Efficient Combination of Rematerialization and Offloading for Training DNNs. NIPS 2021】：同时考虑重物质化和交换的SOTA工作，适用于结构简单的网络。
    3. DTR【27，Marisa Kirisame et al. Dynamic Tensor Rematerialization. ICLR 2021】：使用重物质化技术优化任意网络的SOTA工作，使用MegEngine【1】中的实现。
    4. XLA【46，Amit Sabne. Xla: Compiling machine learning for peak performance. 2020】：使用贪心重物质化算法的SOTA DNN编译器。
    5. TVM (Relay)【9，Tianqi Chen et al. Tvm: An Automated End-to-End Optimizing Compiler for Deep Learning. OSDI 2018】【43，Jared Roesch et al. Relay: A high-level compiler for deep learning. 2019】：执行基本内存节省的SOTA DNN编译器。
    6. Torch-Inductor (TI)【2，torch.compiler - PyTorch 2.1 documentation. 2023】：利用OpenAI Triton【50】的SOTA DNN编译器，执行基本内存节省。
  • 工作负载 (Workloads)：评估使用了多种DNN的训练过程，如表2所示，涵盖了从语言模型到视觉模型，从大型模型到小型模型的各种类型。数据类型对GPT-Neo和BTLM为bf16，其他为tf32。

表 2. 评估使用的工作负载

* 硬件平台：Intel工作站，配备20核Intel(R) Xeon(R) Silver 4210R CPU，一张NVIDIA GeForce RTX 3090 GPU。
* 软件配置：CUDA 11.6, cuDNN 8.4.0, PyTorch 2.1.0, MegEngine 1.12.3, TensorFlow 2.15.0, TVM 0.14.0。MAGIS优化时间预算为3分钟。为公平比较，所有基线都先应用TASO规则进行初步优化。

实验结果 (总结)

内存优化（带延迟约束）
  • 实验内容：评估在10%和5%延迟开销约束下，MAGIS和各基线的内存优化效果。
  • 实验结果 (图9)：
    • 在10%延迟开销下，MAGIS将峰值内存优化到PyTorch的15%~60%，优于所有基线。相比POFO、DTR和XLA，MAGIS的内存使用量是它们的15%~85%。TVM和TI效果接近PyTorch基线。
    • 在5%延迟开销下，MAGIS将峰值内存优化到PyTorch的25%~70%，同样优于其他基线。
  • 分析结论：
    • 对于结构简单的ResNet，MAGIS的优势不明显。
    • 对于内部结构复杂的BERT和ViT，MAGIS表现更好。
    • 对于跨层连接复杂的U-Net和U-Net++，MAGIS的优势最大，因为复杂的结构为图变换提供了更多优化空间。
    • 对于大型语言模型GPT-Neo和BTLM，基线大多出现OOM（内存不足），而MAGIS能成功运行并将内存优化到PyTorch的40%以下。


图 9. 与未优化的PyTorch相比的峰值内存比率（越低越好）。"OOM"表示内存使用超出实验平台内存限制。

延迟优化（带内存约束）
  • 实验内容：比较在80%和40%峰值内存限制下，MAGIS和其他工作的延迟优化效果。
  • 实验结果 (图10)：
    • 在80%内存限制下，MAGIS的延迟开销≤5%，优于POFO(≤40%)、DTR(≤15%)和XLA(≤20%)。
    • 在40%内存限制下，MAGIS的延迟开销≤15%，而POFO≤40%，DTR最高达70%，XLA最高达70%。许多基线在U-Net++、GPT-Neo和BTLM上无法满足40%的内存限制（标记为"FAILURE"）。
  • 分析结论：MAGIS在严格的内存约束下能实现更低的延迟。在U-Net上，相比DTR实现了1.25倍的加速。对于大模型，只有POFO能在40%内存限制下运行，但延迟远高于MAGIS。


图 10. 与未优化的PyTorch相比的延迟开销（越低越好）。"FAILURE"表示无法将内存比率优化到满足约束。

延迟与内存的权衡曲线
  • 实验内容：绘制各方法在不同约束下的内存-延迟权衡曲线。
  • 实验结果 (图11)：MAGIS的曲线在绝大多数情况下都位于其他基线曲线的下方。
  • 分析结论：MAGIS在内存和延迟的双目标优化中实现了更好的帕累托边界。这意味着，在给定延迟约束下，MAGIS能达到更低的内存消耗；或在给定内存约束下，能达到更低的延迟。当内存约束不紧张时，MAGIS的曲线斜率较低；但在极严格的内存限制下，曲线变得陡峭，因为即使是F-Trans也会因破坏局部性而产生巨大开销。


图 11. MAGIS与基线的延迟和内存曲线。MAGIS几乎在所有情况下都能实现帕累托最优。

与微批次（Micro-batching）的比较
  • 实验内容：为了验证图变换的有效性，将一种简单的F-Trans（即微批次）应用于ViT，然后将分割后的子图输入POFO进行优化。
  • 实验结果 (图12)：图变换确实增强了POFO在严格内存约束下的性能。不同的微批次因子在不同内存限制下表现最优，说明图变换和调度之间存在不同的权衡空间。
  • 分析结论：MAGIS由于能更好地协调变换和调度，其性能优于所有“微批次+POFO”的组合。


图 12. MAGIS与POFO的比较。POFO使用的网络已经过微批次预处理（使用不同因子）。

启发式策略消融实验
  • 实验内容：对MAGIS中的三个主要启发式策略进行消融研究：H1) F-Tree构建，H2) 基于调度的规则启发式，H3) F-Tree最大层级L。
  • 实验结果 (图13)：
    • 禁用H1（随机选择F-Trans）的naïve-fission性能最差。
    • 禁用H2（全图匹配调度规则）的naïve-sch-rules搜索收敛慢。
    • max-level=4（默认设置）表现最好。max-level=2因搜索空间太小而潜力不足，max-level=8因搜索空间太大而难以在有限时间内找到优解。
  • 分析结论：论文提出的启发式策略对于高效搜索和取得良好优化结果至关重要。


图 13. 在3分钟内优化BERT工作负载时MAGIS的启发式策略分解，约束条件与§7.2.1和§7.2.2相同。曲线上的菱形"⋄"是其优化结果满足约束的时间点。曲线上的方形"□"是获得最佳优化结果的时间点。

增量调度评估
  • 实验内容：比较增量调度（IS）与完整调度（FS）在10个随机生成的类NASNet网络上的调度速度和结果质量。
  • 实验结果 (图14)：IS相比FS实现了4~30倍（平均10倍）的加速。在100次测试中，有94次IS达到了与FS相同的优化质量。
  • 分析结论：增量调度在大幅降低调度时间开销的同时，几乎不损失优化质量。


图 14. 增量调度（IS）与完整调度（FS）的比较。(a) IS相对于FS的调度时间加速比。(b) IS与FS的调度结果质量比较。

优化时间成本
  • 实验内容：分析MAGIS优化ViT（batch 64）1分钟的时间成本构成。
  • 实验结果 (图15)：在60秒的总时间中，变换（Trans.）耗时2.52秒，调度（Sched.）耗时3.7秒，模拟（Simul.）耗时8.71秒，哈希测试（Hash）耗时44.82秒。
  • 分析结论：哈希测试有效过滤了大量重复图，是总开销的主要部分，但它显著减少了其他过程的开销。

图 15. MAGIS优化ViT (batch 64) 1分钟的优化时间成本分解。"Filtered"表示被哈希测试过滤掉的重复图。

案例研究：UNet
  • 实验内容：展示UNet在PyTorch、MAGIS-1（内存限制80%）和MAGIS-2（内存限制60%）下的执行时间和内存使用曲线。
  • 实验结果 (图16)：PyTorch和MAGIS-1的内存曲线呈现先增后减的趋势。MAGIS-1通过重物质化、交换和F-Trans降低了峰值内存，但延迟增加。MAGIS-2由于应用了覆盖整个图的F-Trans，出现了双内存峰值，进一步降低了峰值内存，但延迟开销更大。
  • 分析结论：该案例直观地展示了MAGIS如何通过不同的变换和调度策略在内存和延迟之间进行权衡。


图 16. UNet的执行时间与内存使用情况。

A7 相关工作

  • 引言：本节简要介绍MAGIS的相关工作，主要包括图调度技术（重物质化、交换和重排序）以及图变换。本节还讨论了一些其他的DNN编译器。

  • 重物质化 (Re-materilization)：重物质化技术通过丢弃一些中间张量，在需要时再重新计算它们。该技术最早由【10，Tianqi Chen et al. Training deep nets with sublinear memory cost. 2016】【17，Andreas Griewank et al. Algorithm 799: revolve: an implementation of checkpointing for the reverse or adjoint mode of computational differentiation. TOMS 2000】【18，Audrunas Gruslys et al. Memory-efficient backpropagation through time. NIPS 2016】应用于深度学习。【28，Ravi Kumar et al. Efficient Rematerialization for Deep Networks. NIPS 2019】【29，Mitsuru Kusumoto et al. A Graph Theoretic Framework of Recomputation Algorithms for Memory-Efficient Backpropagation. NIPS 2019】中使用了图论分析。Checkmate【24，Paras Jain et al. Checkmate: Breaking the Memory Wall with Optimal Tensor Rematerialization. MLSys 2020】使用整数规划（IP）进行优化。DTR【27，Marisa Kirisame et al. Dynamic Tensor Rematerialization. ICLR 2021】使用启发式策略优化动态图的重物质化。MONeT【47，Aashaka Shah et al. Memory Optimization for Deep Networks. ICLR 2022】协同优化重物质化和算子实现。

  • 交换 (Swapping)：交换技术将一些张量存储在外部存储中，在需要时再加载回来。vDNN【42，Minsoo Rhu et al. vDNN: Virtualized deep neural networks for scalable, memory-efficient neural network design. MICRO 2016】、Capuchin【38，Xuan Peng et al. Capuchin: Tensor-based gpu memory management for deep learning. ASPLOS 2020】和SuperNeurons【57，Linnan Wang et al. Superneurons: dynamic GPU memory management for training deep neural networks. PPoPP 2018】将其用于GPU上的DNN训练。SwapAdvisor【22，Chien-Chin Huang et al. Swapadvisor: Pushing deep learning beyond the gpu memory limit via smart swapping. ASPLOS 2020】协同优化重排序、内存分配和交换。TFLMS【30，Tung D. Le et al. Tflms: Large model support in tensorflow by graph rewriting. 2019】通过特殊算子和控制流边来表示交换。POET【37，Shishir G. Patil et al. POET: Training Neural Networks on Tiny Devices with Integrated Rematerialization and Paging. ICML 2022】使用IP结合重物质化和交换在移动设备上进行训练。POFO【5，Olivier Beaumont et al. Efficient Combination of Rematerialization and Offloading for Training DNNs. NIPS 2021】使用动态规划（DP）结合重物质化和交换。ZeRO-Offload【41，Jie Ren et al. ZeRO-Offload: Democratizing Billion-Scale Model Training. ATC 2021】将交换与分布式训练相结合。AutoTM【20，Mark Hildebrand et al. Autotm: Automatic tensor movement in heterogeneous memory systems using integer linear programming. ASPLOS 2020】和ZeRO-Infinity【39，Samyam Rajbhandari et al. Zero-infinity: Breaking the gpu memory wall for extreme scale deep learning. SC 2021】使用持久内存作为外部存储。

  • 重排序 (Re-ordering)：重排序寻找合适的DNN拓扑序来优化内存。Serenity【3，Byung Hoon Ahn et al. Ordering Chaos: Memory-Aware Scheduling of Irregularly Wired Neural Networks for Edge Devices. MLSys 2020】使用DP进行优化。SwapAdvisor【22，Chien-Chin Huang et al. Swapadvisor: Pushing deep learning beyond the gpu memory limit via smart swapping. ASPLOS 2020】同时考虑重排序和交换。HMCOS【58，Zihan Wang et al. Hierarchical memory-constrained operator scheduling of neural architecture search networks. DAC 2022】分层搜索最优顺序。Zhong等人【72，Shuzhang Zhong et al. Memory-aware Scheduling for Complex Wired Networks with Iterative Graph Optimization. ICCAD 2023】使用带变量剪枝的IP来加速优化。

  • 图变换 (Graph Transformation)：图变换源于编译器的超级优化【34，Henry Massalin. Superoptimizer: a look at the smallest program. ASPLOS 1987】。它通过每步突变一个子图来逐步优化图。MetaFlow【26，Zhihao Jia et al. Optimizing DNN Computation with Relaxed Graph Substitutions. MLSys 2019】使用回溯算法，TenSAT【62，Yichen Yang et al. Equality Saturation for Tensor Graph Superoptimization. MLSys 2021】采用等价饱和【60，Max Willsey et al. egg: Fast and Extensible Equality Saturation. POPL 2021】进行搜索。TASO【25，Zhihao Jia et al. TASO: optimizing deep learning computation with automatic generation of graph substitutions. SOSP 2019】基于程序合成自动生成变换规则。PET【56，Haojie Wang et al. PET: Optimizing Tensor Programs with Partially Equivalent Transformations and Automated Corrections. OSDI 2021】提出了部分等价变换。Unity【54，Colin Unger et al. Unity: Accelerating DNN Training Through Joint Optimization of Algebraic Transformations and Parallelization. OSDI 2022】将分布式并行优化集成到图变换中。Turner等人【53，Jack Turner et al. Neural Architecture Search as Program Transformation Exploration. ASPLOS 2021】将图变换与神经架构搜索相结合。与先前工作相比，MAGIS可以权衡延迟和内存优化。在变换类型方面，MAGIS研究了分裂变换的形式化和搜索。我们还提出了源自图调度的重物质化和交换规则，增强了图变换和调度之间的协调。

  • 其他DNN编译器：除了上述工作，近年来还提出了许多其他DNN编译器【2，torch.compiler - PyTorch 2.1 documentation. 2023】【8，Jingwei Cai et al. Inter-layer Scheduling Space Definition and Exploration for Tiled Accelerators. ISCA 2023】【9，Tianqi Chen et al. Tvm: An Automated End-to-End Optimizing Compiler for Deep Learning. OSDI 2018】【11，Tianqi Chen et al. Learning to optimize tensor programs. NIPS 2018】【14，Yaoyao Ding et al. Ios: Inter-operator scheduler for cnn acceleration. MLSys 2021】【16，Siyuan Feng et al. TensorIR: An Abstraction for Automatic Tensorized Program Optimization. ASPLOS 2023】【31，Ao Li et al. Automatic Horizontal Fusion for GPU Kernels. CGO 2020】【33，Lingxiao Ma et al. Rammer: Enabling Holistic Deep Learning Compiler Optimizations with rTasks. OSDI 2020】【43，Jared Roesch et al. Relay: A high-level compiler for deep learning. 2019】【46，Amit Sabne. Xla: Compiling machine learning for peak performance. 2020】【49，Yining Shi et al. Welder: Scheduling Deep Learning Memory Access via Tile-graph. OSDI 2023】【50，Philippe Tillet et al. Triton: an intermediate language and compiler for tiled neural network computations. MAPL 2019】【59，Jian Weng et al. UNIT: Unifying Tensorized Instruction Compilation. CGO 2021】【61，Jiarong Xing et al. Bolt: Bridging the Gap between Auto-tuners and Hardwarenative Performance. MLSys 2022】【63–71】【74，Hongyu Zhu et al. ROLLER: Fast and Efficient Tensor Compilation for Deep Learning. OSDI 2022】。例如，AutoTVM【11】、FlexTensor【70】、Ansor【64】和Roller【74】自动生成/探索单个算子或小子图的调优空间；UNIT【59】、AMOS【66】和TensorIR【16】自动将算子映射到具有专用张量指令的硬件加速器上；Rammer【33】、HFuse【31】和IOS【14】融合并行算子以提高硬件利用率；DNNFusion【35】、AStitch【71】和Apollo【63】融合链式算子以减少数据移动；BOLT【61】、Chimera【69】、SET【8】、TileFlow【67】和Welder【49】还探索了计算密集型算子的融合空间。

A5 结论

我们提出了MAGIS，一个用于内存和延迟优化的DNN优化器，其系统化地设计了分裂变换，并实现了图变换和调度之间的有效协调。实验结果表明，与最先进的方法相比，MAGIS在相同的延迟约束下仅使用15%~85%的内存，并获得了更好的内存与延迟帕累托边界。