depyf 自豪地成为 PyTorch 生态系统项目。更多详情请参阅公告博客 https://pytorch.org/blog/introducing-depyf/。

你是否曾因 torch.compile 的复杂性而感到不知所措？深入研究其内部机制，仿佛进入了一片黑魔法领域——字节码和 Python 内部细节让许多用户难以理解，从而阻碍了他们对 torch.compile 的掌握与灵活应用。

如果你也有同样的困扰，那么 depyf 或许会对你有所帮助。正如其标志所示，depyf 是一款利用 Python 高级特性（Python 蛇形符号）来揭示 PyTorch 编译器 torch.compile（PyTorch 标志）内部细节（齿轮符号）的工具，旨在帮助用户深入理解、灵活适应并优化代码（调试器符号），从而充分发挥其性能优势。

:warning: 本项目是在与 PyTorch 团队紧密合作下开发的。因此，它依赖于 PyTorch 中非常新的功能，以更好地支持对 torch.compile 的理解。请务必配合 PyTorch>=2.2.0 使用此项目（推荐使用 PyTorch 夜间版）。有关如何安装不同版本 PyTorch 的信息，请访问 PyTorch 官网。

:warning: 在开发过程中，我们广泛征求社区意见。你可能会在一些讨论论坛或社交媒体平台上找到早期的使用示例。请以最新文档为准，了解如何使用此工具。

为什么选择 depyf？

如果你想深入了解 torch.compile 生成的字节码，那么 depyf 可能是你的不二之选。我们测试了几款现有的反编译工具，它们在跨版本处理简单 Python 字节码时表现不佳，且对 PyTorch 的支持十分有限。

反编译工具	Python 3.8	Python 3.9	Python 3.10	Python 3.11	PyTorch
decompyle3	90.6% (77/85)	×	×	×	×
uncompyle6	91.8% (78/85)	×	×	×	×
pycdc	74.1% (63/85)	74.1% (63/85)	74.1% (63/85)	67.1% (57/85)	19.3% (27/140)
depyf	100% (85/85)	100% (85/85)	100% (85/85)	100% (85/85)	100% (140/140)

安装

稳定版：pip install depyf

夜间版（推荐）：pip install git+https://github.com/thuml/depyf.git

使用方法

主要用法非常简单：只需将你的代码包裹在一个上下文管理器中即可：

import torch
from torch import _dynamo as torchdynamo
from typing import List

@torch.compile
def toy_example(a, b):
    x = a / (torch.abs(a) + 1)
    if b.sum() < 0:
        b = b * -1
    return x * b

def main():
    for _ in range(100):
        toy_example(torch.randn(10), torch.randn(10))

if __name__ == "__main__":
    # main()
    # 将要运行的代码包围在 `with depyf.prepare_debug` 中
    import depyf
    with depyf.prepare_debug("./dump_src_dir"):
        main()

随后，你便可在 ./dump_src_dir 目录下查看 torch.compile 的所有详细信息。这些内容被组织如下：

• 每个使用 torch.compile 的函数对应的 full_code_for_xxx.py
• 每个计算图相关的 Python 代码 __transformed_code_for_xxx.py
• 每个 Python 字节码文件 __transformed_code_for_xxx.py.xxx_bytecode，以代码对象形式转储，可通过 dill.load(open("/path/to/file", "wb")) 加载。请注意，load 函数可能会导入一些模块，如 transformers，请确保已安装这些模块。
• 每个计算图及其优化对应的 __compiled_fn_xxx.py：
  • Captured Graph：原始的前向计算图
  • Joint Graph：来自 AOTAutograd 的前向后向联合图
  • Forward Graph：来自 AOTAutograd 的前向图
  • Backward Graph：来自 AOTAutograd 的后向图
  • kernel xxx：来自 Inductor 的编译后的 CPU/GPU 内核封装。

我们在测试超过 100 个深度学习模型的过程中收集了 所有编译产物，你可以参考这些内容来了解 PyTorch 编译器的工作原理。

若想使用调试器逐步执行上述代码，只需再添加一个上下文管理器（并通过调试器启动脚本）：

import torch
from torch import _dynamo as torchdynamo
from typing import List

@torch.compile
def toy_example(a, b):
    x = a / (torch.abs(a) + 1)
    if b.sum() < 0:
        b = b * -1
    return x * b

def main():
    for _ in range(100):
        toy_example(torch.randn(10), torch.randn(10))

if __name__ == "__main__":
    import depyf
    with depyf.prepare_debug("./dump_src_dir"):
        main()
    # 将要调试的代码包围在 `with depyf.debug()` 中
    with depyf.debug():
        main()

调用 depyf.debug() 会暂停程序，以便你设置断点；之后，你就可以使用调试器在上述指定的 ./dump_src_dir 目录下的文件中命中断点。

联系方式

如对 depyf 有任何疑问，欢迎随时提交问题！我们欢迎任何讨论、问题报告或 Pull Request。如有其他问题，也可联系 youkaichao@gmail.com。

引用 depyf

如果你认为 depyf 对你有所帮助，请在你的论文中引用它。

@article{you2024depyf,
      title={depyf: Open the Opaque Box of PyTorch Compiler for Machine Learning Researchers}, 
      author={Kaichao You and Runsheng Bai and Meng Cao and Jianmin Wang and Ion Stoica and Mingsheng Long},
      year={2024},
      eprint={2403.13839},
      journal={arXiv},
      primaryClass={cs.LG},
      url={https://github.com/thuml/depyf}
}

depyf 快速上手指南

depyf 是一个专为 PyTorch 开发者设计的工具，旨在揭开 torch.compile 的黑盒面纱。它能将编译后的字节码反编译为可读的 Python 源码，并导出详细的计算图信息，帮助你深入理解、调试和优化 PyTorch 编译流程。

环境准备

在开始之前，请确保满足以下系统要求：

• Python 版本：支持 3.7, 3.8, 3.9, 3.10, 3.11。
• PyTorch 版本：必须使用 PyTorch >= 2.2.0。
  • 为了获得最佳支持和最新功能，强烈推荐使用 PyTorch Nightly 版本。
  • 安装命令参考：pip install --pre torch --index-url https://download.pytorch.org/whl/nightly/cu121 (根据具体 CUDA 版本调整)。
• 可选依赖：若需加载导出的 bytecode 文件，可能需要安装 dill 及相关模型库（如 transformers）。

提示：国内用户可使用清华或阿里镜像源加速 PyTorch 安装，例如： pip install --pre torch --index-url https://mirror.tuna.tsinghua.edu.cn/pytorch-whl/nightly/cu121

安装步骤

你可以选择安装稳定版或推荐的夜间版（Nightly）：

选项 1：安装稳定版

pip install depyf

选项 2：安装夜间版（推荐，获取最新功能）

pip install git+https://github.com/thuml/depyf.git

国内加速：如果 GitHub 连接缓慢，可使用国内镜像代理：

pip install git+https://ghp.ci/https://github.com/thuml/depyf.git

基本使用

depyf 的核心用法非常简单，只需使用上下文管理器包裹你的代码即可。

1. 导出编译细节

以下示例展示了如何捕获 torch.compile 生成的源码、字节码和计算图，并保存到指定目录。

import torch
from torch import _dynamo as torchdynamo
from typing import List

@torch.compile
def toy_example(a, b):
    x = a / (torch.abs(a) + 1)
    if b.sum() < 0:
        b = b * -1
    return x * b

def main():
    for _ in range(100):
        toy_example(torch.randn(10), torch.randn(10))

if __name__ == "__main__":
    import depyf
    # 使用 prepare_debug 上下文管理器包裹需要运行的代码
    # 编译产物将保存到 ./dump_src_dir 目录
    with depyf.prepare_debug("./dump_src_dir"):
        main()

运行后，./dump_src_dir 目录下将生成以下关键文件：

• full_code_for_xxx.py: torch.compile 装饰函数的完整代码。
• __transformed_code_for_xxx.py: 与每个计算图关联的 Python 代码（反编译结果）。
• __compiled_fn_xxx.py: 包含捕获的计算图（Captured Graph）、联合图（Joint Graph）以及 Inductor 编译后的内核包装器。

2. 启用断点调试

如果你希望使用标准调试器（如 pdb 或 IDE 调试器）逐步跟踪编译后的代码逻辑，可以额外添加 depyf.debug() 上下文管理器。

import torch
from torch import _dynamo as torchdynamo
from typing import List

@torch.compile
def toy_example(a, b):
    x = a / (torch.abs(a) + 1)
    if b.sum() < 0:
        b = b * -1
    return x * b

def main():
    for _ in range(100):
        toy_example(torch.randn(10), torch.randn(10))

if __name__ == "__main__":
    import depyf
    # 第一步：准备调试环境并导出源码
    with depyf.prepare_debug("./dump_src_dir"):
        main()
    
    # 第二步：进入调试模式
    # 程序会在此暂停，允许你在 ./dump_src_dir 生成的文件中设置断点
    with depyf.debug():
        main()

操作提示：

1. 运行脚本，程序会在 depyf.debug() 处暂停。
2. 在你的 IDE 或调试器中，打开 ./dump_src_dir 下生成的 .py 文件。
3. 在感兴趣的代码行设置断点，然后继续执行即可进行单步调试。