清华大学张晨团队开源深度学习编译器MagPy：实现Python程序至计算图提取技术的革新突破

近日，清华大学计算机系PACMAN实验室发布了一款名为MagPy的开源深度学习编译器。这款编译器能够自动将用户用Python编写的深度学习程序转化为计算图表示，从而实现模型的自动加速。MagPy的设计旨在解决当前深度学习编译器需要用户手动将Python程序转换为计算图的问题，提高用户的编程体验和模型运行效率。

MagPy 的优势

MagPy 的一大亮点在于它可以直接处理用户编写的Python程序，无需用户手动将其转化为计算图。这一功能极大地简化了用户的编程流程，使得即使是非专业的程序员也能轻松使用。此外，MagPy 还具备以下特点：

1. 高覆盖率：MagPy 能够处理多种复杂的Python语言特性，相比其他编译器如TorchScript和TorchDynamo，MagPy 的覆盖率更高。
2. 高效性能：MagPy 在多种深度学习模型上进行了测试，结果显示其性能显著优于现有工具，平均加速比可达1.55倍，最高可达2.88倍。
3. 易于集成：MagPy 支持一键编译，用户只需简单配置即可实现模型加速。

研究背景

近年来，深度学习技术在多个领域取得了突破性进展。为了简化编程过程，用户倾向于使用Python编写深度学习模型，并借助PyTorch等张量库进行张量操作。然而，直接执行这些程序往往性能不佳。为了提升模型的运行速度，深度学习编译器通常采用计算图表示进行优化。虽然这种方法能显著提升性能，但需要用户手动将程序转化为计算图，这对于非专业程序员来说是一项挑战。

MagPy 的解决方案

MagPy 的核心思想是通过分析Python解释器中的执行状态信息，让编译器更好地理解用户程序。具体而言，MagPy 基于以下几个观察：

1. 有限动态性：大多数深度学习程序的动态性有限，计算图结构在不同批次间保持不变。MagPy 通过监控张量操作，简化了计算图提取过程。
2. 外部值的影响：程序行为受到外部值的影响较小。MagPy 通过验证外部值是否发生变化，确保计算图结构的一致性。
3. 执行状态分析：MagPy 利用Python解释器保留的执行状态信息，设计了守卫函数和模拟函数，提高了计算图提取的准确性和效率。

实验结果

MagPy 在多个基准测试中表现优异。在对ParityBench中的1191个真实用户程序进行测试时，MagPy 成功将93.40%的程序转化为完整的操作符图，远高于TorchScript（35%）和TorchDynamo（77.2%）。在端到端测试中，MagPy 的加速效果也非常显著，平均加速比达到1.55倍，最高可达2.88倍。

未来展望

MagPy 的发布标志着深度学习编译器技术的新突破。PACMAN实验室将继续在机器学习系统领域深入研究，推动更多创新成果的诞生。用户可通过以下链接获取更多信息：

– 翟季冬教授主页：[https：//pacman.cs.tsinghua.edu.cn/~zjd](https：//pacman.cs.tsinghua.edu.cn/~zjd)
– 论文地址：[https：//www.usenix.org/system/files/atc24-zhang-chen.pdf](https：//www.usenix.org/system/files/atc24-zhang-chen.pdf)
– 项目地址：[https：//github.com/heheda12345/MagPy](https：//github.com/heheda12345/MagPy)

如果您有优秀的工作想要分享，欢迎投稿或联系报道。投稿邮箱：liyazhou@jiqizhixin.com；zhaoyunfeng@jiqizhixin.com

本文来源：