depyf很荣幸成为PyTorch生态系统项目。查看公告博客https://pytorch.org/blog/introducing-depyf/了解更多详情。

你是否曾经对torch.compile的复杂性感到不知所措？深入了解其工作原理可能感觉像是黑魔法，涉及字节码和Python内部细节，许多用户难以理解，阻碍了他们对torch.compile的理解和适应。

如果你也面临这个问题，那么你可能会对depyf感兴趣。正如logo所示，depyf是一个软件工具，利用高级Python特性（Python蛇符号）来揭示PyTorch编译器torch.compile（PyTorch logo）的内部细节（内部齿轮符号），使用户能够理解它，适应它，并调整他们的代码（调试器符号）以获得最大的性能收益。

:warning: 本项目是与PyTorch团队密切合作开发的。因此，它需要PyTorch的最新特性来支持更好地理解torch.compile。请与PyTorch>=2.2.0（推荐使用PyTorch nightly版本）一起使用此项目。访问PyTorch网站了解如何安装不同版本的PyTorch。

:warning: 在开发过程中，我们经常向社区寻求建议。你可能会在一些讨论论坛或社交媒体平台上找到一些早期使用示例。请遵循最新文档来使用此工具。

为什么选择depyf？

如果你想理解torch.compile生成的字节码，那么depyf可能是你唯一的选择。我们测试了几个现有的反编译器，它们在跨版本反编译简单的Python字节码时都遇到了困难，并且对PyTorch的支持很差。

安装

稳定版本：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文件
• 每个图都有一个__transformed_code_for_xxx.py文件，用于相关的Python代码
• __transformed_code_for_xxx.py.xxx_bytecode文件用于Python字节码，转储的代码对象，可以通过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有任何疑问，欢迎开issue联系我们！我们欢迎任何讨论、问题报告或PR。或者如果你有其他问题，可以联系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}
}