MagPy: Compiling Eager Mode DNN Programs by Monitoring Execution States

文章标题：MagPy: 通过监控执行状态编译Eager模式的DNN程序
作者/机构：Chen Zhang, Rongchao Dong, Haojie Wang, Runxin Zhong, Jike Chen, and Jidong Zhai, 清华大学

A1 主要贡献

核心问题：
现实世界中的深度学习程序通常使用像Python这样的动态编程语言开发，这些语言具有复杂的特性（如内置函数、动态类型），并以Eager模式执行，导致性能不佳。深度学习编译器依赖于基于算子的计算图来优化程序，但动态语言的复杂性常常阻碍了程序被完整地转换为算子图，从而导致次优性能。现有的图实例化技术分为两类：基于分析的方法（如TorchDynamo）和基于追踪的方法（如LazyTensor）。基于分析的方法难以支持所有Python特性，导致图碎片化和高昂的解释器开销；基于追踪的方法要么依赖用户保证图的静态性（可能引入静默错误），要么为每次执行重新追踪，开销巨大。如图1所示，现有技术与手动构建完整图的性能差距显著。

研究目标：
本文旨在提出一种新的方法，通过监控Python解释器的执行状态来增强算子图的实例化，从而更有效地将用户编写的、易于使用的Python程序自动转换为编译器友好的完整算子图，以提升深度学习程序的执行性能。

创新点：
本文基于以下三个核心洞察提出了MAGPY系统：
1. 深度学习程序的有限动态性：大多数深度学习程序虽然用Python编写，但其算子图结构（如算子顺序、属性、形状）在不同批次间通常保持不变。因此，关键问题从“图是什么”转变为“图何时会改变”。
2. 外部值决定程序行为：程序的行为（包括算子图和副作用）仅受从外部读取的值（如输入参数、全局变量）的影响。只要这些外部值不变，且所有操作都是确定性的，程序行为就不会改变。因此，只需验证外部读取值是否变化，即可判断图是否可重用，这大大降低了监控的复杂性。
3. 基于执行状态生成Guard和Mock：用于验证输入的Guard函数和用于复现上次运行最终状态的Mock函数，都可以通过分析程序运行时的状态（特别是变量间的引用关系）来确定，而无需深入理解程序逻辑。

基于以上洞察，MAGPY做出了以下主要贡献：
1. 识别了深度学习程序的“有限动态性”，这一特性使得高效的图实例化成为可能。
2. 提出了RefGraph（引用图），用于在程序执行期间记录程序状态并监控潜在的动态性。这显著降低了动态分析的复杂性，为实现更好的图实例化提供了便利。
3. 提出了一种基于注解的方法来处理编程语言的复杂性（包括内置函数、用户自定义函数等），以更好地生成RefGraph。
4. 设计并实现了MAGPY系统，一个用于Python深度学习程序的高效算子图实例化系统。实验证明，MAGPY相比现有技术最高可实现2.88倍（平均1.55倍）的加速，并成功将ParityBench中93.40%的有限动态模型实例化为完整的算子图。

A3 背景知识与挑战

2.1 挑战

将深度学习程序实例化为算子图面临的挑战。深度学习程序的算子图实例化之所以充满挑战，是因为Python和PyTorch为了易于编程而提供了大量灵活的特性。编译器必须正确处理所有这些特性才能生成正确的算子图。通过对ParityBench中272个TorchDynamo无法导出完整算子图的真实用户程序进行手动分析，我们总结出以下主要障碍：
- 动态数据类型：Python是动态类型语言，表达式的类型只能在运行时确定。图3中的示例程序没有类型注解，编译器无法预知哪些操作作用于张量并应被插入到算子图中。几乎所有用户程序都不提供变量的精确数据类型。
- 大量的内置API：Python和PyTorch提供了丰富的API来操作对象，如图3a所示，包括高级容器操作（#1）、灵活的迭代器（#2）和张量元数据处理（#3）。编译器需要知道这些API的精确结果才能构建算子图，但逐一实现数量庞大的内置API是不现实的。47.2%的Dynamo失败源于不支持的内置API。
- 复杂的面向对象编程（OOP）：用户倾向于使用OOP风格定义新类来封装模型的重复部分，如图3b所示。由于Python是动态解释型语言，编译器很难找到实际的对象和函数。34.7%的Dynamo失败是由于处理用户自定义类的能力有限。
- 副作用：变量的可变性在Python程序中普遍存在。如图3c所示，forward函数在第2行写入self.dim属性，在第4行读取该属性。这些修改会改变后续函数调用的行为，必须被正确处理。6.3%的Dynamo失败由此引起。
- 语义鸿沟：Python编程语言与算子图表示之间存在巨大的语义鸿沟。为了追求高性能，图级编译器偏爱简单且静态的图表示。因此，上述复杂的Python特性无法在算子图中表示，必须在图实例化过程中被消除。
尽管存在上述挑战，用户程序的精确行为在不同批次间通常是相同的。因此，只要能够验证程序行为保持不变，MAGPY就可以安全地重用前一次运行的算子图，而无需理解这些程序的精确语义。

2.2 基于Guard的即时编译

利用JIT编译对程序进行特化。鉴于许多变量在不同批次间保持不变，深度学习程序可以通过即时编译（JIT）对这些不变的变量进行特化。例如，当使用一个2x2的Tensor x和run_act=True调用图4(a)中的model函数时，JIT编译器会将程序特化为图4(b)中的一个mock函数。这个特化后的mock函数会执行一个没有run_act分支的静态算子图，并复现model函数中的副作用（更新全局变量z）。为了保证特化的正确性，JIT编译器还会生成一个guard函数来验证特化程序的假设。当新输入能够通过guard检查时，特化后的mock函数将产生与原始用户程序相同的结果。我们将guard函数、mock函数和算子图的集合称为一个“记录”（record）。

JIT系统的工作流程。如图5所示，编译好的记录保存在JIT系统的缓存中。当被编译的函数被调用时，JIT系统首先在缓存中搜索一个其guard能够成功通过的记录。如果找到匹配的记录，JIT系统将通过调用与该记录关联的mock来替代原始函数调用。否则，JIT系统将重新编译用户程序以生成一个新的记录，并调用图级深度学习编译器来优化该记录中的算子图。然后，该记录被保存到缓存中以备将来使用。

A2 方法细节

3 概览

MAGPY系统概览。MAGPY的概览如图6所示。当在缓存中找不到匹配的记录时，MAGPY会通过使用本地语言执行器（如Python解释器）执行程序并监控其执行过程来重新编译用户程序。我们称此过程为用户程序的“监控运行”。与传统编译器基于控制流图等结构分析程序逻辑不同，MAGPY只关心程序状态。MAGPY提出了一个RefGraph（引用图）来保存运行时状态信息，主要是用户程序运行时变量之间的引用关系。在程序执行期间，MAGPY从本地执行器捕获每条指令的执行状态，并根据捕获的执行状态更新RefGraph。当程序执行完毕后，MAGPY通过分析RefGraph生成一个新的记录，并将其存入记录缓存中。

4 设计

本节将介绍RefGraph的定义（§4.1），以及基于RefGraph生成guard、mock和算子图的算法（§4.2-§4.3）。

4.1 RefGraph的定义

RefGraph的构成。RefGraph的定义如图8所示，它包含以下元素：
- RefGraph：RefGraph保存运行时信息，以ShadowNodes（与每个运行时变量相关联）为图节点，以运行时变量的引用关系为图边。MAGPY可以使用运行时变量或ShadowNode ID来检索ShadowNode。图7展示了MAGPY在监控图9中深度学习程序执行时RefGraph的更新过程示例。RefGraph在程序执行期间被更新，其详细过程将在后续部分解释。
- ShadowNode：每个运行时变量（包括Tensor和非Tensor变量）在RefGraph中由一个ShadowNode表示。引入了一个特殊的ShadowNode SN #0来保存外部运行时的状态，例如Python中的全局变量和参数。ShadowNodes仅在解释器显式访问运行时变量时才被懒惰地创建。通过这种方式，MAGPY可以消除为所有变量创建ShadowNodes的开销，并避免在生成guard时对未使用的变量进行过度特化。图7(1)显示了在运行时从外部读取x[0]和y之后，执行x[0] += y之前的RefGraph。MAGPY仅为三个使用过的变量x (SN #1)、x[0](SN #2)和y (SN #3)创建了ShadowNodes，而忽略了未使用的参数z。
- ShadowVersion：MAGPY引入ShadowVersion来保存变量的更新历史。更新历史有助于为生成guard恢复初始输入值，并为生成mock找到原地更新操作。尽管MAGPY只需要运行时变量的初始和最终状态，但为了使解释更流畅，本文中我们假定MAGPY记录了变量的所有版本。当程序运行时发生原地操作时，会创建一个新版本。对于每个版本，MAGPY记录该版本的值以及它所持有的引用的变量。例如，步骤(2)中的原地加法x[0] += y为x[0] (SN #2)创建了一个新版本，步骤(3)中的x.append()为x (SN #1)创建了一个新版本。
- 引用关系：许多复合变量持有对其他变量的引用。例如，Python中的列表持有其元素的引用。引用关系影响了需要被保护（guard）以保持不变的变量集合以及在mocking期间需要复现的变量集合。因此，MAGPY显式地保存这些信息，通过从复合变量的特定ShadowVersion指向被引用变量的整个ShadowNode的边来表示。这些边保存在ShadowVersion的ref字段中。详细的引用关系作为边的属性保存。为简化起见，图7仅描绘了每一步中新创建边的详细关系。MAGPY只需要引用的状态，这可以通过检查运行时变量轻松获得，而无需关心这种状态是如何形成的。例如，在图7的步骤(3)中，MAGPY可以通过观察运行时x的值知道列表x包含作为x[0]的SN #2和作为x[1]的SN #4，而无需知道Python中的append是用于向列表中添加新元素的；在步骤(4)中，尽管元组(SN #5)是从列表x(SN #1)创建的，但这两个ShadowNodes是独立的，没有像传统数据流图中那样的边。从变量x到其所引用目标y的引用关系可以被懒惰地生成，如果目标只能被显式访问，例如用户定义变量的属性。然而，如果目标变量可能被隐式访问，比如容器的包含变量，那么在访问x时就应该生成这种关系。
- 算子图：除了生成的RefGraph，MAGPY还记录了已执行程序的算子图，其中包含所有访问张量数据的Tensor操作，如图10所示。RefGraph中的ShadowNode ID和version ID也为每个Tensor记录下来。通过使用RefGraph，MAGPY可以知道从哪里加载输入Tensor以及将输出Tensor存储到哪里。RefGraph记录了生成用于验证的guard函数和用于正确调用已记录算子图的mock函数所需的所有信息。由于RefGraph主要保存运行时变量的精确引用关系，MAGPY可以通过检查执行器的运行时状态来构建图，而无需过多了解复杂的语言语义。

4.2 Guard和Mock的生成

基于RefGraph搜索关键变量。由于ShadowNode和引用关系的懒惰创建，MAGPY中的RefGraph自然地筛选出了对程序输出有实际影响的关键变量集合。然后，MAGPY使用算法1通过在RefGraph上搜索来收集用于guard和mock的关键变量。

Guard生成。guard的目标是验证程序开始时初始状态是否未改变。它需要检查运行时在监控运行期间从外部状态显式读取的变量。所有这些变量在函数入口处都存在，并且可以通过从外部状态（SN #0）仅通过版本0的引用关系在RefGraph中到达。MAGPY通过算法1中的GetGuardNodes收集这些节点。图9示例代码的待保护变量和guard函数分别显示在图11a和图11c中。

Mock生成。mock的目标是在不执行用户程序的情况下，复现监控运行的最终程序状态。mock需要复现所有可能在编译区域之外使用的变量。这些变量可分为两类，都可以通过从外部状态SN #0在RefGraph中搜索来收集，如算法1中的GetMockNodes所示：
1. 在程序执行期间创建并显式存储到外部的变量，例如示例程序中的a。这些变量可以通过从外部节点SN #0开始，使用最后一个版本中的引用关系作为边进行搜索来收集（算法1的第16、19-21行）。然而，如果一个变量在函数入口处已存在且未被原地更新（其ShadowNode中只有一个版本），MAGPY可以跳过复现它，而是使用其初始值。
2. 在函数入口处存在并被原地修改的变量，例如示例程序中的x和x[0]。这些变量可能会在编译区域之外被访问，所以即使在最后一个版本的引用关系中没有从外部节点SN #0到这些变量的路径，它们仍然需要在mock中被更新。这些变量通过算法1的第17、22-24行收集。

mock将首先调用由图级编译器加速的算子图以获取输出Tensors（详见§4.3）。然后，mock将像在监控运行中一样复现变量。类型1的变量可以从头创建，类型2的变量需要被原地更新。图11b和图11d分别显示了示例代码收集到的节点和生成的mock代码。

指针别名分析。指针别名分析是编译器设计中的一个挑战。然而，在MAGPY中，指针别名关系自然地在RefGraph中表示出来，即从SN #0到达特定节点的不同路径。除了上述的值守卫（value guards），MAGPY还验证指针别名关系与监控运行相匹配。这可以通过检查RefGraph中所有到达同一节点的路径是否都到达同一个变量来实现。mock也需要复现引用关系，这可以通过简单地复现RefGraph中的引用关系来实现。

4.3 算子图生成

确定算子图的输入和输出。对于记录的算子图，MAGPY将确定输入和输出节点，以及这些张量在mocking期间应从运行时加载或存储到何处。输入节点是算子图中没有入边的Tensor节点。这些Tensor变量不是由张量操作创建的，因此它们应该在函数入口处就存在。因此，在RefGraph中存在从SN #0使用版本0的边到相应节点的引用路径。输出节点是需要被mock的Tensor变量，并在mock生成期间确定。输出节点也可以通过从RefGraph中的SN #0搜索得到，其路径表示输出Tensors的存储位置。MAGPY可以确保中间节点将来不会被使用，并将此信息传递给图级编译器。因此，图级编译器可以安全地执行诸如死代码消除或算子融合等优化。然而，基于追踪的框架无法实现这一点，因为它们无法知道一个Tensor将来是否会被使用。

5 RefGraph生成

RefGraph的生成过程。MAGPY通过分析监控运行的中间运行时状态来生成RefGraph。具体来说，MAGPY在监控运行期间捕获每条指令的执行状态（§5.1），获取该指令预定义的操作属性（§5.2），并根据该属性使用不同的RefGraph更新规则（§5.3）来更新RefGraph。

5.1 执行状态捕获接口

从运行时收集信息。动态语言通常会保存运行时变量的高级信息（如类型和变量结构）以进行动态解释。这些信息对MAGPY分析运行时状态非常有价值。因此，MAGPY从运行时收集每条指令的执行状态，其中包含该指令所有相关的变量。执行状态包含以下元素：
- OP：将要执行的操作，如加载全局变量或将两个变量相加。
- input[x]：按索引x访问的操作输入。例如，函数调用的输入是其参数列表，赋值操作的输入是将要存储的值。
- output[x]：按索引x访问的操作输出。例如，函数调用的输出是其返回值。

5.2 操作属性注解

处理原生代码实现的复杂性。出于性能考虑，动态语言使用原生机器码来实现某些操作。例如，Python用C语言实现了一些函数。这些低级机器码不保留高级变量信息，难以分析。为了收集所有算子并正确运行MAGPY，需要对这些函数进行属性注解。请注意，这些只是注解，比基于分析的方法中重新实现这些函数要容易得多。MAGPY也鼓励对动态语言实现的常用函数进行注解，以便MAGPY可以跳过深入函数实现细节，从而加快监控运行速度。所需的注解如下：

函数属性 (Function Properties) 为函数注解：
- AsOpNode：一个将操作转换为算子图的算子节点的函数。如果操作不能表示为算子节点，则返回None。
- PureClosure：一个布尔属性，如果对于相同的输入参数，返回值相同，并且没有像隐式更新全局变量那样的外部效应，则可以注解为True。当PureClosure为True时，允许对输入参数进行原地更新。

AsOpNode用于生成算子图。PureClosure识别Python中的动态行为，如系统调用和随机数生成器，它们的注解将为False。

输入属性 (Input Properties) 为函数的每个输入变量注解：
- ValueRead：一个布尔属性，当函数读取该输入变量的值时为True。
- InPlace：一个布尔属性，指示该变量可能被函数原地修改。

例如，像add或sub这样的数学运算的所有参数都将被注解为ValueRead=True。对于向列表中添加变量的函数list.append，第一个参数（列表）被注解为ValueRead=True，因为它的内部结构被函数访问。相反，第二个参数（要添加的变量）被注解为ValueRead=False，因为函数只将其引用添加到列表中，而不访问其值。此外，第一个参数被注解为InPlace=True，因为列表被指定的值原地修改了，而第二个参数被注解为InPlace=False。

输出属性 (Output Properties) 为每个输出变量注解：
- InPlace：一个布尔属性，指示该变量可能被函数原地修改。
- Relation：标记需要在输入和输出节点之间创建的关系，以及详细的关系内容。需要进一步指定这是通过现有引用进行的读取（例如，读取v.x）还是创建新引用的写入（例如，赋值v.x = 1.0）。

注解的用户工作量。进行此类注解的用户工作量是可承受的。尽管总共有6个属性，但只有操作的PureClosure以及输入和输出变量的InPlace被频繁注解。AsOpNode只需为图级编译器支持的有限数量的操作进行注解。ValueRead仅在函数可能读取Tensor时才必要，这是由于§5.3中的图更新规则。Relation仅用于特殊的懒惰创建引用关系，如显式加载属性。此外，许多操作具有相同的注解，例如+、-、×、÷操作，这允许批量生产注解。

5.3 图更新规则

更新RefGraph和算子图的流程。MAGPY首先为输入和输出变量中没有对应节点的变量生成ShadowNodes。然后，MAGPY检索操作的属性注解，并根据注解和捕获的执行状态更新RefGraph和算子图。

5.3.1 算子图更新和动态性检测

处理不同类型函数以更新图和检测动态行为。AsOpNode、PureClosure和ValueRead注解用于更新算子图和检测动态性。其工作流程如图12所示。
1. 注解检索。MAGPY首先尝试检索函数的注解。虽然大多数内置函数和常用库函数都提供了注解，但某些函数（如用户自定义函数）可能缺少注解。MAGPY根据是否存在注解进入步骤(2)或(3)。
2. 处理无注解函数。如果找不到注解，MAGPY将函数调用视为无条件跳转到该函数，并尝试内联该函数（$0）。如果函数代码不可用（例如，用原生语言实现），MAGPY会将该函数视为动态函数并停止图实例化（$1）。
3. 使用AsOpNode生成算子图。如果AsOpNode成功为函数生成一个算子节点（例如，torch.relu），MAGPY会将该节点添加到算子图中（$2）。否则，跳转到(4)。 4. **使用PureClosure检测动态函数**。如果`PureClosure`为False，MAGPY会将该函数视为动态函数并停止图实例化（$3）。如果为True，MAGPY将假定该函数的行为在不同批次中保持不变，并跳转到(5)。
5. 使用ValueRead检测动态张量计算。MAGPY扫描输入参数，以查找是否存在一个输入位置被标记为ValueRead=True且其运行时值为Tensor。如果找到这样的参数，MAGPY会将该函数视为动态函数并停止图实例化（$4）。原因是MAGPY允许Tensor数据在不同批次中发生变化，因此这些操作会因读取Tensor数据而产生不同的结果。典型情况是图级编译器不支持的张量操作，如Inductor的nn.LSTM。这些操作未被AsOpNode转换为算子节点，因此会进入此分支。否则，操作仅读取非张量变量的值（如list.append仅读取列表），可以确保提供固定的结果（$5）。在这种情况下，无需进行任何更改。
当检测到动态函数并停止图实例化时，MAGPY会使用动态函数调用作为分割点，将用户程序切分为子程序。MAGPY可以通过监控用户程序的单次运行来分别编译每个子程序，并生成急切执行动态函数的编译后代码。

5.3.2 RefGraph的更新

使用Relation和InPlace注解更新RefGraph。Relation和InPlace注解用于更新RefGraph。
- Relation用于创建引用边。如果两个变量之间注解了引用关系，MAGPY将在RefGraph中对应的ShadowNodes之间添加引用边。对于像图13a中v.x=2.0这样的写操作，MAGPY在源ShadowNode中创建一个新版本，并仅从该版本添加一条引用边。对于读操作，如果关系已存在于源ShadowNode的最新版本中（例如，图13a中的a = v.x），MAGPY将不创建边；如果关系不存在（例如，图13a中的b = v.y），则会从源ShadowNode的所有版本创建到目标ShadowNode的边，因为该引用从程序执行开始就一直存在。
- InPlace用于版本管理。如果InPlace=True，MAGPY将向相应的ShadowNode添加一个新版本，并复制该版本的所有引用关系（如果它们仍然存在）。

6 检测和处理动态行为

处理动态行为。尽管大多数深度学习程序满足有限动态性，但有些程序仍然存在动态行为，如动态值标量、动态形状的Tensor和动态控制流。MAGPY可以自动检测它们并尽力处理。

6.1 动态性检测

检测动态性来源。动态函数的处理已在§5.3中讨论。另一个动态性来源是输入参数的动态性，这将导致guard失败。
- 为不同输入缓存记录。当guard失败时，MAGPY将重新生成一个编译记录并将其保存到缓存中。当一个guard匹配时，MAGPY将更新缓存，以便在后续调用中更早地尝试该记录。
- Guard失败原因检测。当记录数量超过一个阈值时，MAGPY将假定记录被过度特化，并尝试做出更弱的假设。MAGPY将定位导致guard失败的非恒定变量，并尝试为这些变量生成更弱的guard（§6.2.1）。

6.2 动态性处理

6.2.1 提升非恒定变量

将动态变量视为图节点。当MAGPY检测到一个非恒定输入变量，且其数据类型受底层图级编译器支持时，MAGPY会将其视为算子图中的一个“Tensor节点”，并像处理Tensor一样处理该变量。MAGPY将记录该变量的所有操作到算子图中，并重新计算所有依赖于它的变量，从而允许该变量的值发生变化。典型情况包括向用户程序提供动态标量和动态形状。如果底层图级编译器不支持该变量类型，MAGPY会将用户程序切分成不访问该变量的适当片段。

6.2.2 动态控制流

处理静态与动态控制流。用户代码混合了静态和动态控制流。
- 静态控制流通常用于尝试不同的架构或超参数，但模型结构在不同批次间是固定的。这些控制流操作通常使用静态标量来决定跳转目标。MAGPY会为这些标量添加适当的guard（如§4中所述），并在生成的图和mock中移除控制流操作。
- 动态控制流通常用于在运行时获取模型架构，并通常使用动态Tensor值来决定跳转目标。当目标图级编译器支持动态控制流时，MAGPY可以将这些动态跳转转换为算子图中的控制流算子。
- 对于循环操作，MAGPY会监控所有迭代的执行。如果在不同迭代中，外部读取的值和算子图都相同，MAGPY可以安全地将记录的算子图转换为Loop节点的主体。
- 对于分支操作，MAGPY会像AutoGraph [【28，Autograph: Imperative-style coding with graph-based performance，2019】]中那样，通过自动修改源代码来强制运行时执行两个分支。MAGPY通过监控两个分支的执行来保证正确性。具体来说，MAGPY首先执行程序不会跳转到的分支，记录该分支造成的所有影响，并确保所有这些影响都是可恢复的。如果MAGPY检测到不可恢复的影响，它将在动态跳转前切分用户程序，并急切地运行此跳转。然后，MAGPY恢复所有影响，运行程序将要跳转的分支，通过合并两个分支的算子图来生成If节点，并合并两个分支的RefGraph。

7 实现

MAGPY的实现细节。MAGPY基于Python和PyTorch实现，代码量约4000行。MAGPY中的算子图被导出为与主流图编译器兼容的torch.fx格式。要启用MAGPY，只需一行MagPy.compile代码，如图14所示。然后，在每次调用已编译模型时，MAGPY会尝试匹配一个已编译的记录，如果找不到匹配项，则执行基于监控的重新编译。

监控与代码注入机制。MAGPY使用Python中sys.settrace的per-bytecode回调来监控每个Python字节码的执行解释器状态。状态捕获是通过分析在sys.settrace回调期间传递给MAGPY的frame来实现的。guard匹配和mock函数的调用是通过使用Python的Frame Evaluation API修改字节码来实现的。

RefGraph的具体实现。MAGPY将Python中的每个对象视为一个拥有自己ShadowNode的变量。RefGraph中的get_node_by_obj接口（图8b）是通过一个对象ID到ShadowNode的映射实现的。为了使对象ID唯一，MAGPY持有所有对象的引用以避免垃圾回收。MAGPY为29种常用的Python数据类型实现了图11中的ShadowNode.match和ShadowNode.mock，并允许在监控运行期间存在其他数据类型，只要它们不需要被guard和mock。match在类型和值都与监控运行匹配时返回true。对于标量，值指的是确切的值。对于容器，MAGPY检查容器中的所有对象是否与监控运行期间相同位置的值相等。对于Tensor，MAGPY验证元数据没有改变，但不检查Tensor的值。如果MAGPY发现这些guard被过度特化，它会尝试做出更弱的假设，如§6.2.1所述。

A4 实验环境

• 硬件配置：
  • GPU：1块 NVIDIA A100-PCIE-40GB
  • CPU：2颗 AMD EPYC 7742
• 软件配置：
  • OS/环境：Python 3.9, CUDA 11.8, GCC 11.4
  • 依赖库：PyTorch v2.0.1, torch.xla v2.0
• 基线系统 (Baselines)：
  • 基于分析的框架：TorchDynamo [1], TorchScript [2]
  • 基于追踪的框架：LazyTensor
  • 未比较的系统：torch.jit.trace [2] 和 AutoGraph [28] 因可能产生静默错误未比较；Janus [20] 和 Terra [23] 因未开源未比较。
• 基准测试集 (Benchmarks)：
  • ParityBench：包含1418个来自Github（超过100星）的PyTorch模型，用于评估对真实用户代码中Python特性的覆盖率。
  • 8个代表性DNN模型：用于性能评测，覆盖CNN、Transformer等典型架构以及多模态、量化等场景。具体模型信息见下表1。

表1: 模型信息。BS代表“批大小”。引用数基于截至2024年1月8日的Google Scholar统计。Source中的xxx/http://yyy。

模型	输入形状	引用数	来源
ALIGN	文本长度64，图像[BS, 3, 289, 289]	2029	huggingface/transformers v4.29.1
Bert	文本长度256	88345	huggingface/transformers v4.29.1
DeBERTa	文本长度256	1595	huggingface/transformers v4.29.1
DenseNet	图像[BS, 3, 224, 224]	41359	pytorch/vision v0.4.1
MonoDepth	图像[BS, 3, 256, 256]	3151	OniroAI/MonoDepth-PyTorch b7bb004
Quantized	图像[BS, 3, 224, 224]	330	eladhoffer/quantized-pytorch d6fc447
ResNet	图像[BS, 3, 224, 224]	195511	pytorch/vision v0.4.1
TridentNet	图像[BS, 3, 224, 224]	932	open-mmlab/mmdetection v2.28.2, mmcv v1.7.11

A4 实验结果

8.2 端到端评估

• 实验内容：比较不同图实例化方法（MAGPY, TorchDynamo, TorchScript, LazyTensor）在不同后端编译器（TorchInductor, TorchScript, XLA）上的端到端推理性能。
• 实验结果：如图15所示，MAGPY相比Eager执行平均提速1.73倍（最高2.93倍）。与使用相同后端编译器的基线相比，MAGPY分别取得了最高6.25倍（对TorchDynamo-Inductor）、2.88倍（对TorchScript）和8倍（对LazyTensor-XLA）的加速，平均加速分别为1.68倍、1.56倍和1.78倍。
• 分析结论：MAGPY的性能提升主要源于两点：1) 减少了Python解释器开销，只需执行快速的guard和mock函数；2) 生成了更完整的算子图（如表2所示，MAGPY对所有模型都生成了1个图，而其他方法会产生多个图碎片），这使得后端图编译器能够进行更全面的优化。

表2：导出的算子图数量

8.3 Python特性覆盖率

• 实验内容：使用ParityBench中1191个理论上静态的模型，评估MAGPY、TorchScript和TorchDynamo将模型完整编译为单个算子图的能力。
• 实验结果：如表3所示，MAGPY成功将93.4%的模型（1112个）实例化为完整的算子图，失败率仅为6.6%。相比之下，TorchDynamo的失败率为22.8%，TorchScript的失败率高达64.5%。MAGPY将不支持Python特性的比率相比最先进的TorchDynamo降低了3.44倍。
• 分析结论：MAGPY对复杂的Python特性具有更强的扩展性和支持能力。主要的失败原因包括：Python中不同表达式生成的布尔标量共享相同的对象ID，导致无法提升动态布尔值；以及部分输出类型（如deque, Counter）的mock生成尚未支持。

表3：ParityBench上的结果

8.4 开销分析

• 实验内容：分析MAGPY在“监控运行”（首次编译生成记录）和“匹配运行”（执行已编译记录）中的时间开销构成。
• 实验结果：如图16所示，在监控运行中，MAGPY自身的分析时间（MagPy profile）平均为2.91秒，占总时间的23%，而图编译时间（TorchInductor）占77%。在匹配运行中，guard验证和mock执行的开销非常小，平均分别占总时间的2%和98%。
• 分析结论：MAGPY的编译时开销是可接受的，且主要由后端图编译器决定。其运行时（匹配运行时）开销极低，保证了推理性能。

8.5 动态场景评估

• 实验内容：评估MAGPY在两种常见动态场景下的性能：动态形状输入和动态控制流。
• 实验结果：
  • 动态形状（图17）：对于简单模型（Bert, ResNet），MAGPY性能与TorchDynamo相当。对于复杂模型（DeBERTa, DenseNet），MAGPY比TorchDynamo分别快1.06倍、1.16倍和2.05倍。
  • 动态控制流（图18）：在LSTM和BlockDrop模型上，MAGPY相比TorchDynamo平均提速5.96倍，最高达到11.59倍。
• 分析结论：MAGPY提出的动态处理机制（§6.2）是有效的，能够为动态形状和动态控制流的程序生成高效的算子图，并在复杂模型上显著优于现有技术。

A7 补充细节

9 相关工作

• 从用户代码导出算子图：
  • 基于分析：Janus [20], TorchScript [2], TorchDynamo [1]直接分析用户代码，但仅支持Python的一个子集，对复杂特性支持不佳。
  • 基于追踪：torch.jit.trace [2] 和 torch.fx [31] 创建一次性编译的静态图。AutoGraph [28] 和 JAX.jit [12] 能在输入参数变化时重编译，但无法检测全局变量等外部值的变化，可能产生静默错误。LazyTensor [33], Terra [23], 和 Torchy [26] 通过每次运行时重新追踪来保证正确性，但这会引入巨大的运行时开销。
• 图级深度学习编译器：MAGPY等系统导出的算子图会被图级编译器进一步优化。这些技术包括图替换（如TASO [22], PET [38], EinNet [43]）和核融合（如Rammer [27], DNNFusion [29], Welder [32]）。TVM [8], FlexTensor [45], Ansor [42]等工作则专注于优化深度学习算子。
• 领域特定语言（DSL）：Triton [37] 和 FreeTensor [36] 等DSL使用Python作为前端，图实例化技术可用于将Python程序转换为它们的中间表示，有助于避免静默错误并支持更灵活的语法。
• 通用程序的即时编译：基于运行时信息的JIT编译在通用程序中广泛使用，如Python的PyPy [5]和Numba [24]，JAVA的HotSpot [30]和trace-JIT [19]。这些工具将动态程序转换为低级机器码，而MAGPY则创建高级算子图。
• 多阶段编程：图实例化和编译过程类似于多阶段编程（如MetaML [34]），但多阶段编程语言需要用户手动指定阶段，且不能在外部环境变化时自动重编译。

A5 结论

本文提出了MAGPY，一个利用深度学习程序中固有的“有限动态性”来实现高效算子图实例化的系统。MAGPY引入了RefGraph来记录程序状态，并通过监控影响程序行为的外部值来降低图实例化的复杂性。评估结果显示，MAGPY能够将复杂的深度学习程序加速高达2.88倍（平均1.55倍），并成功地将1191个有限动态的用户程序中的93.40%实例化为完整的算子图。

A6 附录

A 工件附录

• 摘要：此工件有助于复现ATC'24论文“MAGPY: 通过监控执行状态编译Eager模式的DNN程序”中的结果。MAGPY的输入是一个PyTorch程序，它会自动将程序导出为torch.fx格式，并使用论文中描述的方法添加适当的guard。然后，MAGPY使用用户提供的图级编译器来编译导出的图。
• 范围：该工件可用于复现论文中的实验，包括引言（图1）、端到端比较（图15）、开销分解（图16和表2）、动态模型（图17和18）以及覆盖率（表3）。
• 内容：此工件包括MAGPY的代码、实验输入数据、实验环境设置指南以及运行实验的脚本。
• 托管：http://MAGPY。
• 要求：此工件需要一台至少配备一块NVIDIA A100 GPU的机器，并已正确安装NVIDIA驱动程序。用户可以按照安装指南设置软件环境以复现结果。