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

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

@@ -0,0 +1,105 @@
+> 我们推出了一个新的系列，对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运行时,然后回来继续编译，然后获取另一个子图。
+
+
+![](https://files.mdnice.com/user/53043/af3ed74d-b19c-40f7-a11c-4938475f7906.png)
+
+为什么我们需要Export？现在想象一下那些必须完全不依赖Python的环境，比如有各种硬件限制的设备，或者根本不想涉及Python的服务器。那么，在这些情况下该怎么办呢？整个图的捕捉通常是可取的。想象一下，如果你有一个模型执行计算的完全表示。如果您想要应用所有特殊目的优化，您可能不仅要优化您的代码，还要在代码中注入一些通用编译器永远不会考虑的特定假设。这里你可能需要考虑很多东西。
+
+![](https://files.mdnice.com/user/53043/35ab850e-0c33-4fb2-8e37-9f1e4a3018c1.png)
+
+此外，我们还要考虑——为什么我们需要soundness稳健性？我们来看看一个例子：我们通常希望基于一些示例输入一次性捕获模型，然后将其运行在可能很多不同的输入上，怎样才能保证你的代码不会崩溃呢？如果真的崩溃了呢？想象一下，在遥远的某个时刻调试运行时错误，你都不记得你的代码是何时何地编译的了。简直是个噩梦。就算你最终找到了根本原因，并修复了你的代码，你也必须重新进行整个流程，部署代码。所以，我们所说的“soundness”是什么意思呢？
+
+![](https://files.mdnice.com/user/53043/658c0f5f-07c2-4148-9f8d-f2f0c306cec5.png)
+
+简而言之，它意味着如果你通过代码输入有效的输入，你将得到正确的输出；我们可以根据示例输入推断出什么是有效的输入。
+
+好了，关于为什么的部分就讲到这里。让我们进入到我们是如何构建export的。这既是容易又是困难的。容易的一部分是 Torch Compile 已经存在，并且其中包含了一些真正出色的技术。艰难的部分在于我们要违背PyTorch的动态思路。在某种程度上，我们要违背即时执行的精神。
+
+
+![](https://files.mdnice.com/user/53043/cda0403a-723c-4a15-86ef-4f9680570727.png)
+
+因此，我们现在要讨论graph。首先，可用性是必须关注的问题——我们如何确保您能够获得一个导出的模型呢？其次，标准化必须也是值得关注的问题——运行时如何确保它们能够运行您的导出模型呢？所以我们基本上复用了Dynamo来进行整个图的捕获。您提供示例输入，它会追踪代码，并返回给您一个graph。但有一个注意事项，有时您可能需要重写代码。这部分是因为追踪编译器的工作方式，如控制流。我们只会通过您的代码走一条可能的路径，如果不清楚应该选择哪条路径，我们实际上会发出警告。因此，如果您想捕捉代码中的更多路径，您必须使用特殊的控制流程运算符。为了指导您，我们建立了一个名为ExportDB的示例数据库，希望能够帮助您。随时间变化逐渐增加。这些示例包含了一些被支持和不被支持的用例，一些预期会正常工作和预期不会正常工作的事情。当出现错误时，通常会链接到这些示例，以便您能够正确使用它们。
+总的来说，我们致力于使您能够轻松导出。现在让我们来谈谈可靠性。没错，Dynamo的设计初衷就是为了保证可靠性。它不仅根据所采取的跟踪路径生成代码，还返回所谓的"guards"，即跟踪路径上的条件。这些条件保证代码的正确性，它们是必要条件。
+
+![](https://files.mdnice.com/user/53043/82c8dfbc-8cb4-4196-935e-ddc29bd75994.png)
+
+这些guards容易由机器检查，但非人类可读，这没关系，因为在eager模式下，我们基本上使用它们来决定何时重新编译。当我们检查这些guards时，如果它们失败，我们会重新编译。
+不过，对于导出而言，我们希望关注于输入张量形状上的条件。这些被称为shape guards。我们希望使用这些shape guards与用户建立一个约定。我们首先简化这些shape guards，使其成为用户可以理解的条件，然后我们在编译时验证这些条件。这个过程可能会导致错误，但幸好我们始终有可操作的错误信息。
+一旦这个过程完成，我们将把这个约定转化为运行时断言。我们还默认使用静态形状，也就是说，未来输入的形状必须完全匹配当前示例输入的形状，就像静态类型推断一样。但有时候仅靠静态形状还不够，对吧？就像你可能想为程序添加泛型一样，你可能希望使用动态形状导出你的模型。让我们看一个具体的例子。
+
+![](https://files.mdnice.com/user/53043/d89468cf-af7b-4d42-bb94-b6ddc37076d9.png)
+
+假设你的代码中有一个MatMul运算符，是最常见的运算符之一，那么你肯定需要找到这些符号，对吧？4，4是静态推导的shape，但是我们想要得到a,b,c
+希望这已经满足了您的需求。那么，我们接着看看：
+
+![](https://files.mdnice.com/user/53043/8bfdc83a-473c-4dab-a4e3-d1db528dd245.png)
+
+好的，正如我之前所说，我们在编译时假设静态形状，你所做的基本上就是将其传递给export函数；传入一些示例输入，您将获得一个导出的程序，然后您可以使用与示例完全相同shape的其他输入进行调用。就是将模型导出的所有需要做的事情。事实上，这就是我们在开源项目中自动导出一大堆模型的方式。
+当然，当无法捕获整个图形时，可能会出现错误。您的错误将指向我刚刚谈到的示例。您可以查看这些示例并修复代码，但通常仅此而已。但有时候您可能希望将代码推广到不同形状的输入。在这种情况下，您将使用动态形状，对吗？一个非常常见的用例是当批次维度可以随时间变化时。
+
+![](https://files.mdnice.com/user/53043/42eccea8-b9ca-4642-acb0-6e84f2e612bf.png)
+
+所以，您所做的就是创建一个名为batch的维度。这实际上只是一个符号，并使用它来创建一个动态形状规范，用于传递给您的导出调用的输入。然后，唰，您导出的程序可以在具有不同形状的输入上进行调用。以上是在导出时使用动态形状的API例子。最后，假设你得到了一个exported 程序。
+
+![](https://files.mdnice.com/user/53043/3cc20b21-6596-4b9d-9895-df912b202956.png)
+
+通常你可能希望在这个export程序上编写自己的自定义pass。底层实际上只是一个普通的fx图，还有一些额外的信息，比如源代码级元数据、静态和动态形状等等。对于这个你可以随心所欲地做任何事情，可以用这个图，可以替换ops，你可以将其lower为你自己喜欢的IR，或者你可能希望将编写这些自定义传递的任务交给后端完成后再load，如果你想强制前后端分离。
+
+![](https://files.mdnice.com/user/53043/774dcd70-9215-4cdd-a97c-ff48b6ef1c07.png)
+
+你可以像保存和加载PyTorch模型一样保存和加载你的导出程序，对吧？不过这里我们不仅谈论模型权重，还有模型代码和我刚才提到的所有额外的信息，还有版本信息。
+
+
+![](https://files.mdnice.com/user/53043/36a5eb68-8e7e-427b-9af3-6217b597ec42.png)
+
+最后来看一些关于backend工作人员可能会感兴趣的东西。默认情况下，您的export程序将包含full ops以简化转换。backend开发者可以针对这些opset的core子集进行定位，以最大限度地提高覆盖范围。
+将export程序转换为core子集只需简单的API调用。我们称之为run_decomposition。您还可以使用它来插入您自己的特殊操作符实现。比如说，如果您有专门用于卷积的硬件，您可以告诉API避免对该操作符进行分解，然后用您自己的替代它。还有很多关于此话题可以谈论。例如，我们还支持自定义ops。
+
+![](https://files.mdnice.com/user/53043/d80e3a54-e628-4fca-95aa-9fe4572103a2.png)
+
+这就是我今天关于导出的要说的内容。请尝试一下。它可以通过nightly获取。不过还没有关于稳定性的保证。我们正在努力解决这个问题，我们希望能尽快进行正式发布。不过，这不仅仅是一个原型。我们在Meta内部已经成功地使用了Xport一段时间了。事实上，它正在支持Meta RL设备和现实实验室设备上的AI inference。
+
+![](https://files.mdnice.com/user/53043/d2a38b16-657d-43cc-b83f-0c33ad64ccd9.png)
+
+它已经随着Ray-Ban的Meta智能眼镜一起发货，并将很快推出Quest 3虚拟现实头显。它还在推动我们的一系列下一代服务用例，并且通过与ExecuTorch和AOTInductor的集成实现了所有可能性。你可能听过这些讨论。
+我们还有更广泛的社区：TensorRT和XLA的集成，我们也很兴奋能够与像Apple、ARM和Qualcomm这样的行业领导者在各种集成上合作。那就是这样。请让我们知道你使用export功能的体验。我们在这里倾听、支持和改进。