今天，我们很高兴地宣布，一项全新的高级 CUDA 功能——CUDA Graphs 已被引入 PyTorch。现代深度学习框架具有复杂的软件栈，在向 GPU 提交每个操作时都会产生显著的开销。当深度学习负载被强扩展（strong-scaled）到多个 GPU 以提高性能时，每个 GPU 操作所需的时间缩短至仅几微秒。在这种情况下，框架的高工作提交延迟往往会导致 GPU 利用率低下。随着 GPU 变得更快，工作负载扩展到更多设备，工作负载受到这些因启动导致停顿的可能性也在增加。为了克服这些性能开销，NVIDIA 工程师与 PyTorch 开发人员合作，在 PyTorch 中原生启用了 CUDA graph 执行。这一设计对于将 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 集体通信库 (NCCL) 后端。

CUDA Graphs

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

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

图 1. 使用 CUDA graphs 的优势

NCCL 对 CUDA graphs 的支持

上述降低启动开销的优势同样适用于 NCCL 内核启动。NCCL 支持基于 GPU 的集体通信和 P2P 通信。通过 NCCL 对 CUDA graphs 的支持，我们可以消除 NCCL 内核启动开销。

此外，由于各种 CPU 负载和操作系统因素，内核启动时间可能不可预测。这种时间偏斜（time skews）可能对 NCCL 集体操作的性能产生不利影响。使用 CUDA graphs，内核被聚类在一起，从而使分布式工作负载中各 rank 之间的性能保持一致。这在大型集群中特别有用，因为即使是一个慢速节点也可能拖累整体集群性能。

对于分布式多 GPU 工作负载，NCCL 用于集体通信。如果我们观察利用数据并行性的神经网络训练，如果没有 NCCL 对 CUDA graphs 的支持，我们需要为前向/反向传播和 NCCL AllReduce 分别进行启动。相比之下，借助 NCCL 对 CUDA graphs 的支持，我们可以将前向/反向传播和 NCCL AllReduce 合并到一次图启动中，从而降低启动开销。

图 2. 观察一个典型的神经网络，NCCL AllReduce 的所有内核启动都可以捆绑到一个图中，以减少启动开销时间。

PyTorch CUDA Graphs

从 PyTorch v1.10 开始，CUDA graphs 功能以一系列测试版 API 的形式提供。

API 概览

PyTorch 支持使用 流捕获 (stream capture) 构建 CUDA graphs，这会将 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 工作。捕获前，通过运行几次饥饿迭代来预热（warm up）要捕获的工作负载。预热必须在辅助流上进行。由于图在每次重播中都从相同的内存地址读取和写入，因此在捕获期间，您必须保持对保存输入和输出数据的张量的长期引用。要在新输入数据上运行图，请将新数据复制到捕获的输入张量中，重播图，然后从捕获的输出张量中读取新输出。

如果整个网络是捕获安全的，则可以像以下示例那样捕获并重播整个网络。

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 端逻辑），您可以饥饿地运行不安全的部分，并使用 torch.cuda.make_graphed_callables 仅对捕获安全的部分进行制图。下面演示了这一点。

make_graphed_callables 接受可调用对象（函数或 nn.Module），并返回已制图的版本。默认情况下，由 make_graphed_callables 返回的可调用对象是感知自动求导（autograd-aware）的，可以在训练循环中直接替换您传递的函数或 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 倍的加速。

	GPU 数量	CUDA-graphs 带来的加速
Mask R-CNN	272	1.70×
BERT	4096	1.12×

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

Mask R-CNN

深度学习框架使用 GPU 加速计算，但仍有大量代码在 CPU 核心上运行。CPU 核心处理张量形状等元数据，以准备启动 GPU 内核所需的参数。处理元数据是固定成本，而 GPU 所做的计算工作成本与批次大小正相关。对于大批次，CPU 开销在总运行时间成本中所占比例可忽略不计；但在小批次下，CPU 开销可能超过 GPU 运行时间。当发生这种情况时，GPU 在内核调用之间会处于空闲状态。这个问题可以在图 3 的 NSight 时间轴图中识别出来。下方的图显示了制图前每 GPU 批次大小为 1 的 Mask R-CNN 的“骨干”部分。绿色部分显示 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 中，只有特定子集的标记（tokens）参与损失函数，这是由预生成的掩码张量决定的。从该掩码中提取有效标记的索引，并使用这些索引来收集贡献损失的标记，会产生一个动态形状的张量，即形状在迭代中不是恒定的。为了确保张量大小是静态的，我们没有在损失计算中使用动态形状张量，而是使用静态形状张量，其中使用掩码来指示哪些元素是有效的。结果是，所有张量形状都是静态的。动态形状也需要 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，他的工作涵盖了广泛的深度学习和 AI 应用，包括语音、语言和视觉处理以及推荐系统。

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

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

Sukru Burc Eryilmaz 高级架构师，NVIDIA

Sukru 获得了斯坦福大学博士学位和比尔肯特大学学士学位。他目前致力于提高神经网络训练的端到端性能，涵盖单节点规模和超级计算机规模。

Vartika Singh 深度学习框架和库技术合作伙伴主管，NVIDIA

Vartika 领导了跨云与分布式计算、扩展和 AI 领域的团队，影响了主要企业的设计和战略。她目前与 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 在麻省理工学院和斯坦福大学学习计算机科学，之后加入 Facebook。他是 PyTorch 核心团队的一员，也是 PyTorch 的主要贡献者之一。