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 今天确实支持许多对 DL 性能至关重要的操作，并且支持的操作数量将在后续 PyTorch 版本中增加。nvFuser 能够为其支持的操作生成高度专业化和优化的 GPU 函数。这意味着 nvFuser 能够为 TorchDynamo 和 FuncTorch 等新的 PyTorch 系统提供支持，将 PyTorch 以其灵活性著称的特性与无与伦比的性能结合起来。
nvFuser 性能
在深入了解如何使用 nvFuser 之前，本节将展示 nvFuser 为来自 HuggingFace Transformers 和 PyTorch 图像模型 (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)，dtype=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（“通道优先”）和 NCHW（“通道最后”）格式时的加速。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 也可以通过 nvFuser 启用融合。推理性能如图 2 和图 3 所示。推理只在 float16 AMP 下运行，因为在 full float32 精度下运行推理工作负载并不常见。
图 2：启用 CUDA Graphs 和 CUDA Graphs + nvFuser 相对于未启用 CUDA Graphs 和 nvFuser 的原生 PyTorch 在 TIMM 模型上使用 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：启用 CUDA Graphs 和 CUDA Graphs + nvFuser 相对于未启用 CUDA Graphs 和 nvFuser 的原生 PyTorch 在 TIMM 模型上使用 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。