depyf: Open the Opaque Box of PyTorch Compiler for Machine Learning Researchers

发表时间: 2024-03 · arXiv:2403.13839

原文: https://arxiv.org/abs/2403.13839

速读

一句话结论 本文开发了开源工具 `depyf`，通过字节码反编译和函数执行劫持技术，将 PyTorch 2.x 编译器生成的不可读字节码还原为可调试的 Python 源代码，使研究者能使用标准调试器逐行检查编译后的计算图。

要解决什么问题 PyTorch 2.x 引入内置编译器 `torch.compile` 以优化硬件执行效率，但该编译器对研究者如同黑盒，导致了严重的调试卡点。具体而言，编译器的前端 Dynamo 负责将用户代码拆分为纯 Python 代码和构成计算图的纯 PyTorch 代码。这一拆分过程直接在 Python 字节码层面运行，例如处理加载、跳转和调用等底层指令，而绝大多数算法研究者并不具备直接阅读这些底层字节码的能力。更致命的卡点在于编译器的后端机制：后端在提取出计算图后，会将其优化并编译成适用于各类硬件的二进制可执行文件。在 Python 运行环境中，这个被编译出来的计算图表现为一个动态生成的单一函数，系统强制要求该函数必须被作为一个整体完整执行。这种机制彻底切断了传统调试路径，意味着用户根本无法使用常规调试器去逐行分析这个动态函数。当模型在训练或推理中出现非数字（NaN）等棘手的数值错误时，研究者无法通过逐行跟踪代码来定位具体是哪一步张量操作导致了异常，极大地拖慢了问题排查效率。此外，现有 Python 反编译器大多只兼容旧版 Python，且专门针对由静态源代码编译出的字节码设计，完全无法处理 PyTorch 在程序运行时动态生成的复杂字节码。

怎么做的 为绕开动态字节码无法阅读和计算图无法逐行调试的卡点，`depyf` 的核心思路是将内存中动态生成的晦涩代码对象反向映射回磁盘上人类可读且可交互的源代码文件。该方法主要由字节码反编译、底层逻辑重写以及函数执行劫持三个关键部件构成。首先是字节码反编译部件，其职责是将 PyTorch 编译器生成的底层字节码转换回等效的 Python 源代码。为克服现有工具无法处理动态生成字节码的缺陷，作者采用对字节码进行符号执行的全新策略。这种方法通过模拟字节码执行路径推导高层语义，仅需处理大约两百种基础 Python 字节码类型，从而兼容了 PyTorch 支持的所有 Python 版本。其次是底层逻辑重写部件，`depyf` 在内部使用纯 Python 重新实现了 PyTorch 编译器中原本由 C 语言编写的核心组件，向用户透明展示编译器底层的分析和转换逻辑。最后是函数执行劫持部件，由于 Python 调试器只能对磁盘上的源代码文件进行逐行追踪，该部件利用 Python 动态特性，在运行时精准拦截并替换 PyTorch 内部的关键函数调用，将原本不可拆分的动态生成函数强行替换为包含完整调试信息且已落盘的对应源代码函数。在工程实现上，这套机制被封装为两个非侵入式的上下文管理器。用户用 `with depyf.prepare_debug()` 包裹目标代码，工具会自动捕获编译器内部细节，将计算图和反编译源码转储到本地目录；随后用户通过 `with depyf.debug()` 即可暂停程序，在转储的源码中设置断点，并呼出标准 Python 调试器（如 pdb）逐行执行和检查张量操作。

效果如何 实验验证主要围绕工具的正确性与兼容性展开。测试环境覆盖了 PyTorch 官方支持的所有 Python 版本（Python 3.8 至 3.11），并针对 PyTorch 夜间构建版本实施持续集成测试。测试涵盖超过八十个 Python 语法用例，并覆盖三大模型库的核心架构：TorchBench 的 Segment Anything 等；Huggingface Transformers 的 LLaMA、BERT、GPT 家族等；以及 TIMM 的 ResNet、ViT 等视觉模型。在对比基线方面，作者将 `depyf` 与多种现有的 Python 反编译器（代表传统的静态字节码反编译路线）进行了横向对比。量化测试结果表明，在面对不同 Python 版本和海量 PyTorch 模型的复杂动态字节码时，现有反编译器均出现了不同程度的解析失败，而 `depyf` 是唯一一款成功通过所有正确性与兼容性测试的反编译器。通过该工具，研究者成功实现了对上述复杂模型编译后计算图的逐行断点调试，作者还借此收集并开源了这些模型在训练和推理过程中张量形状变化的全部计算细节。本文未明确交代实验所依赖的具体硬件配置，且该工具主要聚焦于解决编译器的可解释性与调试卡点，并未在文中报告其对模型实际运行性能可能产生的额外开销。

depyf：为机器学习研究者打开PyTorch编译器的不透明盒子

作者/机构:
Kaichao You† *
Runsheng Bai†
Meng Cao§
Jianmin Wang†
Ion Stoica‡
Mingsheng Long† B

† 软件学院, BNRist, 清华大学, 北京 100084, 中国
§ AIML, Apple
‡ 计算机科学部, 加州大学伯克利分校, CA 94720-1776, 美国

A1 主要贡献

本文旨在解决机器学习研究者在使用 PyTorch 2.x 内置编译器 torch.compile 时遇到的理解和调试困难。

  • 核心问题: 随着大型语言模型（LLM）的发展，深度学习对计算资源的需求日益增长，催生了GPU和TPU等专用硬件。然而，机器学习研究者通常专注于算法开发，缺乏充分利用这些硬件所需的底层编程知识。PyTorch 2.x 引入的 torch.compile 编译器旨在弥合这一差距，但其工作在Python字节码层面，对用户来说如同一个“不透明的盒子”，难以理解其内部机制、进行调试或优化。

  • 研究目标: 本文的目标是揭开 PyTorch 编译器的神秘面纱，让研究者能够清晰地理解其内部工作流程。

  • 创新点与主要贡献: 为了实现这一目标，本文开发并介绍了一款名为 depyf 的工具。其主要贡献如下：

    1. 字节码反编译: depyf 能够将 PyTorch 编译器生成的字节码反编译为等效的、人类可读的Python源代码。这使得研究者可以直观地看到他们的代码被转换成了什么样子。
    2. 源代码关联与调试: 该工具能够将内存中动态生成的代码对象与其在磁盘上对应的源代码文件关联起来。这一特性允许用户使用标准的Python调试器（如pdb）逐行执行和检查编译器生成的代码，极大地增强了对底层过程的理解和调试能力，尤其是在处理 NaN（非数字）错误等棘手问题时。
    3. 非侵入式和易用性: depyf 的使用非常简单，主要通过两个方便的上下文管理器 (with depyf.prepare_debug() 和 with depyf.debug()) 来实现其核心功能，对用户的现有代码库几乎没有侵入性。
    4. depyf 项目已开源，并被PyTorch官方认可为生态系统项目。

A3 理解PyTorch编译器的挑战

Dynamo：PyTorch编译器的前端

PyTorch编译器中最复杂的部分是其名为Dynamo的前端。Dynamo的关键功能是将用户代码分离成纯Python代码和构成计算图的纯PyTorch代码。其操作过程涉及三个主要步骤：首先，识别出第一个无法在计算图中表示、但又需要图中已计算张量值的操作，例如打印张量值或在if语句中使用张量值进行控制流判断。其次，将此前的操作划分为两部分：一个只包含张量计算的计算图和一个专门操作Python对象的Python代码。最后，将后续操作处理为一个或多个新的函数（称为恢复函数），并对这些函数递归地重新启动上述分析过程。Dynamo在Python字节码层面运行（如图1中的LOAD, JUMP, CALL指令），这是一个比Python源代码更基础的层面，很少有机器学习研究者能够熟练解读这种字节码。

图1: PyTorch编译器的工作流程（左）以及depyf如何提供帮助（右）。
图1: PyTorch编译器的工作流程（左）以及depyf如何提供帮助（右）。

PyTorch编译器的后端

在前端提取出计算图后，后端会优化这个图，并最终生成适用于CPU、GPU和TPU硬件的二进制可执行文件。在Python中，一个计算图是一个动态生成的函数，这意味着它必须被完整执行。因此，用户无法使用调试器对该函数进行逐行分析。当计算结果出现NaN（非数字）错误时，这个问题变得尤其具有挑战性，因为它排除了逐行跟踪代码以确定导致数值问题的具体操作的可能性。

A2 解决方案

字节码反编译

本文的主要目标是让机器学习研究者摆脱字节码的复杂性。 将字节码转换回源代码的过程称为“反编译”。在depyf之前，现有的Python反编译器虽然可以将Python字节码转换为源代码，但它们存在显著的局限性：它们通常只支持旧版本的Python，兼容性有限；并且，它们是为反编译由源代码编译而来的字节码而设计的，难以处理像PyTorch那样由程序生成的字节码。为了克服这些问题，我们通过对字节码进行符号执行，创建了一个新的Python字节码反编译器。这种方法仅需处理大约两百种Python字节码类型，确保了与PyTorch支持的所有Python版本的兼容性。此外，depyf内部用Python复现了PyTorch编译器的C语言核心组件，以便为用户阐明其底层机制。

函数执行劫持

为了能让用户使用调试器逐行执行代码，Python执行的字节码必须来源于磁盘上的源代码文件。 我们利用了Python的高级特性来拦截和替换PyTorch中的关键函数调用。这种替换涉及将动态生成的函数替换为包含调试信息的对应函数。

使用方法

depyf的使用方式简单直接且非侵入。 用户只需将他们的代码包裹在with depyf.prepare_debug()上下文管理器中。这一操作使得depyf能够捕获该上下文中所有PyTorch的内部细节，包括反编译的源代码和计算图。对于希望使用调试器逐行执行反编译代码的用户，可以使用另一个上下文管理器with depyf.debug()。更多使用细节见附录A。

概览

图1概述了depyf的处理流程。 更全面的细节可以在我们的文档页面找到。depyf的优势体现在三个方面：首先，它提供了一个与PyTorch的C语言实现相对应的Python实现，帮助用户理解PyTorch编译器的逻辑（参见图2中的full code xxx.py）。其次，它包含一个Python字节码反编译器，能将字节码转换为等效的源代码，帮助用户理解PyTorch转换后的字节码（参见图2中的transformed xxx.py）。最后，它劫持了PyTorch中的关键函数，使用户能够使用调试器逐行执行计算图函数（参见图2中的compiled xxx.py）。

图2: depyf的两种用法。
图2: depyf的两种用法。

A4 实验环境与结果

实验内容与结果

本文通过一系列实验验证了depyf的正确性和兼容性。表1展示了多种现有反编译器与depyf在不同Python版本和PyTorch测试集上的兼容性情况。测试的详细描述见附录B和C。从结果可以看出，depyf是唯一成功通过所有测试的反编译器。我们的测试方法采用持续集成的方式，即每一个新的代码提交都会针对所有支持的Python版本下的PyTorch夜间版进行测试。这种主动策略使我们能够在新的PyTorch版本发布前识别并解决任何兼容性问题。此外，我们还与PyTorch团队进行讨论，以提出维持这种兼容性的解决方案。

表1: Python和PyTorch测试中各反编译器的正确性。
表1: Python和PyTorch测试中各反编译器的正确性。

我们还收集了所有这些实验的输出，这为PyTorch编译器的新手提供了一个宝贵的资源，帮助他们深入了解常见深度学习模型的计算细节。更多信息可以在附录D中找到。

实验环境

  • 模型与数据集: 实验中使用的模型覆盖了深度学习领域的多个方面，主要来自三个模型库：

    • TorchBench 【Constable et al., TorchBench: A collection of open source benchmarks for PyTorch performance and usability evaluation, 2020, https://github.com/pytorch/benchmark】：收集了来自著名机器学习项目的模型， 如Segment Anything 【Kirillov et al., Segment Anything, 2023】和SuperSloMo 【Jiang et al., Super slomo: High quality estimation of multiple intermediate frames for video interpolation, 2018, CVPR】。
    • Huggingface Transformers 【Wolf et al., Transformers: State-of-the-art natural language processing, 2020, EMNLP】：最流行的Transformer模型库，包括LLaMA 【Touvron et al., LLaMA: Open and Efficient Foundation Language Models, 2023】和BERT 【Devlin et al., BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding, 2019, NAACL】。
    • TIMM 【Wightman, PyTorch Image Models, 2023】：最流行的计算机视觉模型库，包括ResNet 【He et al., Deep residual learning for image recognition, 2016, CVPR】和ViT 【Dosovitskiy et al., An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale, 2021, ICLR】。
      （详细模型列表见附录B）

  • 软件配置: depyf 在所有PyTorch支持的Python版本上进行了测试（根据表1，包括Python 3.8, 3.9, 3.10, 3.11），并针对PyTorch的夜间构建版本进行持续集成测试。

  • 硬件配置: 论文未明确指定实验所用的具体硬件配置。

A5 结论

在本文中，我们介绍了depyf，一个旨在打开PyTorch编译器这个不透明盒子（opaque box）的新工具，以帮助机器学习研究人员理解和适应torch.compile。

A6 附录

附录A. 使用方法

我们为用户提供了两个便捷的上下文管理器：with depyf.prepare_debug() 和 with depyf.debug()。 第一个管理器会捕获所有对使用torch.compile的函数的调用，并将许多内部细节转储到用户指定的目录中（即with depyf.prepare_debug()的参数）。第二个管理器会暂停程序，以便用户在转储的源代码中设置断点，并且任何与torch.compile相关的函数调用都可以使用标准的Python调试器逐行执行。depyf转储的源代码有三种类型：计算图（前缀为compiled）、反编译的源代码（前缀为transformed）和描述性源代码（前缀为full_code）。

图2: depyf的两种用法。
图2: depyf的两种用法。

附录B. 测试的PyTorch模型

本节列出了我们在表1中测试的所有PyTorch模型。 这些模型来自三个深度学习模型套件：TorchBench 【Constable et al., TorchBench: A collection of open source benchmarks for PyTorch performance and usability evaluation, 2020】 收集了来自著名机器学习仓库（按 https://paperswithcode.com/ 排名的高引用项目）的模型，如Segment Anything 【Kirillov et al., Segment Anything, 2023】 和SuperSloMo 【Jiang et al., Super slomo: High quality estimation of multiple intermediate frames for video interpolation, 2018, CVPR】；Huggingface Transformers 【Wolf et al., Transformers: State-of-the-art natural language processing, 2020, EMNLP】 是最流行的Transformer模型库，包括LLaMA 【Touvron et al., LLaMA: Open and Efficient Foundation Language Models, 2023】 和BERT 【Devlin et al., BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding, 2019, NAACL】；TIMM 【Wightman, PyTorch Image Models, 2023】 是最流行的计算机视觉模型库，包括ResNet 【He et al., Deep residual learning for image recognition, 2016, CVPR】 和ViT 【Dosovitskiy et al., An image is worth 16x16 words: Transformers for image recognition at scale, 2021, ICLR】。

具体模型包括（部分列举）：
* BERT家族: BertForMaskedLM, BertForQuestionAnswering, BERT pytorch, hf Bert 【Devlin et al., 2019】
* Albert: AlbertForMaskedLM, AlbertForQuestionAnswering 【Lan et al., 2020】
* Longformer: AllenaiLongformerBase 【Beltagy et al., 2020】
* Bart家族: BartForCausualLM, BartForConditionalGeneration, hf Bart 【Lewis et al., 2019】
* Blenderbot家族: BlenderbotForCausualLM, BlenderbotForConditionalGeneration, BlenderbotSmallForCausualLM, BlenderbotSmallForConditionalGeneration 【Shuster et al., 2022】
* CamemBert: CamemBart 【Martin et al., 2020】
* Deberta家族: DebertaForMaskedLM, DebertaForQuestionAnswering, DebertaV2ForMaskedLM, DebertaV2ForQuestionAnswering 【He et al., 2021】
* DistilBert/GPT: DistilBertForMaskedLM, DistilBertForQuestionAnswering, DistilGPT 【Sanh et al., 2020; Radford et al., 2019】
* Electra: ElectraForCausalLM, ElectraForQuestionAnswering 【Clark et al., 2020】
* GPT家族: GPT2ForSequenceClassification, hf GPT2 【Radford et al., 2019】, GPTJForCausalLM, GPTJForQuestionAnswering 【Radford et al., 2022】, GPTNeoForCausalLM, GPTNeoForSequenceClassification 【Gao et al., 2020】
* LayoutLM: LayoutLMForMaskedLM, LayoutLMForSequenceClassification 【Xu et al., 2020】
* Transformer变体: M2M100 【Fan et al., 2020】, MBart 【Liu et al., 2020】, MT5 【Xue et al., 2021】, MegatronBert 【Fan et al., 2020】, MobileBert 【Sun et al., 2020】, OPT 【Zhang et al., 2022】, PLBart 【Ahmad et al., 2021】, Pegasus 【Zhang et al., 2020b】, RoBERTa 【Liu et al., 2019】, S2T2 【Lin and Ng, 2022】, T5 【Raffel et al., 2023】, TrOCR 【Li et al., 2022】, XGLM 【Lin et al., 2022】, XLNet 【Yang et al., 2020】, YituTechConvBert 【Jiang et al., 2021】
* Inception: gluon inception v3, inception v3 【Szegedy et al., 2015】, adv inception v3 【Kurakin et al., 2018】
* Vision Transformers (ViT) 及其变体: beit 【Bao et al., 2022】, cait 【Touvron et al., 2021c】, coat 【Xu et al., 2021】, convit 【d’Ascoli et al., 2022】, convnext 【Liu et al., 2022】, crossvit 【Chen et al., 2021a】, deit 【Touvron et al., 2021b】, levit 【Graham et al., 2021】, mobilevit 【Mehta and Rastegari, 2022】, pit 【Heo et al., 2021】, swin 【Liu et al., 2021b】, tnt 【Han et al., 2021b】, twins 【Chu et al., 2021】, visformer 【Chen et al., 2021b】, vit 【Dosovitskiy et al., 2021a】, volo 【Yuan et al., 2021】, xcit 【El-Nouby et al., 2021】
* CNN模型: botnet 【Srinivas et al., 2021】, eca_botnext/sebotnet 【Srinivas et al., 2021; Wightman et al., 2021】, convmixer 【Ng et al., 2022】, cspdarknet53 【Wang et al., 2019; Bochkovskiy et al., 2020】, dla102 【Yu et al., 2019】, nfnet 【Brock et al., 2021】, dpn107 【Chen et al., 2017】, eca_halonext 【Vaswani et al., 2021; Wightman et al., 2021】, ese_vovnet 【Lee and Park, 2020】, fbnet 【Wu et al., 2019】, gernet 【Lin et al., 2020a】, ghostnet 【Han et al., 2020a】, hrnet 【Wang et al., 2020】, jx_nest 【Zhang et al., 2021】, lcnet 【Cui et al., 2021】, mixnet 【Tan and Le, 2019】, mnasnet 【Tan et al., 2019】, mobilenetv2/v3 【Sandler et al., 2019; Howard et al., 2019】, pnasnet5large 【Liu et al., 2018】, poolformer 【Yu et al., 2022】, regnet 【Radosavovic et al., 2020】, repvgg 【Ding et al., 2021】, res2net/resnet/resnext 【Gao et al., 2021; He et al., 2016; Xie et al., 2017】, resmlp 【Touvron et al., 2021a】, resnest 【Zhang et al., 2020a】, rexnet 【Han et al., 2021a】, selecsls42b 【Mehta et al., 2020】, spnasnet 【Stamoulis et al., 2019】, tf_efficientnet 【Tan and Le, 2020; Xie et al., 2020】, tinynet 【Han et al., 2020b】
* MLP-based模型: mixer/gmixer 【Tolstikhin et al., 2021】, gmlp 【Liu et al., 2021a】
* 其他经典模型: Background Matting 【Lin et al., 2020b】, LearningToPaint 【Huang et al., 2019】, alexnet 【Krizhevsky et al., 2017】, dcgan 【Radford et al., 2016】, densenet121 【Huang et al., 2018】, nvidia_deeprecommender 【Kuchaiev and Ginsburg, 2017】, pytorch_unet 【Ronneberger et al., 2015】, shufflenet_v2_x1_0 【Zhang et al., 2017】, squeezenet1_1 【Iandola et al., 2016】, vgg16 【Simonyan and Zisserman, 2015】

附录C. 测试的Python语法

我们还收集了上述模型中常用的Python特性，并将它们存储在一个简单的Python测试文件中。 该文件包含超过80个测试用例，地址为 https://github.com/thuml/depyf/blob/master/tests/test.py 。

附录D. 收集的输出

我们将PyTorch的所有输出收集在 https://github.com/thuml/learn_torch.compile 。 它包括了许多常用模型、PyTorch如何转换它们，以及在训练和推理过程中张量的形状。所有细节都在自包含的脚本中。