《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功能的体验。我们在这里倾听、支持和改进。