发表时间: 2024-07
文章标题:MagPy: 通过监控执行状态编译Eager模式的DNN程序
作者/机构:Chen Zhang, Rongchao Dong, Haojie Wang, Runxin Zhong, Jike Chen, and Jidong Zhai, 清华大学
一句话结论 本文提出了MAGPY系统,通过在运行时监控Python解释器的执行状态并构建引用图,将具有复杂动态特性的Eager模式深度学习程序自动且完整地转化为算子图,从而大幅提升了模型的执行性能。
要解决什么问题 深度学习编译器依赖完整的静态算子图来进行底层优化,但现实中的模型通常使用Python以Eager模式开发。Python的动态数据类型使得编译器无法预知哪些操作作用于张量,海量的内置API处理、复杂的面向对象编程封装以及普遍存在的变量副作用,导致传统的图实例化技术卡在图的完整性与开销上。具体而言,基于静态分析的路线(如TorchDynamo和TorchScript)难以解析所有Python复杂特性,遇到不支持的语法就会打断计算图,导致图碎片化并引入高昂的Python解释器回退开销;而基于追踪的路线(如LazyTensor)要么强依赖用户保证图的静态性从而容易引发静默错误,要么在每次执行时都重新追踪,带来巨大的运行时开销。这种高级动态语言语义与底层静态算子图之间的巨大语义鸿沟,使得包含复杂特性的真实用户程序很难被完整编译,无法确定变量在未来是否会被使用,进而阻碍了死代码消除或算子融合等图级优化,最终导致模型执行效率低下。
怎么做的 方法的核心思路是利用深度学习程序的“有限动态性”——即尽管Python语法高度动态,但模型在不同批次间的算子执行顺序、属性和形状通常是不变的。因此,只要从外部读取的值(如输入参数和全局变量)不变且操作确定,程序行为就不会变。基于此,系统放弃了对复杂程序逻辑的静态分析,转而通过监控Python解释器的底层执行状态来捕获变量间的引用关系。关键设计由引用图(RefGraph)、状态捕获与注解系统、以及代码生成器三个部件构成。首先,系统在程序监控运行期间构建RefGraph,图的节点是代表运行时变量的ShadowNode(其中特殊的节点编号0用于保存外部运行时状态),边是变量间的引用关系,同时引入ShadowVersion来记录变量的原地更新历史。为了降低开销,ShadowNode仅在变量被显式访问时才懒惰创建。其次,为了处理C语言实现的底层内置函数,系统引入了轻量级的操作属性注解,通过标记函数是否能转为算子节点(AsOpNode)、是否纯粹(PureClosure)、是否读取张量值(ValueRead)以及是否原地修改(InPlace)等属性,来指导RefGraph的更新并检测动态行为。最后,代码生成器基于构建好的RefGraph生成三个产物:一是算子图,提取所有张量计算交由底层编译器优化;二是Guard函数,通过在RefGraph中从外部状态节点出发搜索第一个版本的引用关系,收集所有显式读取的外部变量,用于在后续运行时验证初始输入状态是否发生改变;三是Mock函数,通过搜索最后一个版本的引用关系以及被原地修改的变量,用于在不执行原始Python代码的情况下,直接复现程序的最终状态和副作用。当新输入通过Guard验证时,系统直接调用优化后的算子图和Mock函数,完全绕开解释器;对于动态标量或动态控制流,系统会自动将其提升为算子图中的张量节点或强制执行分支合并来处理。
效果如何 实验在配备NVIDIA A100 GPU的硬件环境下进行,使用了包含1418个真实用户模型的ParityBench数据集以及8个代表性深度学习模型(涵盖Bert、ResNet、DeBERTa等)进行评估。对比基线包括代表基于静态分析路线的TorchDynamo和TorchScript,以及代表基于追踪路线的LazyTensor,底层编译器搭配了TorchInductor和XLA。量化结果显示,在端到端推理性能上,该系统相比Eager模式平均提速1.73倍,最高达2.93倍;在相同底层编译器下,相比TorchDynamo搭配Inductor最高加速6.25倍,相比LazyTensor搭配XLA最高加速8倍。在Python特性覆盖率测试中,该系统成功将ParityBench中93.40%的模型实例化为单一完整的算子图,失败率仅为6.6%,而TorchDynamo和TorchScript的失败率分别高达22.8%和64.5%,证明了该方法能有效消除图碎片化。开销方面,匹配运行时的Guard验证仅占总时间的2%,运行时开销极低。在动态控制流场景下,该系统在LSTM等模型上相比TorchDynamo平均提速5.96倍。作者承认的局限性在于,由于Python中不同表达式生成的布尔标量会共享相同的底层对象ID,导致系统目前无法将动态布尔值提升为图节点,且部分特殊输出类型(如双端队列和计数器)的Mock代码生成尚未得到支持。
核心问题:
现实世界中的深度学习程序通常使用像Python这样的动态编程语言开发,这些语言具有复杂的特性(如内置函数、动态类型),并以Eager模式执行,导致性能不佳。深度学习编译器依赖于基于算子的计算图来优化程序,但动态语言的复杂性常常阻碍了程序被完整地转换为算子图,从而导致次优性能。现有的图实例化技术分为两类:基于分析的方法(如TorchDynamo)和基于追踪的方法(如LazyTensor)。基于分析的方法难以支持所有Python特性,导致图碎片化和高昂的解释器开销;基于追踪的方法要么依赖用户保证图的静态性(可能引入静默错误),要么为每次执行重新追踪,开销巨大。如图1所示,现有技术与手动构建完整图的性能差距显著。
研究目标:
本文旨在提出一种新的方法,通过监控Python解释器的执行状态来增强算子图的实例化,从而更有效地将用户编写的、易于使用的Python程序自动转换为编译器友好的完整算子图,以提升深度学习程序的执行性能。
创新点:
本文基于以下三个核心洞察提出了MAGPY系统:
Guard函数和用于复现上次运行最终状态的Mock函数,都可以通过分析程序运行时的状态(特别是变量间的引用关系)来确定,而无需深入理解程序逻辑。基于以上洞察,MAGPY做出了以下主要贡献:
将深度学习程序实例化为算子图面临的挑战。深度学习程序的算子图实例化之所以充满挑战,是因为Python和PyTorch为了易于编程而提供了大量灵活的特性。编译器必须正确处理所有这些特性才能生成正确的算子图。通过对ParityBench中272个TorchDynamo无法导出完整算子图的真实用户程序进行手动分析,我们总结出以下主要障碍:
forward函数在第2行写入self.dim属性,在第4行读取该属性。这些修改会改变后续函数调用的行为,必须被正确处理。6.3%的Dynamo失败由此引起。尽管存在上述挑战,用户程序的精确行为在不同批次间通常是相同的。因此,只要能够验证程序行为保持不变,MAGPY就可以安全地重用前一次运行的算子图,而无需理解这些程序的精确语义。
利用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系统将重新编译用户程序以生成一个新的记录,并调用图级深度学习编译器来优化该记录中的算子图。然后,该记录被保存到缓存中以备将来使用。
MAGPY系统概览。MAGPY的概览如图6所示。当在缓存中找不到匹配的记录时,MAGPY会通过使用本地语言执行器(如Python解释器)执行程序并监控其执行过程来重新编译用户程序。我们称此过程为用户程序的“监控运行”。与传统编译器基于控制流图等结构分析程序逻辑不同,MAGPY只关心程序状态。MAGPY提出了一个RefGraph(引用图)来保存运行时状态信息,主要是用户程序运行时变量之间的引用关系。在程序执行期间,MAGPY从本地执行器捕获每条指令的执行状态,并根据捕获的执行状态更新RefGraph。当程序执行完毕后,MAGPY通过分析RefGraph生成一个新的记录,并将其存入记录缓存中。
本节将介绍RefGraph的定义(§4.1),以及基于RefGraph生成guard、mock和算子图的算法(§4.2-§4.3)。
RefGraph的构成。RefGraph的定义如图8所示,它包含以下元素:
x[0]和y之后,执行x[0] += y之前的RefGraph。MAGPY仅为三个使用过的变量x (SN #1)、x[0](SN #2)和y (SN #3)创建了ShadowNodes,而忽略了未使用的参数z。x[0] += y为x[0] (SN #2)创建了一个新版本,步骤(3)中的x.append()为x (SN #1)创建了一个新版本。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搜索关键变量。由于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所示:
a。这些变量可以通过从外部节点SN #0开始,使用最后一个版本中的引用关系作为边进行搜索来收集(算法1的第16、19-21行)。然而,如果一个变量在函数入口处已存在且未被原地更新(其ShadowNode中只有一个版本),MAGPY可以跳过复现它,而是使用其初始值。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中的引用关系来实现。
确定算子图的输入和输出。对于记录的算子图,MAGPY将确定输入和输出节点,以及这些张量在mocking期间应从运行时加载或存储到何处。输入节点是算子图中没有入边的Tensor节点。这些Tensor变量不是由张量操作创建的,因此它们应该在函数入口处就存在。因此,在RefGraph中存在从SN #0使用版本0的边到相应节点的引用路径。输出节点是需要被mock的Tensor变量,并在mock生成期间确定。输出节点也可以通过从RefGraph中的SN #0搜索得到,其路径表示输出Tensors的存储位置。MAGPY可以确保中间节点将来不会被使用,并将此信息传递给图级编译器。因此,图级编译器可以安全地执行诸如死代码消除或算子融合等优化。然而,基于追踪的框架无法实现这一点,因为它们无法知道一个Tensor将来是否会被使用。
RefGraph的生成过程。MAGPY通过分析监控运行的中间运行时状态来生成RefGraph。具体来说,MAGPY在监控运行期间捕获每条指令的执行状态(§5.1),获取该指令预定义的操作属性(§5.2),并根据该属性使用不同的RefGraph更新规则(§5.3)来更新RefGraph。
从运行时收集信息。动态语言通常会保存运行时变量的高级信息(如类型和变量结构)以进行动态解释。这些信息对MAGPY分析运行时状态非常有价值。因此,MAGPY从运行时收集每条指令的执行状态,其中包含该指令所有相关的变量。执行状态包含以下元素:
处理原生代码实现的复杂性。出于性能考虑,动态语言使用原生机器码来实现某些操作。例如,Python用C语言实现了一些函数。这些低级机器码不保留高级变量信息,难以分析。为了收集所有算子并正确运行MAGPY,需要对这些函数进行属性注解。请注意,这些只是注解,比基于分析的方法中重新实现这些函数要容易得多。MAGPY也鼓励对动态语言实现的常用函数进行注解,以便MAGPY可以跳过深入函数实现细节,从而加快监控运行速度。所需的注解如下:
函数属性 (Function Properties) 为函数注解:
AsOpNode用于生成算子图。PureClosure识别Python中的动态行为,如系统调用和随机数生成器,它们的注解将为False。
输入属性 (Input Properties) 为函数的每个输入变量注解:
例如,像add或sub这样的数学运算的所有参数都将被注解为ValueRead=True。对于向列表中添加变量的函数list.append,第一个参数(列表)被注解为ValueRead=True,因为它的内部结构被函数访问。相反,第二个参数(要添加的变量)被注解为ValueRead=False,因为函数只将其引用添加到列表中,而不访问其值。此外,第一个参数被注解为InPlace=True,因为列表被指定的值原地修改了,而第二个参数被注解为InPlace=False。
输出属性 (Output Properties) 为每个输出变量注解:
v.x)还是创建新引用的写入(例如,赋值v.x = 1.0)。注解的用户工作量。进行此类注解的用户工作量是可承受的。尽管总共有6个属性,但只有操作的PureClosure以及输入和输出变量的InPlace被频繁注解。AsOpNode只需为图级编译器支持的有限数量的操作进行注解。ValueRead仅在函数可能读取Tensor时才必要,这是由于§5.3中的图更新规则。Relation仅用于特殊的懒惰创建引用关系,如显式加载属性。此外,许多操作具有相同的注解,例如+、-、×、÷操作,这允许批量生产注解。
更新RefGraph和算子图的流程。MAGPY首先为输入和输出变量中没有对应节点的变量生成ShadowNodes。然后,MAGPY检索操作的属性注解,并根据注解和捕获的执行状态更新RefGraph和算子图。
处理不同类型函数以更新图和检测动态行为。AsOpNode、PureClosure和ValueRead注解用于更新算子图和检测动态性。其工作流程如图12所示。
AsOpNode成功为函数生成一个算子节点(例如,torch.relu),MAGPY会将该节点添加到算子图中($2)。否则,跳转到(4)。
4. **使用PureClosure检测动态函数**。如果`PureClosure`为False,MAGPY会将该函数视为动态函数并停止图实例化($3)。如果为True,MAGPY将假定该函数的行为在不同批次中保持不变,并跳转到(5)。ValueRead=True且其运行时值为Tensor。如果找到这样的参数,MAGPY会将该函数视为动态函数并停止图实例化($4)。原因是MAGPY允许Tensor数据在不同批次中发生变化,因此这些操作会因读取Tensor数据而产生不同的结果。典型情况是图级编译器不支持的张量操作,如Inductor的nn.LSTM。这些操作未被AsOpNode转换为算子节点,因此会进入此分支。否则,操作仅读取非张量变量的值(如list.append仅读取列表),可以确保提供固定的结果($5)。在这种情况下,无需进行任何更改。当检测到动态函数并停止图实例化时,MAGPY会使用动态函数调用作为分割点,将用户程序切分为子程序。MAGPY可以通过监控用户程序的单次运行来分别编译每个子程序,并生成急切执行动态函数的编译后代码。
使用Relation和InPlace注解更新RefGraph。Relation和InPlace注解用于更新RefGraph。
v.x=2.0这样的写操作,MAGPY在源ShadowNode中创建一个新版本,并仅从该版本添加一条引用边。对于读操作,如果关系已存在于源ShadowNode的最新版本中(例如,图13a中的a = v.x),MAGPY将不创建边;如果关系不存在(例如,图13a中的b = v.y),则会从源ShadowNode的所有版本创建到目标ShadowNode的边,因为该引用从程序执行开始就一直存在。InPlace=True,MAGPY将向相应的ShadowNode添加一个新版本,并复制该版本的所有引用关系(如果它们仍然存在)。处理动态行为。尽管大多数深度学习程序满足有限动态性,但有些程序仍然存在动态行为,如动态值标量、动态形状的Tensor和动态控制流。MAGPY可以自动检测它们并尽力处理。
检测动态性来源。动态函数的处理已在§5.3中讨论。另一个动态性来源是输入参数的动态性,这将导致guard失败。
将动态变量视为图节点。当MAGPY检测到一个非恒定输入变量,且其数据类型受底层图级编译器支持时,MAGPY会将其视为算子图中的一个“Tensor节点”,并像处理Tensor一样处理该变量。MAGPY将记录该变量的所有操作到算子图中,并重新计算所有依赖于它的变量,从而允许该变量的值发生变化。典型情况包括向用户程序提供动态标量和动态形状。如果底层图级编译器不支持该变量类型,MAGPY会将用户程序切分成不访问该变量的适当片段。
处理静态与动态控制流。用户代码混合了静态和动态控制流。
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所述。
表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 |
表2:导出的算子图数量
表3:ParityBench上的结果
本文提出了MAGPY,一个利用深度学习程序中固有的“有限动态性”来实现高效算子图实例化的系统。MAGPY引入了RefGraph来记录程序状态,并通过监控影响程序行为的外部值来降低图实例化的复杂性。评估结果显示,MAGPY能够将复杂的深度学习程序加速高达2.88倍(平均1.55倍),并成功地将1191个有限动态的用户程序中的93.40%实例化为完整的算子图。