PyTorch 团队很高兴地宣布,我们关于 PyTorch 2 的论文已被接受,并将在 ACM 国际编程语言和操作系统架构支持大会 (ASPLOS) 上发表,该会议定于 2024 年 4 月 27 日至 5 月 1 日在美国加利福尼亚州圣地亚哥举行。
该论文深入探讨了 torch.compile 的实现,并强调了驱动它的关键技术,包括 TorchDynamo(图捕获)、TorchInductor(后端编译器)和动态形状支持。
在 ASPLOS 会议期间,我们将在 4 月 27 日星期六举办一个教程,重点介绍 PyTorch 2 的内部工作原理以及系统研究人员如何利用和在此基础上进行构建。随着活动的临近,请继续关注更多详细信息——我们期待您的参与!
论文预览附在下方
标题: PyTorch 2:通过动态 Python 字节码转换和图编译实现更快的机器学习。 完整论文 PDF
摘要
本文介绍了流行 PyTorch 机器学习框架的两个扩展,TorchDynamo 和 TorchInductor,它们实现了 PyTorch 2 中发布的 torch.compile 功能。TorchDynamo 是一个 Python 级别的即时 (JIT) 编译器,它可以在 PyTorch 程序中实现图编译,而不会牺牲 Python 的灵活性。它通过在执行前动态修改 Python 字节码,并将 PyTorch 操作序列提取到 FX 图中来实现这一点,然后使用众多可扩展后端之一对该图进行 JIT 编译。TorchInductor 是 TorchDynamo 的默认编译器后端,它将 PyTorch 程序转换为适用于 GPU 的 OpenAI Triton 和适用于 CPU 的 C++。结果表明,TorchDynamo 能够比以前的方法更稳健地捕获图,同时增加的开销极小,并且 TorchInductor 能够在 NVIDIA A100 GPU 上,在 180 多个实际模型中,提供 2.27 倍的推理几何平均加速和 1.41 倍的训练几何平均加速,优于其他六个编译器。这些扩展提供了一种在 PyTorch 等急切模式框架中通过编译器应用优化的新方法。
作者
Jason Ansel (Meta); Edward Yang (Meta); Horace He (Meta); Natalia Gimelshein (OpenAI); Animesh Jain (Meta); Michael Voznesensky (Meta); Bin Bao (Meta); Peter Bell (Quansight); David Berard (Meta); Evgeni Burovski Quansight; Geeta Chauhan (Meta); Anjali Chourdia (Meta); Will Constable (Meta); Alban Desmaison (Meta); Zachary DeVito (Meta); Elias Ellison (Meta); Will Feng (Meta); Jiong Gong (Intel); Michael Gschwind (Meta); Brian Hirsh (Meta); Sherlock Huang (Meta); Kshiteej Kalambarkar (Quansight); Laurent Kirsch (Meta); Michael Lazos (Meta); Mario Lezcano (Quansight); Yanbo Liang (Meta); Jason Liang (Meta); Yinghai Lu (Meta); CK Luk (Meta); Bert Maher (Meta); Yunjie Pan (University of Michigan); Christian Puhrsch (Meta); Matthias Reso (Meta); Mark Saroufim (Meta); Marcos Yukio Siraichi (Quansight); Helen Suk (Meta); Michael Suo (Meta); Phil Tillet (OpenAI); Eikan Wang (Intel); Xiaodong Wang (Meta); William Wen (Meta); Shunting Zhang (Meta); Xu Zhao (Meta); Keren Zhou (OpenAI & George Mason University); Richard Zou (Meta); Ajit Mathews (Meta); Gregory Chanan (Meta); Peng Wu (Meta); Soumith Chintala (Meta)
ASPLOS’24 – 全天教程日程
ASPLOS’24 PyTorch 2 教程(4 月 27 日星期六)的完整日程可在此处查看:此处