跳转至

《PytorchConference2023翻译系列》22. PT2 Export - 用于PyTorch的全图捕获机制

我们推出了一个新的系列,对PytorchConference2023 的博客进行中文编译,会陆续在公众号发表。建议PC端阅读本文。

PT2 Export - A Sound Full Graph Capture Mechanism for PyTorch

大纲

  1. PT2 Export的需求

  2. PT2 Export的实现方法

  3. Export产出如何

详细要点

1. PT2 Export的需求

  • 需要全图捕获机制支持不依赖Python环境部署

  • 需要soundness保证输入有效就能得到正确输出

2. PT2 Export的实现方法

  • 基于Torch Compile技术提前全面捕获程序运行图

  • 使用Dynamo追踪代码生成图表示

  • 提供ExportDB示例库和错误链接指导使用

3. Export产出如何

  • 提供静态/动态形状接口导出模型

  • 导出产出是一个可以自定义优化的FX graph

  • 支持保存加载与PyTorch模型相同方式

  • 后端可以针对core子集优化和插入自定义算子

我的名字是Avik Chaudhary,我要向你介绍一种机制,它可以为PyTorch程序实现全图捕获。这是建立在Torch Compile技术突破的基础上的一项工作,不再仅仅是在程序运行时进行即时编译,而是完全提前进行。 我们从几个类别来理解:为什么我们需要export,我们是如何构建export的,最后export是什么样子。 好的,让我们从为什么开始。为什么PyTorch需要健全的全图捕获机制呢?简单地说,因为它经常是必需的。Tosh.compile通常会与Python运行时进行很多往返交互。当你调用一个函数时,我们会获取字节码然后在一个子图中进行编译。但是如果我们遇到无法编译的东西,通常会退回到Python运行时,然后回来继续编译,然后获取另一个子图。

为什么我们需要Export?现在想象一下那些必须完全不依赖Python的环境,比如有各种硬件限制的设备,或者根本不想涉及Python的服务器。那么,在这些情况下该怎么办呢?整个图的捕捉通常是可取的。想象一下,如果你有一个模型执行计算的完全表示。如果您想要应用所有特殊目的优化,您可能不仅要优化您的代码,还要在代码中注入一些通用编译器永远不会考虑的特定假设。这里你可能需要考虑很多东西。

此外,我们还要考虑——为什么我们需要soundness稳健性?我们来看看一个例子:我们通常希望基于一些示例输入一次性捕获模型,然后将其运行在可能很多不同的输入上,怎样才能保证你的代码不会崩溃呢?如果真的崩溃了呢?想象一下,在遥远的某个时刻调试运行时错误,你都不记得你的代码是何时何地编译的了。简直是个噩梦。就算你最终找到了根本原因,并修复了你的代码,你也必须重新进行整个流程,部署代码。所以,我们所说的“soundness”是什么意思呢?

简而言之,它意味着如果你通过代码输入有效的输入,你将得到正确的输出;我们可以根据示例输入推断出什么是有效的输入。

好了,关于为什么的部分就讲到这里。让我们进入到我们是如何构建export的。这既是容易又是困难的。容易的一部分是 Torch Compile 已经存在,并且其中包含了一些真正出色的技术。艰难的部分在于我们要违背PyTorch的动态思路。在某种程度上,我们要违背即时执行的精神。

因此,我们现在要讨论graph。首先,可用性是必须关注的问题——我们如何确保您能够获得一个导出的模型呢?其次,标准化必须也是值得关注的问题——运行时如何确保它们能够运行您的导出模型呢?所以我们基本上复用了Dynamo来进行整个图的捕获。您提供示例输入,它会追踪代码,并返回给您一个graph。但有一个注意事项,有时您可能需要重写代码。这部分是因为追踪编译器的工作方式,如控制流。我们只会通过您的代码走一条可能的路径,如果不清楚应该选择哪条路径,我们实际上会发出警告。因此,如果您想捕捉代码中的更多路径,您必须使用特殊的控制流程运算符。为了指导您,我们建立了一个名为ExportDB的示例数据库,希望能够帮助您。随时间变化逐渐增加。这些示例包含了一些被支持和不被支持的用例,一些预期会正常工作和预期不会正常工作的事情。当出现错误时,通常会链接到这些示例,以便您能够正确使用它们。 总的来说,我们致力于使您能够轻松导出。现在让我们来谈谈可靠性。没错,Dynamo的设计初衷就是为了保证可靠性。它不仅根据所采取的跟踪路径生成代码,还返回所谓的"guards",即跟踪路径上的条件。这些条件保证代码的正确性,它们是必要条件。

这些guards容易由机器检查,但非人类可读,这没关系,因为在eager模式下,我们基本上使用它们来决定何时重新编译。当我们检查这些guards时,如果它们失败,我们会重新编译。 不过,对于导出而言,我们希望关注于输入张量形状上的条件。这些被称为shape guards。我们希望使用这些shape guards与用户建立一个约定。我们首先简化这些shape guards,使其成为用户可以理解的条件,然后我们在编译时验证这些条件。这个过程可能会导致错误,但幸好我们始终有可操作的错误信息。 一旦这个过程完成,我们将把这个约定转化为运行时断言。我们还默认使用静态形状,也就是说,未来输入的形状必须完全匹配当前示例输入的形状,就像静态类型推断一样。但有时候仅靠静态形状还不够,对吧?就像你可能想为程序添加泛型一样,你可能希望使用动态形状导出你的模型。让我们看一个具体的例子。

假设你的代码中有一个MatMul运算符,是最常见的运算符之一,那么你肯定需要找到这些符号,对吧?4,4是静态推导的shape,但是我们想要得到a,b,c 希望这已经满足了您的需求。那么,我们接着看看:

好的,正如我之前所说,我们在编译时假设静态形状,你所做的基本上就是将其传递给export函数;传入一些示例输入,您将获得一个导出的程序,然后您可以使用与示例完全相同shape的其他输入进行调用。就是将模型导出的所有需要做的事情。事实上,这就是我们在开源项目中自动导出一大堆模型的方式。 当然,当无法捕获整个图形时,可能会出现错误。您的错误将指向我刚刚谈到的示例。您可以查看这些示例并修复代码,但通常仅此而已。但有时候您可能希望将代码推广到不同形状的输入。在这种情况下,您将使用动态形状,对吗?一个非常常见的用例是当批次维度可以随时间变化时。

所以,您所做的就是创建一个名为batch的维度。这实际上只是一个符号,并使用它来创建一个动态形状规范,用于传递给您的导出调用的输入。然后,唰,您导出的程序可以在具有不同形状的输入上进行调用。以上是在导出时使用动态形状的API例子。最后,假设你得到了一个exported 程序。

通常你可能希望在这个export程序上编写自己的自定义pass。底层实际上只是一个普通的fx图,还有一些额外的信息,比如源代码级元数据、静态和动态形状等等。对于这个你可以随心所欲地做任何事情,可以用这个图,可以替换ops,你可以将其lower为你自己喜欢的IR,或者你可能希望将编写这些自定义传递的任务交给后端完成后再load,如果你想强制前后端分离。

你可以像保存和加载PyTorch模型一样保存和加载你的导出程序,对吧?不过这里我们不仅谈论模型权重,还有模型代码和我刚才提到的所有额外的信息,还有版本信息。

最后来看一些关于backend工作人员可能会感兴趣的东西。默认情况下,您的export程序将包含full ops以简化转换。backend开发者可以针对这些opset的core子集进行定位,以最大限度地提高覆盖范围。 将export程序转换为core子集只需简单的API调用。我们称之为run_decomposition。您还可以使用它来插入您自己的特殊操作符实现。比如说,如果您有专门用于卷积的硬件,您可以告诉API避免对该操作符进行分解,然后用您自己的替代它。还有很多关于此话题可以谈论。例如,我们还支持自定义ops。

这就是我今天关于导出的要说的内容。请尝试一下。它可以通过nightly获取。不过还没有关于稳定性的保证。我们正在努力解决这个问题,我们希望能尽快进行正式发布。不过,这不仅仅是一个原型。我们在Meta内部已经成功地使用了Xport一段时间了。事实上,它正在支持Meta RL设备和现实实验室设备上的AI inference。

它已经随着Ray-Ban的Meta智能眼镜一起发货,并将很快推出Quest 3虚拟现实头显。它还在推动我们的一系列下一代服务用例,并且通过与ExecuTorch和AOTInductor的集成实现了所有可能性。你可能听过这些讨论。 我们还有更广泛的社区:TensorRT和XLA的集成,我们也很兴奋能够与像Apple、ARM和Qualcomm这样的行业领导者在各种集成上合作。那就是这样。请让我们知道你使用export功能的体验。我们在这里倾听、支持和改进。


本文总阅读量