Pytorch 2.0如何通过运算符融合和CPU/GPU代码生成来加速深度学习

码农小羊
1,335 阅读19分钟

image.png

计算机编程是很神奇的。我们用人类可读的语言编写代码，就像变魔术一样，它被转化为电流通过硅晶体管，使其表现得像开关，并允许它们实现复杂的逻辑--只是为了让我们能够在互联网上欣赏猫咪视频。在编程语言和运行它的硬件处理器之间，有一项重要的技术--编译器。编译器的工作是将我们人类可读的语言代码翻译并简化为处理器能够理解的指令。

编译器在深度学习中起着非常重要的作用，可以提高训练和推理性能，提高能源效率，并针对不同的AI加速器硬件。在这篇博文中，我将讨论支持PyTorch 2.0的深度学习编译器技术。我将带领大家了解编译过程的不同阶段，并通过代码实例和可视化的方式讨论各种底层技术。

什么是深度学习编译器？

深度学习编译器将深度学习框架中编写的高级代码翻译成优化的低级硬件特定代码，以加速训练和推理。它在深度学习模型中找到机会，通过执行层和运算符融合，更好的内存规划，以及生成目标特定的优化融合内核来减少函数调用开销，从而优化性能。

image.png

与传统的软件编译器不同，深度学习编译器必须与通常在专门的人工智能加速器硬件（GPU、TPU、AWS Trainium/Inferentia、Intel Habana Gaudi等）上加速的高度可并行的代码一起工作。为了提高性能，深度学习编译器必须利用硬件的特定功能，如混合精度支持、性能优化内核，并尽量减少主机（CPU）和AI加速器之间的通信。

在深度学习算法持续快速发展的同时，硬件AI加速器也在同步发展，以跟上深度学习算法的性能和效率需求。我在之前的一篇博文中讨论了算法和AI加速器的共同发展。

在这篇博文中，我将重点讨论软件方面的问题，尤其是更接近硬件的软件子集--深度学习编译器。首先，让我们先来看看深度学习编译器中的不同功能。

PyTorch 2.0中的深度学习编译器

PyTorch 2.0包括新的编译器技术，以提高模型性能和运行效率，并通过一个简单的API针对不同的硬件后端：torch.compile（）。虽然其他博文和文章已经详细 讨论了PyTorch 2.0的性能优势，但在这里我将重点讨论当你调用PyTorch 2.0编译器时，引擎盖下发生了什么。如果你正在寻找量化的性能优势，你可以从huggingface、timm和torchbench找到不同模型的性能仪表板。

在高层次上，PyTorch 2.0深度学习编译器的默认选项执行了以下关键任务：

  1. 图形捕获：为你的模型和函数提供计算图表示。PyTorch技术：TorchDynamo, Torch FX, FX IR
  2. 自动分化：使用自动微分和降低到基元运算符进行向后的图形追踪。PyTorch技术：AOTAutograd、Aten IR
  3. 优化：前向和后向图层面的优化和运算符融合。PyTorch技术：TorchInductor（默认）或其他编译器
  4. 代码生成：生成硬件特定的C++/GPU代码。PyTorch技术：TorchInductor、OpenAI Triton（默认）其他编译器

通过这些步骤，编译器对你的代码进行转换，并生成逐步 "降低 "的中间表示（IR）。降低是编译器词典中的一个术语，指的是通过编译器的自动转换和重写，将一组广泛的操作（如PyTorch API支持的操作）映射为一组狭窄的操作（如硬件支持的操作）。PyTorch 2.0的编译器流程：

image.png

如果你是编译器术语的新手，先不要让这些东西吓到你。我也不是一个编译器工程师。继续读下去，事情会变得很清楚，因为我将用一个简单的例子和可视化的方式来分解这个过程。

通过torch.compile() 编译器过程的演练

注意：整个演练是在Jupyter笔记本中进行的，在这里托管。

为了简单起见，我将定义一个非常简单的函数，并通过PyTorch 2.0编译器过程运行它。你可以用一个深度神经网络模型或nn.Module子类来代替这个函数，但这个例子应该能帮助你更好地理解引擎盖下发生的事情，而不是一个复杂的数百万个参数的模型。

image.png

该函数的PyTorch代码：

def f(x):  return torch.sin(x)**2 + torch.cos(x)**2

如果你在高中的三角学课上注意了，你就会知道，对于所有实值的x，我们的函数的值总是要为1。这在验证函数及其导数的作用时将会很方便。

现在，是时候调用torch.compile() 。首先让我们说服自己，编译这个函数不会改变它的输出。对于同样的1x1000随机向量，我们的函数的输出和1的向量之间的平均平方误差对于已编译和未编译的函数都应该是零（在一定的误差容忍度下）。

image.png

我们所做的只是添加了一行额外的代码torch.compile() 来调用我们的编译器。现在让我们来看看在每个阶段的引擎盖下发生了什么。

图形捕获：为您的PyTorch模型或函数提供计算图示

**PyTorch技术：**TorchDynamo、FX Graphs、FX IR

编译器的第一步是确定要编译的内容。输入TorchDynamo。TorchDynamo拦截你的Python代码的执行，并将其转换为FX中间表示法（IR），并将其存储在一个名为FX Graph的特殊数据结构中。你问这是什么样子的？很高兴你这么问。下面，我们来看看我们用来生成这个的代码，但这里是转换和输出：

image.png

值得注意的是，Torch FX图只是IR的容器，并没有真正指定它应该容纳什么运算符。在下一节中，我们将看到FX图的容器再次出现在一组不同的IR上。如果你比较一下函数代码和FX IR，两者之间几乎没有什么区别。事实上，这是你写的相同的PyTorch代码，但以FX图形数据结构所期望的格式来布置。它们在执行时都会提供相同的结果。

如果你调用torch.compile()而没有任何参数，它将使用默认设置，运行整个编译器栈，其中包括默认的硬件后端编译器，称为TorchInductor。但如果我们现在讨论TorchInductor，就会跳过这个话题，所以我们先把这个话题搁置起来，等我们准备好了再来讨论。首先我们需要讨论图形捕获，我们可以通过拦截Torch.compile()的调用来实现。我们将这样做：torch.compile()也允许你提供你自己的编译器，但因为我不是一个编译器工程师，而且我对如何写一个编译器没有丝毫的了解，我将提供一个假的编译器函数来捕获TorchDynamo生成的FX图IR。

下面是我们的假编译器后端函数，名为 inspect_backend，用于 torch.compile() ，在该函数中我做了两件事：

  1. 打印由TorchDynamo捕获的外汇IR代码
  2. 保存FX图形的可视化

上述代码段的输出是FX IR代码和显示我们功能的图示sin^2(x)+cos^2(x)

image.png

请注意，我们的假编译器 inspect_backend 函数只有在我们用一些数据调用编译后的函数时才会被调用，即当我们调用compiled_model(x) 。在上面的代码片段中，我们只是在评估这个函数，或者用深度学习的术语来说，是在做一个 "前向通过"。在下一节，我们将利用PyTorch的自动微分引擎torch.autograd来计算导数和 "后向通 "图。

自动微分：前向和后向计算图

**PyTorch技术：**AOTAutograd、Core Aten IR

TorchDynamo给我们提供了作为FX图的前向传递函数评估，但后向传递的情况呢？为了完整起见，我将偏离我们的主要话题，谈一谈为什么我们需要评估一个函数相对于其权重的梯度。如果你已经熟悉了数学优化的工作原理，请跳过这一节。

什么是后向传递和后向图？

深度学习和机器学习的 "学习 "部分是一个数学优化问题，简单说来就是：找到一个变量w的值，产生w的某个函数的最低值，或者更简洁的说：

image.png

在机器学习中，f(w)是以权重为参数的损失函数。f(w)可以更清楚地表示为基于训练数据的训练标签和模型的预测标签之间的一些误差措施：

image.png

事实证明，如果我们能计算出损失相对于权重的 "减少率"，我们就能更新我们的权重，使之一步步接近越来越小的损失f(w)。换句话说，我们必须向更适合我们训练数据集的模型靠近。我们可以通过计算损失f(w)在给定w处的最陡峭的斜率来找到下一个权重值，并对w进行扰动，使其朝这个方向发展。一个函数相对于权重的斜率，就是它相对于权重的导数。由于有多个权重值，导数成为一个矢量，称为梯度，它是一个部分导数的矢量，每个权重都有组成部分。权重w在每次迭代时都被梯度的某个函数g()所扰动，如下所示：

image.png

其中函数g(.)取决于优化器（如SGD、SGDM、RMSPROP、ADAM等）。

对于SGD来说，权重更新步骤成为：

image.png

PyTorch 2.0是如何追踪后向通行图的？

首先让我们计算一下我们期望的后向传递图应该是什么样子，然后与PyTorch生成的图进行比较。对于我们的简单函数，前向图和后向图应该实现以下功能。如果sin和cos困扰着你，你可以想象f(x)是应用于神经网络的损失函数。

image.png

PyTorch使用反向模式的自动微分来计算梯度，PyTorch的自动化微分的实现被称为Autograd。PyTorch 2.0引入了AOTAutograd，它提前追踪前向和后向图形，即在执行之前，并生成一个联合前向和后向图形。然后，它将前向图和后向图分割成两个独立的图。前进和后退图都存储在FX图的数据结构中，并可以被可视化，如下图所示。

image.png

你可以通过图上的节点来验证数学的检查结果。AOTAutograd生成的后向传递确实计算了我之前分享的方程中显示的导数，这应该等于零，因为原始函数只产生身份。

我们现在通过扩展我们的假编译器函数 inspect_backend 来运行 AOTAutograd，调用 AOTAutograd 并生成我们的后向图。更新后的 inspect_backed 定义了一个前向（fw）和后向（bw）编译器捕获函数，它从 AOTAutograd 读取前向和后向图形，并打印出降低的 ATen IR，并为前向和后向图形保存 FX 图。

这将产生以下的前向和后向图形。例如，FX图IR中的torch.sin(x)和我们原始代码中的torch.ops.aten.sin.default()已经被替换掉了。如果你还不熟悉的话，你可能会问，这个叫做aten的有趣东西是什么？ATen是A Tensor库的缩写，它是一个名字非常有创意的低级库，有一个C++接口，实现了许多在CPU和GPU上运行的基本操作。

在急切的操作模式下，你的PyTorch操作会被路由到这个库，然后调用相应的CPU或GPU实现。AOTAutograd会自动生成代码，用ATen IR来取代高级别的PyTorch API，用于前进和后退的图形，你可以在下面的输出中看到：

image.png

你还可以看到，除了前向通道的输出外，前向图还输出了一些额外的张量[add, sin, cos, primals_1] 。这些张量被保存在后向通道中用于梯度计算。在前面分享的图中，你也可以从前向和后向的计算图中看到这一点。

PyTorch中的IR有哪些不同类型？

ATen IR是我们在上一节中讨论的ATen库所支持的操作列表，您可以在这里看到ATen库实现的全部操作列表。在PyTorch中还有两个IR概念，您应该了解一下：1/ Core Aten IR 2/ Prims IR。Core Aten IR是更广泛的Aten IR和Prims IR的一个子集，也是Core Aten IR的一个更小的子集。假设你正在设计一个处理器，想在你的硬件上支持PyTorch代码加速。在硬件上支持PyTorch API的全部列表几乎是不可能的，所以你可以做的是建立一个编译器，只支持Core Aten IR或Prims IR中定义的较小的基本运算符子集，并让AOTAutograd将复合运算符分解为核心运算符，我们将在下一节看到。

ATen IR、Core ATen IR、Prims IR之间有什么区别？

image.png

Core Aten IR（以前的canonical Aten IR）是Aten IR的一个子集，可以用来分解Aten IR中的所有其他运算符。针对特定硬件加速器的编译器可以专注于只支持Core Aten IR，并将其映射到他们的低级硬件API。这使得向PyTorch添加硬件支持变得更加容易，因为他们不需要实现对完整的PyTorch API的支持，而PyTorch API将随着越来越多的抽象而继续增长。

Prims IR是Core Aten IR的一个更小的子集，它将Core Aten IR的操作进一步分解为基本操作，使针对特定硬件的编译器更容易支持PyTorch。但是，将运算符分解成越来越低的操作，由于过多的内存写入和函数调用开销，肯定会导致性能下降。但我们期望硬件编译器可以把这些运算符融合在一起，以支持硬件API，从而重新获得性能。

虽然我们不需要将我们的函数进一步分解为Core Aten IR和Prims IR，但我将在下面演示如何分解。

(可选话题）分解为核心Aten IR和Prims IR

如果你在设计硬件或硬件编译器，要在硬件中支持PyTorch API的全部列表几乎是不可能的，尤其是考虑到深度学习和人工智能的发展速度。但对于硬件设计者来说，优势在于大多数深度学习功能可以映射到很少的基本数学操作中，而计算量最大的是矩阵-矩阵和矩阵-矢量操作。像PyTorch API支持的那些复合运算符可以使用AOTAutograd分解成这些基本运算，我们将在本节讨论。如果你不处理低级别的硬件，你可以跳过这一节。

你可以更新AOTAutograd函数，传入一个分解字典，可以把Aten IR降低为Core Aten IR和Prims IR。我在这里只分享相关的代码片段和输出，因为你可以在GitHub上找到完整的笔记本。默认情况下，运算符不会被分解成Core Aten IR或Prims IR，但你可以传递一个分解的字典。

在下面的代码片段中，我将我们的函数f转换成了损失函数f_loss，将均方误差（MSE）的计算纳入了我们的函数。我这样做是为了演示AOTAutograd如何将MSE分解成其基本运算符。

image.png

image.png

分解的结果是，mese_loss被分解成更多的基本运算：减法，power(2)，均值。

image.png

这是因为MSE或两个向量x和y之间的均方误差被定义为以下内容，它只需要这3个操作减去，其中power是一个元素明智的操作。如果你为你的硬件写一个编译器，你很可能已经支持这3个操作，通过分解，你的PyTorch代码将运行，而不需要进一步修改。

image.png

你也可以看到这一点反映在FX图的可视化中

image.png

现在让我们把它进一步分解成Prims IR，这是一个更小的子集~250个操作。同样，我在这里只分享相关的代码片段和输出，因为你可以在GitHub上找到完整的笔记本。

image.png

下面是prim IR分解的输出。所有红色的aten操作被替换或分解为绿色的prim操作。

image.png

图形优化：层和运算符的融合以及C++/GPU代码生成

讨论的PyTorch技术：TorchInductor、OpenAI Triton（默认）其他编译器

在这篇博文的最后一节，我们将讨论运算符融合和使用TorchInductor为CPU和GPU自动生成代码。首先是一些基础知识：

什么是深度学习的优化编译器？

深度学习的优化编译器善于发现代码中的性能差距，并通过改造代码以减少代码属性，如内存访问、内核启动、目标后端的数据布局优化来解决这些问题。TorchInductor是默认的优化编译器，使用torch.compile()可以为使用OpenAI Triton的GPU和使用OpenMP pragma指令的CPU生成优化内核。

什么是深度学习中的运算器融合？

深度学习由许多基本操作组成，如矩阵-矩阵和矩阵-矢量乘法。在PyTorch的急切执行模式下，每个操作都会导致硬件上的单独函数调用或内核启动。这导致了启动内核的CPU开销，并导致内核启动之间更多的内存读写。像TorchInductor这样的深度学习优化编译器可以将多个操作融合到python中的一个复合运算符中，并为其生成低级别的GPU内核或C++/OpenMP代码。由于内核启动次数减少，内存读/写次数减少，这就导致了计算速度加快。

image.png

上一节中AOTAutograd输出的计算图是由许多Aten运算符组成的，以FX图表示。TorchInductor的优化并不改变图中的底层计算，而只是用运算符和层的融合来重组它，并为它生成CPU或GPU代码。由于TorchInductor可以提前看到完整的前向和后向计算图，所以它可以对不相互依赖的操作的失序执行做出决定，并使硬件资源的利用率最大化。

在引擎盖下，对于GPU目标，TorchInductor使用OpenAI的Triton来生成融合的GPU内核。Triton本身是一个独立的基于Python的框架和编译器，用于编写优化的低级GPU代码，否则就用CUDA C/C++编写。但唯一的区别是，TorchInductor将生成Triton代码，并将其编译成低级别的PTX代码。

对于多核CPU目标，TorchInductor会生成C++代码并注入OpenMP pragma指令以生成并行内核。从PyTorch的用户层面来看，这就是IR转换流程：

image.png

当然这只是高层次的观点，我在这里省略了一些细节，我鼓励你阅读TorchInductor论坛的帖子和OpenAI的triton博客帖子来了解triton。

现在我们将扔掉我们在上一节中使用的假编译器，而使用使用TorchInductor的完整的PyTorch编译器栈。

请注意，我已经通过了可选的参数，启用了两个调试功能：

  • trace.enabled:生成中间代码以检查由TorchInductor生成的代码
  • trace.graph_enabled:生成运算符融合后的优化计算图的可视化。

对于我们这个简单的例子，TorchInductor能够将我们函数中的所有中间操作融合成一个单一的自定义运算符，你可以在下面看到它是如何简化前向和后向计算图的。

image.png

你肯定想知道这个融合的运算符在代码中是什么样子的。融合运算符的代码是由TorchInductor自动生成的，根据目标设备--CPU或GPU，它采用C++或Triton。你不需要向TorchInductor明确指定目标设备，它可以从数据和模型设备类型中推断出来。

为了查看生成的代码，你必须使用trace.enabled=True来启用调试，这将创建一个名为torch_compile_debug的主管，其中包含调试信息。

前向和后向图形代码的完整路径是：

  • torch_compile_debug/run_<DATE_TIME_PID>/aot_torchinductor/model__XX_forward_XX/output_code.py
  • torch_compile_debug/run_<DATE_TIME_PID>/aot_torchinductor/model__XX_backward_XX/output_code.py

如果你设置了device='cuda'（假设你的电脑有一个GPU设备），那么前向文件夹中生成的代码在Open AI Triton中

image.png

如果你设置device = 'CPU'，那么生成的代码是用OpenMP pragmas的C++语言。

image.png

回顾一下

深度学习编译器是复杂的，其内部工作的复杂程度可以与瑞士手表相媲美。在这篇博文中，我希望我为您提供了一个关于这个主题的温和而易于理解的入门知识，以及这些技术如何为PyTorch 2.0提供动力。

image.png

  1. 我从一个简单的PyTorch函数开始
  2. 我展示了TorchDynamo是如何捕捉图形并在FX IR中表示它的。
  3. 我展示了AOTAutograd如何生成后向传递图，将PyTorch运算符降低为Aten运算符，并将其表示在FX图容器中。
  4. 我讨论了如何将Aten操作符进一步分解为Core Aten IR和Prims IR，以减少其他编译器可以支持的操作符数量，而不支持全部PyTorch API列表。
  5. 我展示了TorchInductor如何进行运算符融合，并为CPU和GPU目标生成优化代码。

如果你跟上了，你应该能够对以下问题做出高层次的回答：

  • 什么是深度学习编译器？
  • 当你调用torch.compile()时，PyTorch 2.0编译器会做什么？
  • 为什么我们需要一个提前的前向和后向传递图？
  • PyTorch中不同的中间表示法（IR）是什么？
  • ATen IR、Core ATen IR、Prims IR之间有什么区别？
  • 什么是运算符融合，为什么它很重要？
avatar
文章
阅读
粉丝