nvFuser 是一款用于 NVIDIA GPU 的深度学习编译器，它能够自动进行即时（JIT）编译，生成快速且灵活的内核，从而可靠地加速用户的网络。通过在运行时生成快速的自定义“融合”（fusion）内核，它为在 Volta 及之后架构的 CUDA 加速器上运行的深度学习网络提供了显著的速度提升。nvFuser 专为满足 PyTorch 社区的独特需求而设计，支持多种网络架构，以及具有不同形状（shape）和步长（stride）的动态输入程序。在这篇博文中，我们将介绍 nvFuser 及其目前的使用方式，展示它在 HuggingFace 和 TIMM 模型上获得的显著性能改进，并展望 nvFuser 在 PyTorch 1.13 及更高版本中的发展。如果您想了解融合技术如何以及为何能提高深度学习网络的训练速度，请参阅我们之前在 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 的优势在于它能够将装饰器应用于用户的脚本，有效地隔离出应该发送给 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。模型运行时的批大小（batch size）和序列长度分别为 [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）且 dtype=float16 的情况下运行。
虽然这些速度提升非常显著，但必须理解 nvFuser 尚未实现对网络运行加速的全面自动化。例如，对于 HuggingFace Transformers，使用来自 NVIDIA Apex 仓库 的 AdamW 融合优化器非常重要，否则优化器将消耗大量的运行时间。使用融合的 AdamW 优化器来提高网络速度后，会暴露下一个主要的性能瓶颈——内存受限操作。这些操作由 nvFuser 优化，从而提供了另一次巨大的性能提升。在启用融合优化器和 nvFuser 后，这些网络的训练速度提高了 1.12 到 1.5 倍。HuggingFace Transformer 模型使用了 torch.amp 模块运行。（“amp”代表自动混合精度，详见 “PyTorch 混合精度训练中每个用户都应该知道的事” 博客文章。）使用 nvFuser 的选项已添加到 HuggingFace 的 Trainer 中。如果您安装了 TorchDynamo，可以通过将 torchdynamo = ‘nvfuser’ 传递给 Trainer 类来激活它，从而在 HuggingFace 中启用 nvFuser。nvFuser 对自然语言处理 (NLP) 模型中常见的归一化内核及相关融合有很好的支持，建议用户在 NLP 工作负载中尝试使用 nvFuser。
PyTorch Image Models (TIMM) 基准测试
nvFuser 也可以显著缩短 TIMM 网络的训练时间，与即时 PyTorch 相比提升超过 1.3 倍；结合 torch.amp 模块使用时，相比即时 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) 和 channels first 连续输入，圆形标记为 float32 输入，三角形为使用 float16 AMP 和 channels last 连续输入。缺少数据点是因为在追踪时遇到了错误。
在使用 float32 精度运行时，nvFuser 在 TIMM 网络上提供了 1.12 倍的几何平均值 (“geomean”) 加速；在使用 torch.amp 和“channels first”运行时，它提供了 1.14 倍的 geomean 加速。然而，nvFuser 目前并不能加速 torch.amp 和“channels last”训练（有 0.9 倍的 geomean 回归），因此我们建议在这些情况下不要使用它。我们目前正在积极改进“channels last”的性能，很快我们将会有两种额外的优化策略（针对 channels-last 归一化的网格持久化优化和快速转置），我们预计这些策略将在 PyTorch 1.13 及更高版本中提供与“channels first”相当的加速。nvFuser 的许多优化也可以在推理场景中提供帮助。然而，在 PyTorch 中进行小批次推理时，性能通常受限于 CPU 开销，这是 nvFuser 无法完全消除或修复的。因此，通常推理最重要的优化是在可能的情况下启用 CUDA Graphs。一旦启用了 CUDA Graphs，那么通过 nvFuser 启用融合也会是有益的。推理性能如图 2 和图 3 所示。推理仅在 float16 AMP 下运行，因为在完整的 float32 精度下运行推理工作负载并不常见。
图 2：在 TIMM 模型上，启用 CUDA Graphs 以及启用 CUDA Graphs 配合 nvFuser，与不使用 CUDA Graphs 和 nvFuser 的原生 PyTorch 性能相比的增益，批大小分别为 1 和 8。使用 CUDA Graphs 的 geomean 加速为 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 和 **channels last 输入**，批大小分别为 1 和 8。使用 CUDA Graphs 的 geomean 加速为 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。