nvFuser 是一个用于 NVIDIA GPU 的深度学习编译器，它能够自动即时编译快速灵活的内核，以可靠地加速用户的网络。通过在运行时生成快速自定义“融合”内核，它为在 Volta 及更高版本 CUDA 加速器上运行的深度学习网络提供了显著的加速。nvFuser 专门设计用于满足 PyTorch 社区的独特需求，并支持具有不同形状和步幅的动态输入的各种网络架构和程序。在本篇博文中，我们将介绍 nvFuser 及其当前的使用方式，展示它在 HuggingFace 和 TIMM 模型上可以获得的显著性能提升，并展望 PyTorch 1.13 及更高版本中的 nvFuser。如果您想了解融合如何以及为何能提高深度学习网络的训练速度，请参阅我们之前在 GTC 2022 和 GTC 2021 上关于 nvFuser 的演讲。nvFuser 依赖于 PyTorch 操作的图表示来进行优化和加速。由于 PyTorch 采用即时执行模型，用户正在运行的 PyTorch 操作无法直接作为可由 nvFuser 等系统优化的完整程序来访问。因此，用户必须使用基于 nvFuser 构建的系统，这些系统能够捕获用户程序并将其转换为可由 nvFuser 优化的形式。然后，这些更高级的系统将这些捕获的操作传递给 nvFuser，以便 nvFuser 可以优化用户脚本在 NVIDIA GPU 上的执行。有三种系统可以捕获、转换用户程序并将其传递给 nvFuser 进行优化
- TorchScript jit.script
- 该系统直接解析带注释的 Python 脚本的各个部分，将其转换为自己的表示形式，以理解用户正在执行的操作。然后，该系统将自己的自动微分版本应用于图，并将随后的前向和后向图的各个部分传递给 nvFuser 进行优化。
- FuncTorch
- 该系统不直接查看用户 Python 脚本，而是插入一种机制来捕获正在运行的 PyTorch 操作。我们将这种捕获系统称为“跟踪程序获取”，因为我们正在跟踪已执行的操作。FuncTorch 不执行自己的自动微分——它只是直接跟踪 PyTorch 的 autograd 以获取反向图。
- TorchDynamo
- TorchDynamo 是另一种基于 FuncTorch 构建的程序获取机制。TorchDynamo 解析用户脚本生成的 Python 字节码，以选择要使用 FuncTorch 跟踪的部分。TorchDynamo 的好处是它能够将装饰器应用于用户的脚本，有效地隔离应发送给 FuncTorch 的内容，使 FuncTorch 更容易成功跟踪复杂的 Python 脚本。
这些系统可供用户直接交互，同时 nvFuser 自动无缝地优化用户代码中性能关键区域。这些系统自动将解析后的用户程序发送给 nvFuser，以便 nvFuser 可以
- 分析在 GPU 上运行的操作
- 规划这些操作的并行化和优化策略
- 在生成的 GPU 代码中应用这些策略
- 运行时编译生成的优化 GPU 函数
- 在后续迭代中执行这些 CUDA 内核
需要注意的是，nvFuser 尚未支持所有 PyTorch 操作，并且 nvFuser 中仍有一些场景正在积极改进中，本文将对此进行讨论。然而，nvFuser 目前确实支持许多对深度学习性能至关重要的操作，并且在随后的 PyTorch 版本中，支持的操作数量将不断增加。nvFuser 能够为其支持的操作生成高度专业化和优化的 GPU 函数。这意味着 nvFuser 能够为 TorchDynamo 和 FuncTorch 等新 PyTorch 系统提供支持，将 PyTorch 以其灵活性著称的特性与无与伦比的性能相结合。
nvFuser 性能
在深入了解如何使用 nvFuser 之前，本节将展示 nvFuser 为来自 HuggingFace Transformers 和 PyTorch Image Models (TIMM) 仓库的各种模型提供的训练速度改进，我们将讨论 nvFuser 性能中当前正在开发中的不足。本节中的所有性能数据均使用 NVIDIA A100 40GB GPU 测量，并单独使用 FuncTorch 或 FuncTorch 结合 TorchDynamo。
HuggingFace Transformer 基准测试
当与另一个重要的优化相结合时（稍后会详细介绍），nvFuser 可以显著加速 HuggingFace Transformers 的训练。如图 1 所示，在部分流行的 HuggingFace Transformer 网络中，性能提升范围在 1.12 倍到 1.50 倍之间。
图 1：HuggingFace Transformer 仓库中 8 种训练场景的性能提升。深绿色部分的第一波性能提升是由于将优化器替换为 NVIDIA Apex 融合的 AdamW 优化器。浅绿色部分是由于添加了 nvFuser。模型分别以 [64, 128]、[8, 512]、[2, 1024]、[64, 128]、[8, 512]、[8, src_seql=512, tgt_seql=128]、[8, src_seql=1024, tgt_seql=128] 和 [8, 512] 的批大小和序列长度运行。所有网络均启用自动混合精度 (AMP)，数据类型为 float16。
虽然这些加速效果显著，但重要的是要理解 nvFuser（目前）尚未自动化网络快速运行的一切。例如，对于 HuggingFace Transformers，使用来自 NVIDIA 的 Apex 仓库 的 AdamW 融合优化器至关重要，否则优化器会占用大量的运行时间。使用融合的 AdamW 优化器使网络更快地运行，从而暴露出下一个主要的性能瓶颈——内存受限操作。这些操作通过 nvFuser 进行了优化，提供了另一个巨大的性能提升。在启用融合优化器和 nvFuser 的情况下，这些网络的训练速度提高了 1.12 倍到 1.5 倍。HuggingFace Transformer 模型使用 torch.amp 模块 运行。（“amp”代表自动混合精度，详情请参阅 “每个用户都应该了解的 PyTorch 混合精度训练” 博文。）HuggingFace 的 Trainer 中添加了使用 nvFuser 的选项。如果您已经 安装了 TorchDynamo，您可以通过将 torchdynamo = ‘nvfuser’ 传递给 Trainer 类来激活它，从而在 HuggingFace 中启用 nvFuser。nvFuser 对自然语言处理 (NLP) 模型中常见的归一化内核和相关融合有很好的支持，建议用户在他们的 NLP 工作负载中尝试使用 nvFuser。
PyTorch 图像模型 (TIMM) 基准测试
nvFuser 还可以显著缩短 TIMM 网络的训练时间，相较于 eager PyTorch 可提升高达 1.3 倍，当与 torch.amp 模块结合使用时，相较于 eager PyTorch 可提升高达 1.44 倍。图 1 展示了 nvFuser 在不使用 torch.amp 的情况下的加速，以及在使用 torch.amp 且采用 NHWC（“channels last”）和 NCHW（“channels first”）格式时的加速。nvFuser 通过 FuncTorch 直接跟踪（不使用 TorchDynamo）集成到 TIMM 中，可以通过在运行 TIMM 基准测试或训练脚本时添加 --aot-autograd 命令行参数 来使用。
图 1：Y 轴是 nvFuser 相对于不使用 nvFuser 所提供的性能增益。值为 1.0 表示性能没有变化，2.0 表示 nvFuser 快两倍，0.5 表示 nvFuser 运行时间是原来的两倍。方块标记表示使用 float16 自动混合精度 (AMP) 和通道优先连续输入，圆形标记表示 float32 输入，三角形标记表示使用 float16 AMP 和通道最后连续输入。缺失的数据点是由于跟踪时遇到错误。
当以 float32 精度运行时，nvFuser 在 TIMM 网络上提供了 1.12 倍的几何平均 (geomean) 加速，当与 torch.amp 和“通道优先”一起运行时，它提供了 1.14 倍的几何平均加速。然而，nvFuser 目前无法加速 torch.amp 和“通道最后”训练（0.9 倍的几何平均回归），因此我们建议在这些情况下不要使用它。我们正在积极改进“通道最后”的性能，很快我们将有两种额外的优化策略（针对通道最后归一化的网格持久化优化和快速转置），我们预计这些策略将在 PyTorch 1.13 及更高版本中提供与“通道优先”相当的加速。nvFuser 的许多优化也可以在推理情况下提供帮助。然而，在 PyTorch 中，当以小批量大小运行推理时，性能通常受限于 CPU 开销，而 nvFuser 无法完全消除或修复这些开销。因此，通常推理最重要的优化是在可能的情况下启用 CUDA 图。一旦启用了 CUDA 图， then then it can also be beneficial to also enable fusion through nvFuser。推理的性能如图 2 和图 3 所示。推理只运行 float16 AMP，因为在全 float32 精度下运行推理工作负载并不常见。
图 2：在 TIMM 模型上，启用 CUDA Graphs 以及 CUDA Graphs 与 nvFuser 结合的性能增益与未启用 CUDA Graphs 和 nvFuser 的原生 PyTorch 性能进行比较，这些模型均使用 float16 AMP、通道优先输入，批大小分别为 1 和 8。CUDA Graphs 提供了 2.74 倍的几何平均加速，而 CUDA Graphs + nvFuser 分别提供了 2.71 倍的几何平均加速。nvFuser 的最大回归为 0.68 倍，最大性能增益为 2.74 倍（相对于未启用 nvFuser 的 CUDA Graphs）。性能增益是根据 PyTorch 在不使用 CUDA Graphs 和 nvFuser 的情况下每次迭代的平均时间来衡量的。模型按 nvFuser 提供的额外性能排序。
图 3：在 TIMM 模型上，启用 CUDA Graphs 以及 CUDA Graphs 与 nvFuser 结合的性能增益与未启用 CUDA Graphs 和 nvFuser 的原生 PyTorch 性能进行比较，这些模型均使用 AMP、通道最后输入，批大小分别为 1 和 8。CUDA Graphs 提供了 2.29 倍的几何平均加速，而 CUDA Graphs + nvFuser 分别提供了 2.95 倍的几何平均加速。nvFuser 的最大回归为 0.86 倍，最大性能增益为 3.82 倍（相对于未启用 nvFuser 的 CUDA Graphs）。性能增益是根据 PyTorch 在不使用 CUDA Graphs 和 nvFuser 的情况下每次迭代的平均时间来衡量的。模型按 nvFuser 提供的额外性能排序。
到目前为止，nvFuser 的性能尚未针对推理工作负载进行调优，因此其性能优势并非在所有情况下都一致。然而，仍然有许多模型在推理过程中受益于 nvFuser，我们鼓励用户在推理工作负载中尝试 nvFuser，看看今天是否能从中受益。nvFuser 在推理工作负载中的性能将在未来得到提升，如果您对 nvFuser 在推理工作负载中的应用感兴趣，请通过 PyTorch 论坛与我们联系。
开始使用 – 使用 nvFuser 加速您的脚本
我们创建了一个教程，演示如何利用 nvFuser 加速标准 Transformer 块的一部分，以及如何使用 nvFuser 定义快速且新颖的操作。如本博客文章所述，nvFuser 仍然存在一些我们正在努力改进的粗糙之处。然而，我们也展示了在 HuggingFace 和 TIMM 中多个网络训练速度上的巨大改进，我们预计您的网络中存在 nvFuser 可以提供帮助的机会，并且未来将有更多机会。如果您想了解更多关于 nvFuser 的信息，我们建议观看我们在 NVIDIA GTC 会议上的演示：GTC 2022 和 GTC 2021。