作者：Vinh Nguyen, Michael Carilli, Sukru Burc Eryilmaz, Vartika Singh, Michelle Lin, Natalia Gimelshein, Alban Desmaison, Edward Yang

今天，我们很高兴地宣布，一项新的高级 CUDA 功能——CUDA Graphs——已引入 PyTorch。现代深度学习（DL）框架拥有复杂的软件栈，将每个操作提交给 GPU 会产生显著的开销。当为了性能将深度学习工作负载强扩展到多个 GPU 时，每个 GPU 操作所需的时间会缩短到仅几微秒，在这种情况下，框架的高工作提交延迟往往会导致 GPU 利用率低下。随着 GPU 速度加快以及工作负载扩展到更多设备，工作负载受启动引起的停顿影响的可能性会增加。为了克服这些性能开销，NVIDIA 工程师与 PyTorch 开发者合作，在 PyTorch 中原生支持 CUDA 图执行。这一设计对于将 NVIDIA 的 MLPerf 工作负载（在 PyTorch 中实现）扩展到 4000 多个 GPU 并实现创纪录的性能至关重要。

PyTorch 中对 CUDA Graphs 的支持只是 NVIDIA 和 Facebook 工程师之间长期合作的又一个例子。torch.cuda.amp，例如，使用半精度进行训练，同时保持使用单精度实现的网络精度，并在可能的情况下自动利用 Tensor Core。仅需几行代码更改，AMP 即可提供比 FP32 高达 3 倍的性能。类似地，NVIDIA 的 Megatron-LM 是在 PyTorch 上使用多达 3072 个 GPU 进行训练的。在 PyTorch 中，扩展 GPU 训练性能的最高效方法之一是结合使用 torch.nn.parallel.DistributedDataParallel 和 NVIDIA Collective Communications Library (NCCL) 后端。

CUDA Graphs

CUDA Graphs 于 CUDA 10 首次亮相，它允许将一系列 CUDA 内核定义并封装为一个单元，即一个操作图，而不是一系列单独启动的操作。它提供了一种机制，可以通过单个 CPU 操作启动多个 GPU 操作，从而减少启动开销。

CUDA Graphs 的优势可以通过图 1 中的简单示例来展示。图顶部显示，CPU 逐个启动一系列短内核。CPU 启动开销在内核之间产生了显著的间隙。如果我们用 CUDA 图替换这一系列内核，最初我们需要花费一些额外的时间来构建图并在第一次一次性启动整个图，但随后的执行将非常快，因为内核之间的间隙非常小。当相同的操作序列重复多次时，例如在许多训练步骤中，这种差异会更加明显。在这种情况下，构建和启动图的初始成本将在整个训练迭代次数中摊销。关于该主题的更全面介绍，请参阅我们的博客 《CUDA Graphs 入门》和 GTC 演讲 《轻松使用 CUDA Graphs》。

图 1. 使用 CUDA Graphs 的优势

NCCL 对 CUDA Graphs 的支持

前面提到的减少启动开销的优势也适用于 NCCL 内核启动。NCCL 支持基于 GPU 的集合通信和 P2P 通信。借助 NCCL 对 CUDA Graphs 的支持，我们可以消除 NCCL 内核启动开销。

此外，由于各种 CPU 负载和操作系统因素，内核启动时间可能不可预测。这种时间偏差可能损害 NCCL 集合操作的性能。使用 CUDA Graphs，内核被组合在一起，从而使得分布式工作负载中各个进程（ranks）的性能保持一致。这在大规模集群中特别有用，即使一个慢节点也可能拉低整个集群的性能。

对于分布式多 GPU 工作负载，NCCL 用于集合通信。如果我们考虑训练利用数据并行性的神经网络，如果没有 NCCL 对 CUDA Graphs 的支持，前向/后向传播和 NCCL AllReduce 各自都需要单独启动。相比之下，借助 NCCL 对 CUDA Graphs 的支持，我们可以通过将前向/后向传播和 NCCL AllReduce 全部合并到一个图启动中来减少启动开销。

图 2. 对于一个典型的神经网络，所有用于 NCCL AllReduce 的内核启动都可以打包到一个图中，以减少启动开销时间。

PyTorch CUDA Graphs

从 PyTorch v1.10 开始，CUDA Graphs 功能作为一组 Beta 版 API 提供。

API 概述

PyTorch 支持使用 流捕获（stream capture）来构建 CUDA 图，这会将 CUDA 流置于捕获模式。发送给捕获流的 CUDA 工作实际上不会在 GPU 上运行。相反，这些工作被记录在一个图中。捕获后，可以启动该图来运行 GPU 工作，次数不限。每次回放都以相同的参数运行相同的内核。对于指针参数，这意味着使用相同的内存地址。通过在每次回放前用新数据（例如，来自新的批次）填充输入内存，您可以在新数据上重新运行相同的工作。

回放图以牺牲典型 Eager Execution 的动态灵活性为代价，换取大幅降低的 CPU 开销。图的参数和内核是固定的，因此图回放会跳过所有参数设置和内核分派层，包括 Python、C++ 和 CUDA 驱动程序开销。在底层，回放通过一次调用 cudaGraphLaunch 将整个图的工作提交给 GPU。回放中的内核在 GPU 上执行速度也会略快，但消除 CPU 开销是主要优势。

如果您的网络全部或部分是图安全的（通常这意味着静态形状和静态控制流，但请参阅其他约束），并且您怀疑其运行时至少在一定程度上受限于 CPU，那么您应该尝试使用 CUDA Graphs。

API 示例

PyTorch 通过一个原始的 torch.cuda.CUDAGraph 类和两个便捷的包装器 torch.cuda.graph 和 torch.cuda.make_graphed_callables 来暴露图功能。

torch.cuda.graph 是一个简单、通用的上下文管理器，用于捕获其上下文中的 CUDA 工作。在捕获之前，通过运行几次 Eager 迭代来预热要捕获的工作负载。预热必须在辅助流上进行。因为图在每次回放时都从相同的内存地址读取和写入，所以在捕获期间必须保留对持有输入和输出数据的张量的长期引用。要在新输入数据上运行图，请将新数据复制到捕获的输入张量中，回放图，然后从捕获的输出张量中读取新输出。

如果整个网络都是捕获安全的，则可以像以下示例一样捕获和回放整个网络。

N, D_in, H, D_out = 640, 4096, 2048, 1024
model = torch.nn.Sequential(torch.nn.Linear(D_in, H),
                            torch.nn.Dropout(p=0.2),
                            torch.nn.Linear(H, D_out),
                            torch.nn.Dropout(p=0.1)).cuda()
loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)

# Placeholders used for capture
static_input = torch.randn(N, D_in, device='cuda')
static_target = torch.randn(N, D_out, device='cuda')

# warmup
# Uses static_input and static_target here for convenience,
# but in a real setting, because the warmup includes optimizer.step()
# you must use a few batches of real data.
s = torch.cuda.Stream()
s.wait_stream(torch.cuda.current_stream())
with torch.cuda.stream(s):
    for i in range(3):
        optimizer.zero_grad(set_to_none=True)
        y_pred = model(static_input)
        loss = loss_fn(y_pred, static_target)
        loss.backward()
        optimizer.step()
torch.cuda.current_stream().wait_stream(s)

# capture
g = torch.cuda.CUDAGraph()
# Sets grads to None before capture, so backward() will create
# .grad attributes with allocations from the graph's private pool
optimizer.zero_grad(set_to_none=True)
with torch.cuda.graph(g):
    static_y_pred = model(static_input)
    static_loss = loss_fn(static_y_pred, static_target)
    static_loss.backward()
    optimizer.step()

real_inputs = [torch.rand_like(static_input) for _ in range(10)]
real_targets = [torch.rand_like(static_target) for _ in range(10)]

for data, target in zip(real_inputs, real_targets):
    # Fills the graph's input memory with new data to compute on
    static_input.copy_(data)
    static_target.copy_(target)
    # replay() includes forward, backward, and step.
    # You don't even need to call optimizer.zero_grad() between iterations
    # because the captured backward refills static .grad tensors in place.
    g.replay()
    # Params have been updated. static_y_pred, static_loss, and .grad
    # attributes hold values from computing on this iteration's data.

如果您的网络某些部分不安全捕获（例如，由于动态控制流、动态形状、CPU 同步或必要的 CPU 端逻辑），您可以 Eager 方式运行不安全的部分，并使用 torch.cuda.make_graphed_callables 仅对捕获安全的部分生成图。下面将对此进行演示。

make_graphed_callables 接受可调用对象（函数或 nn.Module）并返回图化版本。默认情况下，由 make_graphed_callables 返回的可调用对象是 Autograd 感知的，可以在训练循环中直接替换您传入的函数或 nn.Module。 make_graphed_callables 在内部创建 CUDAGraph 对象，运行预热迭代，并根据需要维护静态输入和输出。因此，（与使用 torch.cuda.graph 不同）您无需手动处理这些。

在下面的示例中，依赖于数据的动态控制流意味着网络无法端到端地捕获，但 make_graphed_callables() 允许我们无论如何都将图安全的部分捕获并作为图运行

N, D_in, H, D_out = 640, 4096, 2048, 1024

module1 = torch.nn.Linear(D_in, H).cuda()
module2 = torch.nn.Linear(H, D_out).cuda()
module3 = torch.nn.Linear(H, D_out).cuda()

loss_fn = torch.nn.MSELoss()
optimizer = torch.optim.SGD(chain(module1.parameters(),
                                  module2.parameters(),
                                  module3.parameters()),
                            lr=0.1)

# Sample inputs used for capture
# requires_grad state of sample inputs must match
# requires_grad state of real inputs each callable will see.
x = torch.randn(N, D_in, device='cuda')
h = torch.randn(N, H, device='cuda', requires_grad=True)

module1 = torch.cuda.make_graphed_callables(module1, (x,))
module2 = torch.cuda.make_graphed_callables(module2, (h,))
module3 = torch.cuda.make_graphed_callables(module3, (h,))

real_inputs = [torch.rand_like(x) for _ in range(10)]
real_targets = [torch.randn(N, D_out, device="cuda") for _ in range(10)]

for data, target in zip(real_inputs, real_targets):
    optimizer.zero_grad(set_to_none=True)

    tmp = module1(data)  # forward ops run as a graph

    if tmp.sum().item() > 0:
        tmp = module2(tmp)  # forward ops run as a graph
    else:
        tmp = module3(tmp)  # forward ops run as a graph

    loss = loss_fn(tmp, target)
    # module2's or module3's (whichever was chosen) backward ops,
    # as well as module1's backward ops, run as graphs
    loss.backward()
    optimizer.step()

示例用例

MLPerf v1.0 训练工作负载

PyTorch CUDA Graphs 功能在将 NVIDIA 的 MLPerf 训练 v1.0 工作负载（在 PyTorch 中实现）扩展到 4000 多个 GPU 并全面刷新纪录方面发挥了关键作用。下面我们展示了两个使用 CUDA Graphs 后观察到最显著提升的 MLPerf 工作负载，速度提升高达约 1.7 倍。

表 1. 使用 PyTorch CUDA Graph 提升 MLPerf 训练 v1.0 性能。

Mask R-CNN

深度学习框架使用 GPU 加速计算，但仍有大量代码在 CPU 核上运行。CPU 核处理张量形状等元数据，以准备启动 GPU 内核所需的参数。处理元数据是固定成本，而 GPU 完成的计算工作的成本与批次大小正相关。对于大批次大小，CPU 开销占总运行时间成本的比例可以忽略不计，但在小批次大小时，CPU 开销可能变得大于 GPU 运行时间。发生这种情况时，GPU 在内核调用之间会处于空闲状态。这个问题可以在图 3 的 NSight 时间线图中识别出来。下图显示了在图化之前，Mask R-CNN 的“骨干（backbone）”部分，每个 GPU 的批次大小为 1。绿色部分显示 CPU 负载，蓝色部分显示 GPU 负载。在此配置文件中，我们看到 CPU 负载达到 100% 满载，而 GPU 大部分时间处于空闲状态，GPU 内核之间有很多空白。

图 3：Mask R-CNN 的 NSight 时间线图

当张量形状是静态的时，CUDA Graphs 可以自动消除 CPU 开销。所有内核调用的完整图在第一步中被捕获，在后续步骤中，整个图通过单个操作启动，消除了所有 CPU 开销，如图 4 所示。

图 4：CUDA Graphs 优化

通过图化，我们看到 GPU 内核紧密排列，GPU 利用率保持很高。图化部分现在在 6 毫秒内运行，而不是 31 毫秒，速度提升了 5 倍。我们没有对整个模型进行图化，主要只是 ResNet 骨干网络，这导致了整体约 1.7 倍的速度提升。为了增加图的范围，我们在软件栈中进行了一些更改，以消除一些 CPU-GPU 同步点。在 MLPerf v1.0 中，这项工作包括更改 torch.randperm 函数的实现，使用 CUB 而非 Thrust，因为后者是一个同步的 C++ 模板库。这些改进在最新的 NGC 容器中可用。

BERT

类似地，通过对模型进行图捕获，我们消除了 CPU 开销以及伴随的同步开销。CUDA Graphs 的实现为我们的最大规模 BERT 配置带来了 1.12 倍的性能提升。为了最大化 CUDA Graphs 的优势，重要的是保持图的范围尽可能大。为了实现这一点，我们修改了模型脚本，在执行过程中移除了 CPU-GPU 同步，以便可以捕获整个模型。此外，我们还确保在图范围内执行期间的张量大小是静态的。例如，在 BERT 中，只有总 token 的特定子集对损失函数有贡献，这由预生成的掩码张量决定。从该掩码中提取有效 token 的索引，并使用这些索引收集对损失有贡献的 token，会导致一个具有动态形状的张量，即其形状在迭代中不是恒定的。为了确保张量大小是静态的，我们在损失计算中没有使用动态形状的张量，而是使用了静态形状的张量，其中使用一个掩码来指示哪些元素是有效的。因此，所有张量形状都是静态的。动态形状也需要 CPU-GPU 同步，因为它必须涉及框架在 CPU 端的内存管理。只使用静态形状时，无需进行 CPU-GPU 同步。这在图 5 中有所展示。

图 5. 通过使用固定大小的张量和布尔掩码（如文中描述），我们能够消除动态大小张量所需的 CPU 同步

NVIDIA 深度学习示例集中的 CUDA Graphs

单 GPU 用例也可以从使用 CUDA Graphs 中受益。对于使用小批次启动许多短内核的工作负载来说尤其如此。一个很好的例子是推荐系统的训练和推理。下面我们展示了 NVIDIA 深度学习示例集中深度学习推荐模型（DLRM）实现的初步基准测试结果。对此工作负载使用 CUDA Graphs 为训练和推理都带来了显著的加速。当使用非常小的批次大小时，效果尤其明显，因为此时 CPU 开销更为突出。

CUDA Graphs 正在积极集成到其他 PyTorch NGC 模型脚本和 NVIDIA Github 深度学习示例中。敬请关注更多关于如何使用它的示例。

图 6：DLRM 模型的 CUDA Graphs 优化。

行动号召：PyTorch v1.10 中的 CUDA Graphs

对于包含许多小型 GPU 内核因而被 CPU 启动开销拖慢的工作负载，CUDA Graphs 可以提供显著的优势。这已在我们的 MLPerf 工作中得到证明，尤其是在优化 PyTorch 模型方面。这些优化中的许多，包括 CUDA Graphs，已经或最终将被集成到我们的 PyTorch NGC 模型脚本集和 NVIDIA Github 深度学习示例中。目前，请查看我们的开源 MLPerf 训练 v1.0 实现，它可以作为一个很好的起点来了解 CUDA Graph 的实际应用。或者，在您自己的工作负载上尝试 PyTorch CUDA Graphs API。

我们感谢许多 NVIDIA 和 Facebook 工程师提供的讨论和建议：Karthik Mandakolathur US, Tomasz Grel, PLJoey Conway, Arslan Zulfiqar US

作者简介

Vinh Nguyen 深度学习工程师，NVIDIA

Vinh 是一名深度学习工程师和数据科学家，发表了 50 多篇科学文章，被引用次数超过 2500 次。在 NVIDIA，他的工作涵盖广泛的深度学习和人工智能应用，包括语音、语言和视觉处理以及推荐系统。

Michael Carilli 高级开发者技术工程师，NVIDIA

Michael 曾在空军研究实验室工作，为现代并行架构优化 CFD 代码。他拥有加州大学圣巴巴拉分校计算物理学博士学位。作为 PyTorch 团队成员，他专注于让 GPU 训练更快、数值更稳定，并对内部团队、外部客户和 PyTorch 社区用户来说更易于使用。

Sukru Burc Eryilmaz 开发架构高级架构师，NVIDIA

Sukru 获得斯坦福大学博士学位和 Bilkent 大学学士学位。他目前致力于提升神经网络训练的端到端性能，包括单节点规模和超级计算机规模。

Vartika Singh 深度学习框架和库技术合作负责人，NVIDIA

Vartika 曾领导团队在云计算和分布式计算、扩展和人工智能的交汇领域工作，影响了各大公司的设计和战略。她目前与 NVIDIA 内部和外部的主要框架、编译器组织和开发者合作，帮助设计在 NVIDIA 硬件上高效最优地运行。

Michelle Lin 产品实习生，NVIDIA

Michelle 目前正在加州大学伯克利分校攻读计算机科学和工商管理本科学位。她目前负责管理市场研究和为 Magnum IO 创建营销材料等项目的执行。

Natalia Gimelshein 应用研究科学家，Facebook

Natalia Gimelshein 曾在 NVIDIA 和 Facebook 工作，负责深度学习工作负载的 GPU 性能优化。她目前是 PyTorch 核心团队成员，与合作伙伴合作，以无缝支持新的软件和硬件功能。

Alban Desmaison 研究工程师，Facebook

Alban 学习了工程学并获得了机器学习和优化博士学位，在此期间，他在加入 Facebook 之前是 PyTorch 的开源贡献者。他的主要职责是维护一些核心库和功能（autograd、optim、nn），并致力于全面改进 PyTorch。

Edward Yang 研究工程师，Facebook

Edward 在 MIT 和斯坦福大学学习计算机科学，之后进入 Facebook 工作。他是 PyTorch 核心团队的一员，也是 PyTorch 的主要贡献者之一。