由 Vinh Nguyen, Michael Carilli, Sukru Burc Eryilmaz, Vartika Singh, Michelle Lin, Natalia Gimelshein, Alban Desmaison, Edward Yang 撰写

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

PyTorch 对 CUDA 图的支持只是 NVIDIA 和 Facebook 工程师长期合作的一个例子。例如，torch.cuda.amp 可以在保持与单精度相同的网络精度的情况下以半精度进行训练，并尽可能地自动利用张量核心。AMP 只需更改几行代码即可将性能提高 3 倍以上。同样，NVIDIA 的 Megatron-LM 也是使用 PyTorch 在最多 3072 个 GPU 上训练的。在 PyTorch 中，最有效的扩展 GPU 训练的方法之一是使用 torch.nn.parallel.DistributedDataParallel 与 NVIDIA 集体通信库（NCCL）后端相结合。

CUDA 图

CUDA 图首次出现在 CUDA 10 中，允许将一系列 CUDA 内核定义并封装为一个单一单元，即操作图，而不是一系列单独启动的操作。它提供了一种通过单个 CPU 操作启动多个 GPU 操作的方法，从而减少了启动开销。

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

图 1. 使用 CUDA 图的好处

支持 CUDA 图的 NCCL

前文提到的减少启动开销的好处也适用于 NCCL 内核启动。NCCL 支持基于 GPU 的集体和 P2P 通信。有了对 CUDA 图的 NCCL 支持，我们可以消除 NCCL 内核启动开销。

此外，由于 CPU 负载和操作系统等因素，内核启动时间可能不可预测。这种时间偏差可能对 NCCL 集体操作的性能有害。使用 CUDA 图，内核被聚集在一起，以确保分布式工作负载中各个 rank 的性能一致。这在大型集群中特别有用，因为即使单个慢节点也可能降低整体集群级别的性能。

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

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

PyTorch CUDA 图

从 PyTorch v1.10 版本开始，CUDA 图功能作为一组测试版 API 提供。

API 概述

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

重放图以降低 CPU 开销为代价，牺牲了典型即时执行（eager execution）的动态灵活性。图的参数和内核是固定的，因此图重放跳过了所有参数设置和内核调度的层次，包括 Python、C++和 CUDA 驱动程序的开销。在底层，重放通过一次调用 cudaGraphLaunch 将整个图的工作提交给 GPU。重放中的内核在 GPU 上执行也略快，但省略 CPU 开销是主要好处。

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

API 示例

PyTorch 通过一个原始类 torch.cuda.CUDAGraph 和两个便利包装器 torch.cuda.graph 和 torch.cuda.make_graphed_callables 公开图。

torch.cuda.graph 是一个简单、通用的上下文管理器，可以捕获其上下文中的 CUDA 工作。在捕获之前，通过运行几个急切迭代来预热要捕获的工作负载。预热必须在侧流上发生。因为图在每次重放时都从相同的内存地址读取并写入，所以您必须在捕获期间维护对包含输入和输出数据的张量的长期引用。要在新的输入数据上运行图，请将新数据复制到捕获的输入张量（s），重放图，然后从捕获的输出张量（s）读取新的输出。

如果整个网络是捕获安全的，就可以像以下示例那样捕获和重放整个网络。

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 返回的可调用对象是自动微分感知的，可以用作训练循环中函数或 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 图功能对于将 NVIDIA 的 MLPerf 训练 v1.0 工作负载（在 PyTorch 中实现）扩展到超过 4000 个 GPU，在各个领域都创造了新的记录至关重要。以下我们展示了两个使用 CUDA 图观察到最大增益的 MLPerf 工作负载，速度提升高达 ~1.7 倍。

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

Mask R-CNN

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

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

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

图 4：CUDA 图优化

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

BERT

同样，通过图捕获模型，我们消除了 CPU 开销和伴随的同步开销。CUDA 图实现使我们的最大规模 BERT 配置性能提升了 1.12 倍。为了最大限度地发挥 CUDA 图的优势，重要的是保持图的范围尽可能大。为了实现这一点，我们修改了模型脚本，在执行过程中移除了 CPU-GPU 同步，以便整个模型可以被图捕获。此外，我们还确保了在图的作用域内，执行过程中的张量大小是静态的。例如，在 BERT 中，只有总令牌的一个特定子集对损失函数有贡献，由预先生成的掩码张量确定。从该掩码中提取有效令牌的索引，并使用这些索引收集对损失有贡献的令牌，结果得到一个具有动态形状的张量，即形状在迭代中不是恒定的。为了确保张量大小是静态的，我们不是在损失计算中使用动态形状的张量，而是使用了静态形状的张量，其中使用掩码来指示哪些元素是有效的。因此，所有张量的形状都是静态的。 动态形状也需要 CPU-GPU 同步，因为它需要在 CPU 端涉及框架的内存管理。对于仅静态形状，不需要 CPU-GPU 同步。这如图 5 所示。

如文中所述，通过使用固定大小的张量和布尔掩码，我们可以消除动态大小张量所需的 CPU 同步。

NVIDIA 深度学习示例库中的 CUDA 图

单 GPU 的使用案例也可以从使用 CUDA 图中获得好处。这对于启动许多短内核和小批次的负载尤其如此。一个很好的例子是推荐系统的训练和推理。以下是我们从深度学习示例库中提供的 NVIDIA 深度学习推荐模型（DLRM）的初步基准测试结果。对于这个负载使用 CUDA 图可以显著提高训练和推理的速度。当使用非常小的批次大小时，这种效果尤为明显，因为 CPU 开销更为明显。

CUDA 图形正在积极集成到其他 PyTorch NGC 模型脚本和 NVIDIA Github 深度学习示例中。敬请期待更多使用示例。

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

呼吁行动：PyTorch v1.10 中的 CUDA 图形。

CUDA 图形可以为包含许多小型 GPU 内核的工作负载提供显著的好处，因为这些工作负载通常被 CPU 启动开销拖累。这一点在我们的 MLPerf 努力中得到了证明，我们正在优化 PyTorch 模型。包括 CUDA 图形在内的许多这些优化最终将被集成到我们的 PyTorch NGC 模型脚本集合和 NVIDIA Github 深度学习示例中。目前，您可以查看我们的开源 MLPerf 训练 v1.0 实现，这可以作为观察 CUDA 图形实际应用的良好起点。或者，尝试在您自己的工作负载上使用 PyTorch CUDA 图形 API。

我们感谢许多 NVIDIA 和 Facebook 的工程师就讨论和建议所做的工作：Karthik Mandakolathur US，Tomasz Grel，PLJoey Conway，Arslan Zulfiqar US

作者简介

Vinh NguyenDL 工程师，NVIDIA

Vinh 是一位深度学习工程师和数据科学家，已发表 50 多篇科学论文，获得 2500 多次引用。在 NVIDIA，他的工作涵盖了广泛的深度学习和 AI 应用，包括语音、语言和视觉处理以及推荐系统。

迈克尔·卡里利 高级技术工程师，英伟达

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

苏克鲁·布尔奇·埃里伊尔马兹 高级架构师，NVIDIA

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

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

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

Michelle Lin，NVIDIA 产品实习生

Michelle Lin 目前在美国加州大学伯克利分校攻读计算机科学和商业管理学士学位。她目前负责执行市场研究以及为 Magnum IO 创建营销资产等项目。

吉梅尔申娜塔莉亚，应用研究科学家，Facebook

吉梅尔申娜塔莉亚曾在 NVIDIA 和 Facebook 从事 GPU 性能优化工作，目前是 PyTorch 核心团队成员，与合作伙伴共同支持新的软件和硬件功能。

阿尔班·德马松，研究工程师，Facebook

阿尔班在机器学习和优化领域攻读博士学位，在此期间成为 PyTorch 开源软件的贡献者，加入 Facebook 后，主要负责维护一些核心库和功能（autograd、optim、nn）以及使 PyTorch 整体变得更好。

杨逸飞，Facebook 研究工程师

杨逸飞在麻省理工学院学习计算机科学后，又在斯坦福大学深造，之后加入 Facebook。他是 PyTorch 核心团队的一员，也是 PyTorch 的主要贡献者之一。