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

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如何转换它们，以及在训练和推理过程中张量的形状。所有细节都在自包含的脚本中。