Pytorch2 如何通过算子融合和 CPU/GPU 代码生成加速深度学习

简介: Pytorch2 如何通过算子融合和 CPU/GPU 代码生成加速深度学习

动动发财的小手,点个赞吧!

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

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

编译器在深度学习中发挥着非常重要的作用,可以提高训练和推理性能,提高能效,并针对多样化的 AI 加速器硬件。在这篇博文中,我将讨论为 PyTorch 2.0 提供支持的深度学习编译器技术。我将引导您完成编译过程的不同阶段,并通过代码示例和可视化讨论各种底层技术。

什么是深度学习编译器?

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

与传统的软件编译器不同,深度学习编译器必须使用高度可并行化的代码,这些代码通常在专门的 AI 加速器硬件(GPU、TPU、AWS Trainium/Inferentia、Intel Habana Gaudi 等)上加速。为了提高性能,深度学习编译器必须利用硬件特定的功能,例如混合精度支持、性能优化的内核以及最小化主机 (CPU) 和 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 编译器流程:

如果您不熟悉编译器术语,请不要让所有这些吓到您。我也不是编译器工程师。继续阅读,事情会变得清晰,因为我将使用一个简单的示例和可视化来分解这个过程。

遍历 torch.compile() 编译器过程

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

该函数的 PyTorch 代码:

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

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

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

我们所做的只是添加一行额外的代码 torch.compile() 来调用我们的编译器。现在让我们来看看每个阶段的幕后情况。

图形捕获:PyTorch 模型或函数的计算图形表示

编译器的第一步是确定编译什么。输入 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(),在该函数中我做了两件事:

  1. 打印 TorchDynamo 捕获的 FX IR 代码
  2. 保存 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 生成优化内核。

相关实践学习
部署Stable Diffusion玩转AI绘画(GPU云服务器)
本实验通过在ECS上从零开始部署Stable Diffusion来进行AI绘画创作,开启AIGC盲盒。
相关文章
|
5天前
|
机器学习/深度学习 人工智能 自然语言处理
探索人工智能的未来:机器学习与深度学习的融合之旅
【9月更文挑战第35天】在这篇文章中,我们将深入探讨人工智能的两大支柱——机器学习和深度学习。我们将通过代码示例和实际应用案例,揭示它们如何相互补充,共同推动AI技术的发展。无论你是初学者还是有经验的开发者,这篇文章都将为你提供宝贵的见解和启示。
22 0
|
13天前
|
机器学习/深度学习 自然语言处理 算法
深度学习与自然语言处理的融合
本文探讨了深度学习在自然语言处理中的应用,通过实例展示了深度学习技术如何提升自然语言处理的效果。我们将介绍深度学习和自然语言处理的基本概念,分析深度学习在自然语言处理中的具体应用,并通过案例展示其效果。最后,我们将总结深度学习对自然语言处理的影响,并展望未来的发展方向。
|
17天前
|
机器学习/深度学习 传感器 自然语言处理
深度学习与未来科技的融合
本文旨在探讨深度学习在现代科技中的应用及其对未来科技发展的潜在影响。通过分析深度学习的基本概念、关键技术以及在不同领域的应用案例,本文揭示了深度学习如何推动技术进步,并提出了未来发展的可能方向和面临的挑战。
|
18天前
|
机器学习/深度学习 人工智能 自然语言处理
探索AI的未来:深度学习与自然语言处理的融合
【9月更文挑战第22天】本文旨在探讨AI技术中深度学习与自然语言处理的结合,以及它们如何共同推动未来技术的发展。我们将通过实例和代码示例,深入理解这两种技术如何相互作用,以及它们如何影响我们的生活和工作。
39 4
|
1月前
|
机器学习/深度学习 测试技术 PyTorch
深度学习之测量GPU性能的方式
在深度学习中,测量GPU性能是一个多方面的任务,涉及运行时间、吞吐量、GPU利用率、内存使用情况、计算能力、端到端性能测试、显存带宽、框架自带性能工具和基准测试工具等多种方法。通过综合使用这些方法,可以全面评估和优化GPU的性能,提升深度学习任务的效率和效果。
73 2
|
14天前
|
机器学习/深度学习 人工智能 自然语言处理
探索人工智能的未来:深度学习与日常生活的融合
【9月更文挑战第26天】在这篇文章中,我们将一起探讨人工智能(AI)的一个子领域——深度学习,以及它如何影响我们的日常生活。通过简单易懂的语言和实际代码示例,我们将了解深度学习的基础知识,并看到它是如何被应用到诸如图像识别、自然语言处理等日常场景中的。无论你是技术新手还是有一定基础的爱好者,这篇文章都将为你提供有价值的见解和知识。
24 0
|
1月前
|
人工智能 自然语言处理 文字识别
MinerU-大语言语料处理神器,CPU/GPU均可跑,开源免费“敲”好用
在7月4日举行的WAIC 2024科学前沿主论坛上,书生·浦语2.5正式发布,面向大模型研发与应用的全链条工具体系同时迎来升级。
MinerU-大语言语料处理神器,CPU/GPU均可跑,开源免费“敲”好用
|
2月前
|
机器学习/深度学习 算法 TensorFlow
深入探索强化学习与深度学习的融合:使用TensorFlow框架实现深度Q网络算法及高效调试技巧
【8月更文挑战第31天】强化学习是机器学习的重要分支,尤其在深度学习的推动下,能够解决更为复杂的问题。深度Q网络(DQN)结合了深度学习与强化学习的优势,通过神经网络逼近动作价值函数,在多种任务中表现出色。本文探讨了使用TensorFlow实现DQN算法的方法及其调试技巧。DQN通过神经网络学习不同状态下采取动作的预期回报Q(s,a),处理高维状态空间。
43 1
|
2月前
|
机器学习/深度学习 并行计算 PyTorch
GPU 加速与 PyTorch:最大化硬件性能提升训练速度
【8月更文第29天】GPU(图形处理单元)因其并行计算能力而成为深度学习领域的重要组成部分。本文将介绍如何利用PyTorch来高效地利用GPU进行深度学习模型的训练,从而最大化训练速度。我们将讨论如何配置环境、选择合适的硬件、编写高效的代码以及利用高级特性来提高性能。
268 1
|
2月前
|
机器学习/深度学习 人工智能 调度
显著提升深度学习 GPU 利用率,阿里云拿下国际网络顶会优胜奖!
显著提升深度学习 GPU 利用率,阿里云拿下国际网络顶会优胜奖!
189 7