关于程序员:Pytorch2-如何通过算子融合和-CPUGPU-代码生成加速深度学习

4次阅读

共计 3417 个字符，预计需要花费 9 分钟才能阅读完成。

动动发财的小手，点个赞吧！

PyTorch 中用于图形捕捉、两头示意、运算符交融以及优化的 C++ 和 GPU 代码生成的深度学习编译器技术入门

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

编译器在深度学习中施展着十分重要的作用，能够进步训练和推理性能，进步能效，并针对多样化的 AI 加速器硬件。在这篇博文中，我将探讨为 PyTorch 2.0 提供反对的深度学习编译器技术。我将疏导您实现编译过程的不同阶段，并通过代码示例和可视化探讨各种底层技术。

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

与传统的软件编译器不同，深度学习编译器必须应用高度可并行化的代码，这些代码通常在专门的 AI 加速器硬件（GPU、TPU、AWS Trainium/Inferentia、Intel Habana Gaudi 等）上减速。为了进步性能，深度学习编译器必须利用硬件特定的性能，例如混合精度反对、性能优化的内核以及最小化主机 (CPU) 和 AI 加速器之间的通信。

在深度学习算法持续疾速倒退的同时，硬件 AI 加速器也在一直倒退，以满足深度学习算法的性能和效率需要。

在这篇博文中，我将重点关注软件方面的事件，尤其是更靠近硬件的软件子集——深度学习编译器。首先，让咱们先看看深度学习编译器中的不同函数。

PyTorch 2.0 包含新的编译器技术，以进步模型性能和运行时效率，并应用一个简略的 API 来针对不同的硬件后端：torch.compile()。尽管其余博客文章和文章曾经具体探讨了 PyTorch 2.0 的性能劣势，但在这里我将重点关注调用 PyTorch 2.0 编译器时产生的事件。如果你正在寻找量化的性能劣势，你能够找到来自 huggingface、timm 和 torchbench 的不同模型的性能仪表板。

在高层次上，PyTorch 2.0 深度学习编译器的默认选项执行以下要害工作：

图形捕捉：模型和函数的计算图形示意。PyTorch 技术：TorchDynamo、Torch FX、FX IR
主动微分：应用主动微分和升高到原始运算符的反向图形跟踪。PyTorch 技术：AOTAutograd、Aten IR
优化：前向和后向图级优化和运算符交融。PyTorch 技术：TorchInductor（默认）或其余编译器
代码生成：生成硬件特定的 C++/GPU 代码。PyTorch 技术：TorchInductor、OpenAI Triton（默认）其余编译器

通过这些步骤，编译器会转换您的代码并生成逐步“升高”的两头示意 (IR)。升高是编译器词典中的一个术语，指的是通过编译器的主动转换和重写将一组宽泛的操作（例如 PyTorch API 反对的）映射到一组狭隘的操作（例如硬件反对的）。PyTorch 2.0 编译器流程：

如果您不相熟编译器术语，请不要让所有这些吓到您。我也不是编译器工程师。持续浏览，事件会变得清晰，因为我将应用一个简略的示例和可视化来合成这个过程。

为了简略起见，我将定义一个非常简单的函数并通过 PyTorch 2.0 编译器过程运行它。您能够将此函数替换为深度神经网络模型或 nn.Module 子类，但与简单的数百万参数模型相比，此示例应该能够帮忙您更好地理解引擎盖下产生的事件。

该函数的 PyTorch 代码：

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

如果你在高中三角学课上留神过，你就会晓得咱们函数的值对于所有实值 x 总是为 1。这意味着它是导数，常数的导数，并且必须等于零。这将有助于验证函数及其派生函数的作用。

当初，是时候调用 torch.compile() 了。首先让咱们压服本人编译这个函数不会扭转它的输入。对于雷同的 1×1000 随机向量，咱们函数的输入与 1s 向量之间的均方误差对于编译函数和未编译函数（在肯定的误差容限下）都应该为零。

咱们所做的只是增加一行额定的代码 torch.compile() 来调用咱们的编译器。当初让咱们来看看每个阶段的幕后状况。

编译器的第一步是确定编译什么。输出 TorchDynamo。TorchDynamo 拦挡您的 Python 代码的执行并将其转换为 FX 两头示意 (IR)，并将其存储在称为 FX Graph 的非凡数据结构中。你问这看起来像什么？很快乐你问。上面，咱们将看一下咱们用来生成它的代码，但这里是转换和输入：

重要的是要留神，Torch FX 图只是 IR 的容器，并没有真正指定它应该蕴含哪些运算符。在下一节中，咱们将看到 FX 图形容器再次出现，并带有一组不同的 IR。如果比拟性能代码和 FX IR，两者之间的差异很小。事实上，它与您编写的 PyTorch 代码雷同，但以 FX 图形数据结构所需的格局进行布局。它们在执行时都将提供雷同的后果。

如果您调用 torch.compile() 时不带任何参数，它将应用运行整个编译器堆栈的默认设置，其中包含名为 TorchInductor 的默认硬件后端编译器。然而，如果咱们当初探讨 TorchInductor 就会跳到后面，所以让咱们临时搁置这个话题，等咱们筹备好后再回来探讨。首先咱们须要探讨图形捕捉，咱们能够通过拦挡来自 torch.compile() 的调用来实现。上面是咱们将如何做到这一点：torch.compile() 也容许你提供本人的编译器，但因为我不是编译器工程师，而且我对如何编写编译器无所不知，所以我会提供一个伪造的编译器函数来捕捉 TorchDynamo 生成的 FX 图形 IR。

上面是咱们的假编译器后端函数，称为 inspect_backend 到 torch.compile()，在该函数中我做了两件事：

打印 TorchDynamo 捕捉的 FX IR 代码
保留 FX 图形可视化

def inspect_backend(gm, sample_inputs):
    code = gm.print_readable()
    with open("forward.svg", "wb") as file:
        file.write(FxGraphDrawer(gm,'f').get_dot_graph().create_svg())
    return gm.forward

torch._dynamo.reset()
compiled_f = torch.compile(f, backend=inspect_backend)

x = torch.rand(1000, requires_grad=True).to(device)
out = compiled_f(x)

上述代码片段的输入是 FX IR 代码和显示函数 sin^2(x)+cos^2(x) 的图表

请留神，咱们的假编译器 inspect_backend 函数仅在咱们应用一些数据调用已编译函数时调用，即当咱们调用 compiled_model(x) 时。在下面的代码片段中，咱们只评估函数或在深度学习术语中，进行“前向流传”。在下一节中，咱们将利用 PyTorch 的主动微分引擎 torch.autograd 来计算导数和“向后传递”图。

TorchDynamo 为咱们提供了作为 FX 图的前向传递函数评估，然而向后传递呢？为了残缺起见，我将偏离咱们的次要主题，谈谈为什么咱们须要依据函数的权重来评估函数的梯度。如果您曾经相熟数学优化的工作原理，请跳过本节。

用于深度学习的优化编译器长于发现代码中的性能差距，并通过转换代码以缩小代码属性（例如指标后端的内存拜访、内核启动、数据布局优化）来解决这些问题。TorchInductor 是带有 torch.compile() 的默认优化编译器，它能够为应用 OpenAI Triton 的 GPU 和应用 OpenMP pragma 指令的 CPU 生成优化内核。

本文由 mdnice 多平台公布

正文完