nvFuser 是一款用于 NVIDIA GPU 的深度学习编译器，可自动进行即时编译，生成快速灵活的内核，从而可靠地加速用户的网络。通过在运行时生成快速的定制“融合”内核，它显著加速了在 Volta 及更新的 CUDA 加速器上运行的深度学习网络。nvFuser 专为满足 PyTorch 社区的独特需求而设计，支持具有不同形状和步幅的动态输入的多种网络架构和程序。在这篇博客文章中，我们将介绍 nvFuser 及其当前的使用方式，展示它在 HuggingFace 和 TIMM 模型上获得的显著性能提升，并展望 PyTorch 1.13 及更高版本中的 nvFuser。如果您想进一步了解融合如何以及为何能提升深度学习网络的训练速度，请观看我们之前在 GTC 2022 和 GTC 2021 上关于 nvFuser 的演讲。nvFuser 依赖于 PyTorch 操作的图表示进行优化和加速。由于 PyTorch 采用 eager execution 模型，用户运行的 PyTorch 操作不能直接作为一个完整的程序被像 nvFuser 这样的系统优化。因此，用户必须利用构建在 nvFuser 之上的系统，这些系统能够捕获用户程序并将其转换为 nvFuser 可优化的形式。这些更高级的系统随后将捕获的操作传递给 nvFuser，以便 nvFuser 可以优化用户脚本在 NVIDIA GPU 上的执行。目前有三种系统负责捕获、转换并将用户程序传递给 nvFuser 进行优化
- TorchScript jit.script
- 该系统直接解析标注的 Python 脚本片段，将其转换为自己的内部表示，从而理解用户正在做什么。然后，该系统对图应用自己的自动微分版本，并将后续前向和后向图的部分传递给 nvFuser 进行优化。
- FuncTorch
- 该系统不直接查看用户 Python 脚本，而是插入一种机制来捕获正在运行的 PyTorch 操作。我们将这种捕获系统称为“程序轨迹获取”（trace program acquisition），因为它跟踪的是已执行的操作。FuncTorch 不执行自己的自动微分，它直接跟踪 PyTorch 的 autograd 来获取后向图。
- TorchDynamo
- TorchDynamo 是构建在 FuncTorch 之上的另一种程序获取机制。TorchDynamo 解析用户脚本生成的 Python 字节码，以便选择要使用 FuncTorch 跟踪的部分。TorchDynamo 的优势在于它能够对用户脚本应用装饰器（decorators），有效地隔离应发送给 FuncTorch 的内容，从而使 FuncTorch 更容易成功跟踪复杂的 Python 脚本。
用户可以直接与这些系统交互，而 nvFuser 会自动无缝地优化用户代码中性能关键的区域。这些系统会自动将解析后的用户程序发送给 nvFuser，以便 nvFuser 可以：
- 分析在 GPU 上运行的操作
- 为这些操作规划并行化和优化策略
- 在生成的 GPU 代码中应用这些策略
- 运行时编译生成的优化 GPU 函数
- 在后续迭代中执行这些 CUDA 内核
需要注意的是，nvFuser 尚未支持所有 PyTorch 操作，并且本文讨论的一些场景仍在积极改进中。然而，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 优化器加速网络后，会暴露下一个主要的性能瓶颈——内存密集型操作（memory bound operations）。这些操作由 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 Image Models (TIMM) 基准测试
nvFuser 也可以显著缩短 TIMM 网络的训练时间，相较于 eager PyTorch 可提升 1.3 倍以上，与 torch.amp 模块结合使用时，相较于 eager PyTorch 可提升高达 1.44 倍。图 1 显示了 nvFuser 在不使用 torch.amp 以及与 torch.amp 结合使用 NHWC（“通道在后”）和 NCHW（“通道在前”）格式时的加速效果。nvFuser 通过 FuncTorch 直接跟踪（不使用 TorchDynamo）的方式集成到 TIMM 中，在运行 TIMM 基准测试或训练脚本时，通过添加 –aot-autograd 命令行参数即可使用。
图 1：Y 轴表示 nvFuser 相较于不使用 nvFuser 时的性能提升。值为 1.0 表示性能没有变化，2.0 表示 nvFuser 快一倍，0.5 表示 nvFuser 运行时间增加一倍。方形标记表示使用 float16 自动混合精度 (AMP) 和通道在前（contiguous）的输入，圆形标记表示 float32 输入，三角形标记表示使用 float16 AMP 和通道在后（contiguous）的输入。缺失的数据点是由于跟踪时遇到错误。
在使用 float32 精度运行时，nvFuser 在 TIMM 网络上提供了 1.12 倍的几何平均 (geomean) 加速，而当与 torch.amp 和“通道在前”结合运行时，则提供了 1.14 倍的几何平均加速。然而，nvFuser 目前并不能加速 torch.amp 和“通道在后”的训练（几何平均回归到 0.9 倍），因此我们不建议在这些情况下使用。我们目前正在积极改进“通道在后”的性能，很快我们将推出两种额外的优化策略（针对通道在后归一化的 grid persistent 优化和快速转置），我们期望这些策略在 PyTorch 1.13 及更高版本中能够提供与“通道在前”相当的加速效果。nvFuser 的许多优化也可以在推理场景中提供帮助。然而，在 PyTorch 中，当使用小批次大小进行推理时，性能通常受限于 CPU 开销，这是 nvFuser 无法完全消除或修复的。因此，对于推理而言，最重要的优化通常是尽可能启用 CUDA Graphs。一旦启用 CUDA Graphs，通过 nvFuser 启用融合也可能带来益处。推理性能如 图 2 和 图 3 所示。推理仅在 float16 AMP 下运行，因为在完整的 float32 精度下运行推理工作负载并不常见。
图 2：与未使用 CUDA Graphs 和 nvFuser 的原生 PyTorch 性能相比，在 TIMM 模型上启用 CUDA Graphs 以及启用 CUDA Graphs + nvFuser 的性能提升，这些模型使用 float16 AMP、通道在前输入，批次大小分别为 1 和 8。启用 CUDA Graphs 的几何平均加速为 2.74 倍，启用 CUDA Graphs + nvFuser 的几何平均加速为 2.71 倍。nvFuser 的最大性能下降（regression）为 0.68 倍，最大性能提升为 2.74 倍（相对于未启用 nvFuser 的 CUDA Graphs）。性能提升是相对于 PyTorch 在不使用 CUDA Graphs 和 nvFuser 的情况下每次迭代的平均时间来衡量的。模型按 nvFuser 提供的额外性能提升量排序。
图 3：与未使用 CUDA Graphs 和 nvFuser 的原生 PyTorch 性能相比，在 TIMM 模型上启用 CUDA Graphs 以及启用 CUDA Graphs + nvFuser 的性能提升，这些模型使用 AMP、通道在后输入，批次大小分别为 1 和 8。启用 CUDA Graphs 的几何平均加速为 2.29 倍，启用 CUDA Graphs + nvFuser 的几何平均加速为 2.95 倍。nvFuser 的最大性能下降（regression）为 0.86 倍，最大性能提升为 3.82 倍（相对于未启用 nvFuser 的 CUDA Graphs）。性能提升是相对于 PyTorch 在不使用 CUDA Graphs 和 nvFuser 的情况下每次迭代的平均时间来衡量的。模型按 nvFuser 提供的额外性能提升量排序。
目前，nvFuser 的性能尚未针对推理工作负载进行调优，因此其性能优势在所有情况下并非一致。然而，仍有许多模型在推理过程中能从 nvFuser 中显著受益，我们鼓励用户在推理工作负载中尝试使用 nvFuser，看看今天是否能从中受益。nvFuser 在推理工作负载中的性能未来会得到提升，如果您对在推理工作负载中使用 nvFuser 感兴趣，请通过 PyTorch 论坛联系我们。
入门 - 使用 nvFuser 加速您的脚本
我们创建了一个教程，演示如何利用 nvFuser 来加速标准 transformer block 的一部分，以及如何使用 nvFuser 定义快速且新颖的操作。正如本文所述，nvFuser 目前仍有一些不足之处，我们正在努力改进。但我们也展示了在 HuggingFace 和 TIMM 的多个网络上训练速度的显著提升，我们预计您的网络中目前已有 nvFuser 可以提供帮助的机会，未来还会有更多。如果您想了解更多关于 nvFuser 的信息，我们建议观看我们在 NVIDIA GTC 大会 GTC 2022 和 GTC 2021 上的演讲。
