nvFuser 是一款适用于 NVIDIA GPU 的深度学习编译器，可自动即时编译快速且灵活的内核，以可靠地加速用户的网络。它通过在运行时生成快速的自定义“融合”内核，显著加速在 Volta 及更高版本的 CUDA 加速器上运行的深度学习网络。nvFuser 专为满足 PyTorch 社区的独特需求而设计，它支持各种网络架构和具有不同形状和步幅的动态输入的程序。在这篇博文中，我们将介绍 nvFuser 及其今天的用途，展示它在 HuggingFace 和 TIMM 模型上可以获得的显著性能提升，并展望 PyTorch 1.13 及更高版本中的 nvFuser。如果您想了解更多关于融合如何以及为何提高深度学习网络训练速度的信息，请参阅我们之前在 GTC 2022 和 GTC 2021 上关于 nvFuser 的演讲。nvFuser 依赖于 PyTorch 操作的图形表示来优化和加速。由于 PyTorch 具有 eager 执行模型，用户正在运行的 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 中仍在积极改进一些场景，本文将对此进行讨论。但是，nvFuser 今天确实支持许多 DL 性能关键操作，并且受支持的操作数量将在后续 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 Transformer 的训练。如图 1 所示，性能提升范围在 1.12 倍到 1.50 倍之间，涵盖了流行的 HuggingFace Transformer 网络子集。
图 1：来自 HuggingFace Transformer 仓库的 8 个训练场景的性能增益。深绿色的第一个性能提升是由于将优化器替换为 NVIDIA Apex fused 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 Image Models (TIMM) 基准测试
nvFuser 还可以显著减少 TIMM 网络的训练时间，与 eager PyTorch 相比，最高可提升 1.3 倍以上，与 eager PyTorch 结合 torch.amp 模块相比，最高可提升 1.44 倍。图 1 显示了 nvFuser 在不使用 torch.amp 以及将 torch.amp 与 NHWC（“channels last”）和 NCHW（“channels first”）格式一起使用时的加速效果。nvFuser 通过 FuncTorch 跟踪直接集成到 TIMM 中（不使用 TorchDynamo），并且可以通过在运行 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：启用 CUDA Graphs 以及 CUDA Graphs 与 nvFuser 相比，对于具有 float16 AMP、channels first 输入以及批次大小分别为 1 和 8 的 TIMM 模型，相对于没有 CUDA Graphs 和 nvFuser 的原生 PyTorch 性能的性能增益。CUDA Graphs 的 geomean 加速为 2.74 倍，CUDA Graphs + nvFuser 的 geomean 加速为 2.71 倍。nvFuser 提供了 0.68 倍的最大回归和 2.74 倍的最大性能增益（相对于没有 nvFuser 的 CUDA Graphs）。性能增益是相对于 PyTorch 在没有 CUDA Graphs 和 nvFuser 的情况下每次迭代的平均时间来衡量的。模型按 nvFuser 提供的额外性能大小排序。
图 3：启用 CUDA Graphs 以及 CUDA Graphs 与 nvFuser 相比，对于具有 AMP、channels last 输入以及批次大小分别为 1 和 8 的 TIMM 模型，相对于没有 CUDA Graphs 和 nvFuser 的原生 PyTorch 性能的性能增益。CUDA Graphs 的 geomean 加速为 2.29 倍，CUDA Graphs + nvFuser 的 geomean 加速为 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 上的演讲。
