今天，我们很高兴宣布 PyTorch 引入了一项新的高级 CUDA 功能：CUDA Graphs。现代深度学习框架具有复杂的软件堆栈，这会产生与向 GPU 提交每个操作相关的显著开销。当深度学习工作负载为了性能而强大地扩展到多个 GPU 时，每个 GPU 操作所用的时间会减少到仅几微秒，在这些情况下，框架的高工作提交延迟通常会导致 GPU 利用率低下。随着 GPU 速度越来越快，工作负载扩展到更多设备，工作负载遭受这些启动引起的停顿的可能性也会增加。为了克服这些性能开销，NVIDIA 工程师与 PyTorch 开发者合作，在 PyTorch 中原生启用了 CUDA graph 执行。此设计对于扩展 NVIDIA 的 MLPerf 工作负载（在 PyTorch 中实现）至 4000 多个 GPU 以实现破纪录的性能至关重要。
PyTorch 中对 CUDA graphs 的支持只是 NVIDIA 和 Facebook 工程师长期合作的又一个例证。例如，torch.cuda.amp 以半精度进行训练，同时保持以单精度实现的网络准确性，并尽可能自动利用张量核心。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 操作的机制，从而减少启动开销。
CUDA graphs 的优势可以通过图 1 中的简单示例来演示。在顶部，CPU 逐个启动一系列短内核。CPU 启动开销在内核之间创建了明显的间隔。如果我们将这一系列内核替换为 CUDA graph，最初我们需要花费一些额外的时间来构建 graph，并在第一次执行时一次性启动整个 graph，但后续执行将非常快，因为内核之间几乎没有间隔。当相同的操作序列重复多次时，例如在许多训练步骤中，这种差异会更加明显。在这种情况下，构建和启动 graph 的初始成本将在整个训练迭代次数中摊销。有关该主题的更全面的介绍，请参阅我们的博客 CUDA Graphs 入门 和 GTC 演讲 轻松使用 CUDA Graphs。
图 1. 使用 CUDA graphs 的优势
NCCL 支持 CUDA graphs
前面提到的减少启动开销的优势也扩展到了 NCCL 内核启动。NCCL 支持基于 GPU 的集体通信和 P2P 通信。借助 NCCL 对 CUDA graphs 的支持，我们可以消除 NCCL 内核启动开销。
此外，由于各种 CPU 负载和操作系统因素，内核启动时序可能无法预测。这种时间偏差可能对 NCCL 集体操作的性能有害。借助 CUDA graphs，内核被聚集在一起，从而使分布式工作负载中各个 rank 的性能保持一致。这在大型集群中尤其有用，即使是单个缓慢的节点也可能降低整个集群级别的性能。
对于分布式多 GPU 工作负载，NCCL 用于集体通信。如果我们查看利用数据并行性的神经网络训练，在没有 NCCL 对 CUDA graphs 的支持的情况下，我们需要为前向/反向传播和 NCCL AllReduce 各自单独启动。相比之下，借助 NCCL 对 CUDA graphs 的支持，我们可以通过将前向/反向传播和 NCCL AllReduce 全部集中在一个 graph 启动中来减少启动开销。
图 2. 查看典型的神经网络，NCCL AllReduce 的所有内核启动都可以捆绑到一个 graph 中，以减少启动开销时间。
PyTorch CUDA Graphs
从 PyTorch v1.10 开始，CUDA graphs 功能作为一组 beta API 提供。
API 概述
PyTorch 支持使用 流捕获 来构建 CUDA graphs，这会将 CUDA 流置于捕获模式。发布到捕获流的 CUDA 工作实际上不会在 GPU 上运行。相反，该工作会被记录在 graph 中。捕获后，可以启动 graph 以根据需要多次运行 GPU 工作。每次重放都使用相同的参数运行相同的内核。对于指针参数，这意味着使用相同的内存地址。通过在每次重放之前用新数据（例如，来自新批次）填充输入内存，您可以在新数据上重新运行相同的工作。
重放 graph 会牺牲典型 eager 执行的动态灵活性，以换取大大减少的 CPU 开销。graph 的参数和内核是固定的，因此 graph 重放会跳过参数设置和内核分发的所有层，包括 Python、C++ 和 CUDA 驱动程序开销。在底层，重放通过对 cudaGraphLaunch 的单次调用将整个 graph 的工作提交给 GPU。重放中的内核在 GPU 上执行速度也稍快，但消除 CPU 开销是主要优势。
如果您的网络全部或部分是 graph 安全的（通常这意味着静态形状和静态控制流，但请参阅其他约束），并且您怀疑其运行时至少在某种程度上受 CPU 限制，则应尝试 CUDA graphs。
API 示例
PyTorch 通过原始 torch.cuda.CUDAGraph 类和两个便捷包装器 torch.cuda.graph 和 torch.cuda.make_graphed_callables 公开 graph。
torch.cuda.graph 是一个简单、通用的上下文管理器，可在其上下文中捕获 CUDA 工作。在捕获之前，通过运行几次 eager 迭代来预热要捕获的工作负载。预热必须在辅助流上进行。由于 graph 在每次重放中都从相同的内存地址读取和写入，因此您必须在捕获期间维护对保存输入和输出数据的张量的长期引用。要在新输入数据上运行 graph，请将新数据复制到捕获的输入张量，重放 graph，然后从捕获的输出张量读取新输出。
如果整个网络都是捕获安全的，则可以像以下示例中那样捕获和重放整个网络。
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 仅对捕获安全的部分进行 graph 化。接下来将对此进行演示。
make_graphed_callables 接受可调用对象（函数或 nn.Module），并返回 graph 化的版本。默认情况下，make_graphed_callables 返回的可调用对象是自动微分感知的，并且可以在训练循环中用作您传递的函数或 nn.Module 的直接替代品。make_graphed_callables 在内部创建 CUDAGraph 对象，运行预热迭代，并根据需要维护静态输入和输出。因此，（与 torch.cuda.graph 不同）您无需手动处理这些。
在以下示例中，数据相关的动态控制流意味着网络无法端到端捕获，但 make_graphed_callables() 允许我们捕获和运行 graph 安全的部分，无论如何都作为 graph
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 至关重要，从而在各个方面创造了新的记录。我们在下面说明了两个 MLPerf 工作负载，其中使用 CUDA graphs 观察到了最显著的收益，产生了高达约 1.7 倍的加速。
| GPU 数量 | CUDA graph 的加速 | |
|---|---|---|
| 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 时间线图中识别出此问题。下图显示了在进行 graph 化之前，每个 GPU 批次大小为 1 的 Mask R-CNN 的“backbone”部分。绿色部分显示 CPU 负载，而蓝色部分显示 GPU 负载。在此配置文件中，我们看到 CPU 的负载已达到 100% 的峰值，而 GPU 大部分时间处于空闲状态，GPU 内核之间存在大量空白空间。
图 3：Mask R-CNN 的 NSight 时间线图
当张量形状为静态时，CUDA graphs 可以自动消除 CPU 开销。在第一步中捕获所有内核调用的完整 graph，在后续步骤中，使用单个操作启动整个 graph，从而消除所有 CPU 开销，如图 4 所示。
图 4：CUDA graphs 优化
通过 graph 化，我们看到 GPU 内核紧密地打包在一起，并且 GPU 利用率保持较高水平。graph 化的部分现在运行时间为 6 毫秒，而不是 31 毫秒，加速了 5 倍。我们没有对整个模型进行 graph 化，主要是 resnet backbone，这带来了约 1.7 倍的总体加速。为了增加 graph 的范围，我们对软件堆栈进行了一些更改，以消除一些 CPU-GPU 同步点。在 MLPerf v1.0 中，这项工作包括更改 torch.randperm 函数的实现，以使用 CUB 而不是 Thrust，因为后者是同步 C++ 模板库。这些改进在最新的 NGC 容器中可用。
BERT
同样，通过 graph 捕获模型，我们消除了 CPU 开销和随之而来的同步开销。CUDA graphs 实现为我们的最大规模 BERT 配置带来了 1.12 倍的性能提升。为了最大限度地发挥 CUDA graphs 的优势，保持 graph 的范围尽可能大非常重要。为了实现这一点，我们修改了模型脚本以删除执行期间的 CPU-GPU 同步，以便可以 graph 捕获整个模型。此外，我们还确保执行期间的张量大小在 graph 范围内是静态的。例如，在 BERT 中，只有一部分总 token 会对损失函数做出贡献，这由预生成的掩码张量决定。从此掩码中提取有效 token 的索引，并使用这些索引来收集对损失有贡献的 token，会导致张量具有动态形状，即形状在迭代中不是恒定的。为了确保张量大小是静态的，我们在损失计算中使用了静态形状张量，其中掩码用于指示哪些元素是有效的，而不是使用动态形状张量。因此，所有张量形状都是静态的。动态形状还需要 CPU-GPU 同步，因为它必须涉及 CPU 端的框架内存管理。使用仅静态形状，则不需要 CPU-GPU 同步。这在图 5 中显示。
图 5. 通过使用固定大小的张量和布尔掩码（如文本中所述），我们能够消除动态大小张量所需的 CPU 同步
NVIDIA DL 示例集合中的 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
CUDA graphs 可以为包含许多小型 GPU 内核的工作负载提供实质性的好处，因此可以减轻 CPU 启动开销的负担。这已在我们的 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 英伟达 DL 工程师
Vinh 是一位深度学习工程师和数据科学家，已发表 50 多篇科学文章，被引用超过 2500 次。在 NVIDIA，他的工作涵盖广泛的深度学习和 AI 应用，包括语音、语言和视觉处理以及推荐系统。
Michael Carilli 英伟达高级开发者技术工程师
Michael 曾在空军研究实验室工作，致力于优化现代并行架构的 CFD 代码。他拥有加州大学圣巴巴拉分校计算物理学博士学位。作为 PyTorch 团队的成员，他专注于使 GPU 训练对内部团队、外部客户和 Pytorch 社区用户来说更快、数值更稳定且更容易（更轻松）。
Sukru Burc Eryilmaz 英伟达 Dev Arch 高级架构师
Sukru 获得了斯坦福大学博士学位和比尔肯特大学理学学士学位。他目前致力于提高单节点规模和超级计算机规模的神经网络训练的端到端性能。
Vartika Singh 英伟达 DL 框架和库技术合作伙伴主管
Vartika 领导的团队在云和分布式计算、扩展和 AI 的交汇处工作，影响着大型企业的设计和战略。她目前与 NVIDIA 内外的主要框架和编译器组织及开发者合作，帮助设计在 NVIDIA 硬件上高效且最佳地工作。
Michelle Lin 英伟达产品实习生
Michelle 目前在加州大学伯克利分校攻读计算机科学和工商管理本科 degree。她目前正在管理项目执行，例如进行市场调研和为 Magnum IO 创建营销资产。
Natalia Gimelshein Facebook 应用研究科学家
Natalia Gimelshein 曾在 NVIDIA 和 Facebook 从事深度学习工作负载的 GPU 性能优化工作。她目前是 PyTorch 核心团队的成员，与合作伙伴合作，无缝支持新的软件和硬件功能。
Alban Desmaison Facebook 研究工程师
Alban 学习了工程学，并获得了机器学习和优化博士学位，在此期间，他在加入 Facebook 之前是 PyTorch 的 OSS 贡献者。他的主要职责是维护一些核心库和功能（autograd、optim、nn），并致力于使 PyTorch 总体上变得更好。
Edward Yang Facebook 研究工程师
Edward 在麻省理工学院和斯坦福大学学习计算机科学，之后加入 Facebook。他是 PyTorch 核心团队的一员，也是 PyTorch 的主要贡献者之一。
