CUDAGraph 树¶

背景¶

CUDAGraph¶

有关 CUDAGraph 的更详细背景信息，请阅读使用 CUDAGraph 加速 PyTorch。

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

CUDA 图可以提供显著的加速，特别是对于 CPU 开销高或计算量小的模型。但也存在一些限制，例如要求使用相同的参数、依赖关系和内存地址来运行相同的核函数。

不支持控制流
触发主机到设备同步（例如 .item()）的核函数会出错
核函数的所有输入参数都固定为其录制时的值
CUDA 内存地址是固定的，但这些地址上的内存值可以改变
不支持基本的 CPU 操作或 CPU 副作用

PyTorch CUDAGraph 集成¶

PyTorch 提供了对 CUDAGraphs 的便捷封装，用于处理与 PyTorch 缓存分配器的一些棘手交互。

CachingAllocator 为所有新分配使用一个独立的内存池。在 CUDAGraph 录制期间，内存的记账、分配和释放与即时执行（eager run）期间完全相同。在重放时，仅调用核函数，并且分配器没有变化。初始录制之后，分配器不知道用户程序中哪些内存正在被积极使用。

如果在即时分配和 cudagraph 分配之间使用独立的内存池，如果两者都分配了大量内存，可能会增加程序的内存使用量。

生成图形化可调用对象¶

Make Graphed Callables 是 PyTorch 的一个抽象层，用于在一系列可调用对象之间共享单个内存池。图形化可调用对象利用了 CUDA Graph 录制时缓存分配器对内存进行精确记账的事实，从而安全地在独立的 CUDA Graph 录制之间共享内存。在每次调用中，输出被保留为活动内存，防止一个可调用对象覆盖另一个的活动内存。图形化可调用对象只能按单一顺序调用；第一次运行的内存地址会被“烧录”到第二次运行中，依此类推。

TorchDynamo 先前 CUDA 图集成¶

使用 cudagraph_trees=False 运行时，内存不会在独立的图捕获之间重用，这可能导致较大的内存退化。即使对于没有图中断的模型，这也会有问题。前向传播和后向传播是独立的图捕获，因此前向和后向的内存池不共享。特别是，前向传播中保存的激活内存无法在后向传播中回收。

CUDAGraph Trees 集成¶

与图形化可调用对象类似，CUDA Graph Trees 在所有图捕获之间使用单个内存池。然而，CUDA Graph Trees 不要求按单一序列调用，而是创建独立的 CUDA Graph 捕获树。让我们看一个示例

@torch.compile(mode="reduce-overhead")
def foo(x):
    # GRAPH 1
    y = x * x * x
    # graph break triggered here
    if y.sum() > 0:
        # GRAPH 2
        z = y ** y
    else:
        # GRAPH 3
        z = (y.abs() ** y.abs())
    torch._dynamo.graph_break()
    # GRAPH 4
    return z * torch.rand_like(z)

# the first run warms up each graph, which does things like CuBlas or Triton benchmarking
foo(torch.arange(0, 10, device="cuda"))
# The second run does a CUDA Graph recording, and replays it
foo(torch.arange(0, 10, device="cuda"))
# Finally we hit the optimized, CUDA Graph replay path
foo(torch.arange(0, 10, device="cuda"))

在这个示例中，我们通过函数有两种独立的路径：1 -> 2 -> 4，或 1 -> 3 -> 4。

通过构建一个 CUDA Graph 录制序列（本例中为 1 -> 2 -> 4），我们在独立的录制之间共享单个内存池中的所有内存。我们添加不变量，以确保内存始终位于录制时的相同位置，并且用户程序中不存在可能被覆盖的活动张量。

CUDA 图的相同限制仍然适用：必须使用相同的参数（静态大小、地址等）调用相同的核函数
录制和重放之间必须观察到相同的内存模式：如果在录制过程中，一个图的张量输出在另一个图之后失效，那么在重放时也必须如此。
CUDA 池中的活动内存会强制两个录制之间存在依赖关系
这些录制只能按单一顺序调用 1 -> 2 -> 4

所有内存都共享在单个内存池中，因此与即时执行相比没有额外的内存开销。现在，如果我们遇到新路径并运行图 3 会发生什么？

图 1 被重放，然后我们遇到尚未录制的图 3。在图重放时，私有内存池不会更新，因此 y 未在分配器中反映。如果不加注意，我们会覆盖它。为了支持在重放其他图后重用相同的内存池，我们将内存池检查点回图 1 结束时的状态。现在我们的活动张量已反映在缓存分配器中，我们可以安全地运行新图了。

首先，我们将命中已在图 1 中录制过的优化过的 CUDAGraph.replay() 路径。然后我们将命中图 3。就像之前一样，我们需要在录制前预热图一次。在预热运行中，内存地址不固定，因此图 4 也会回退到 inductor，非 cudagraph 调用。

第二次命中图 3 时，我们已经预热完毕并准备好录制。我们录制图 3，然后再次录制图 4，因为输入内存地址已改变。这创建了一个 CUDA Graph 录制树。一个 CUDA Graph Tree！

  1
 / \\
2   3
 \\   \\
  4   4

输入变异支持¶

输入变异函数是指对输入张量进行原地写入的函数，如下所示

def foo(x, y):
    # mutates input x
    x.add_(1)
    return x + y

输入变异函数通常给 CUDAGraph Trees 带来挑战。由于 CUDAGraph 对 CUDA 内存地址的静态要求，对于每个输入张量 x，CUDAGraph Trees 可能会分配一个静态内存地址 x’。执行期间，CUDAGraph Trees 首先将输入张量 x 复制到静态内存地址 x’，然后重放录制的 CUDAGraph。对于输入变异函数，x’ 会原地更新，但这不会反映到输入张量 x 上，因为 x 和 x’ 位于不同的 CUDA 内存地址。

仔细观察输入变异函数，会发现有三种类型的输入

来自 eager 的输入：我们假设这些张量的输入地址在每次执行时都会变化。由于 cudagraphs 会冻结内存地址，因此在图录制和执行之前，我们需要将这些输入复制到一个静态地址张量。
参数和缓冲区：我们假设（并通过运行时检查）这些张量在每次执行时具有相同的张量地址。我们不需要复制其内容，因为录制的内存地址与执行的内存地址相同。
作为 CUDAGraph Trees 先前输出的张量：由于 cudagraph 的输出张量地址是固定的，如果我们运行 CUDAGraph1，然后运行 CUDAGraph2，则从 CUDAGraph1 进入 CUDAGraph2 的输入将具有固定的内存地址。这些输入与参数和缓冲区一样，不需要复制到静态地址张量。我们在运行时检查以确保这些输入是稳定的，如果不是，我们将重新录制。

CUDAGraph Trees 支持对参数和缓冲区以及作为 CUDAGraph Trees 先前输出的张量进行输入变异。对于来自 eager 的输入变异，CUDAGraph Trees 将在不使用 CUDAGraph 的情况下运行函数，并发出 Skipping due to mutated inputs（因变异输入而跳过）日志。以下示例展示了 CUDAGraph Trees 对作为 CUDAGraph Trees 先前输出的张量的支持。

import torch

@torch.compile(mode="reduce-overhead")
def foo(x):
    return x + 1

@torch.compile(mode="reduce-overhead")
def mut(x):
    return x.add_(2)

# Enable input mutation support
torch._inductor.config.triton.cudagraph_support_input_mutation = True

for i in range(3):
    torch.compiler.cudagraph_mark_step_begin()
    inp = torch.rand([4], device="cuda")

    # CUDAGraph is applied since `foo` does not mutate `inp`
    tmp = foo(inp)
    # Although `mut` mutates `tmp`, which is an output of a CUDAGraph
    # managed function. So CUDAGraph is still applied.
    mut(tmp)


torch.compiler.cudagraph_mark_step_begin()
inp = torch.rand([4], device="cuda")

tmp = foo(inp)
# While `tmp` is a CUDAGraph Tree managed function's output, `tmp.clone()`
# is not. So CUDAGraph is not applied to `mut` and there is a log
# `skipping cudagraphs due to mutated inputs`
mut(tmp.clone())

要为变异 eager 输入的函数启用 CUDAGraph Trees，请重写函数以避免输入变异。

注意

通过将 torch._inductor.config.cudagraph_support_input_mutation 设置为 True 来启用“减少开销”模式下的输入变异支持。

动态形状支持¶

动态形状意味着输入张量在函数调用之间具有不同的形状。由于 CUDAGraph 要求固定的张量地址，CUDAGraph Trees 会为输入张量的每个唯一形状重新录制一个 CUDAGraph。这导致单个 inductor 图对应多个 CUDAGraph。当形状有限时（例如，推理中的批量大小），重新录制 CUDAGraph 是有利的。然而，如果输入张量形状频繁变化，甚至在每次调用时都变化，则重新录制 CUDAGraph 可能不划算。在 CUDA 12.4 和驱动版本 550+ 之前，Nvidia 在 CUDAGraph 中每次核函数启动使用 64 KB 设备内存。多次重新录制 CUDAGraph 会导致显著的内存开销。

对于输入张量形状频繁变化的函数，我们建议将输入张量填充到几个固定的张量形状，以便仍然受益于 CUDAGraph。此外，设置 torch._inductor.config.triton.cudagraph_skip_dynamic_graphs=True 允许跳过对具有动态形状输入的函数进行 cudagraphing，而只对具有静态输入张量形状的函数进行 cudagraphing。

NCCL 支持¶

CUDAGraph Trees 支持包含 nccl 算子的函数。虽然 CUDAGraph Trees 对 CUDAGraph 进行逐设备录制，但 NCCL 支持跨设备通信。

@torch.compile(mode="reduce-overhead")
def func(x):
    y = x * x
    y = torch.distributed.all_reduce(y, op=torch.distributed.ReduceOp.SUM)
    x = torch.nn.functional.silu(x)
    return x * y

跳过 CUDAGraph 的原因¶

由于 CUDAGraph 有静态输入张量地址和不支持 CPU 算子等要求，CUDAGraph Trees 会检查函数是否满足这些要求，并在必要时跳过 CUDAGraph。这里列出了跳过 CUDAGraph 的常见原因。

输入变异：CUDAGraph Trees 跳过原地变异 eager 输入的函数。仍然支持原地变异参数和缓冲区，或由 CUDAGraph Tree 管理的函数输出的张量。有关更多详细信息，请参阅“输入变异支持”部分。
CPU 算子：跳过包含 CPU 算子的函数。请将函数拆分成多个函数，并仅对包含 GPU 算子的函数应用 CUDAGraph Trees。
多设备算子：如果函数包含跨多个设备的算子，则会跳过。目前，CUDAGraph 是基于每个设备应用的。请使用 NCCL 等支持的库进行跨设备通信。有关更多详细信息，请参阅“NCCL 支持”部分。
未支持的自由符号：未支持的自由符号通常在动态形状期间发生。CUDAGraph Trees 当前会为每个唯一的输入张量形状录制一个 CUDAGraph。有关更多详细信息，请参阅“动态形状支持”。
不兼容的算子：如果函数包含不兼容的算子，则 CUDAGraph Trees 会跳过该函数。请将函数中的这些算子替换为支持的算子。我们提供了不兼容算子的完整列表

aten._fused_moving_avg_obs_fq_helper.default
aten._fused_moving_avg_obs_fq_helper_functional.default
aten.multinomial.default
fbgemm.dense_to_jagged.default
fbgemm.jagged_to_padded_dense.default
run_and_save_rng_state
run_with_rng_state
aten._local_scalar_dense
aten._assert_scalar

当 torch.are_deterministic_algorithms_enabled() 时，以下算子不兼容。

aten._fused_moving_avg_obs_fq_helper.default
aten._fused_moving_avg_obs_fq_helper_functional.default
aten.multinomial.default
fbgemm.dense_to_jagged.default
fbgemm.jagged_to_padded_dense.default
run_and_save_rng_state
run_with_rng_state
aten._local_scalar_dense
aten._assert_scalar

限制¶

由于 CUDA Graph 固定内存地址，CUDA Graph 没有很好的方法处理来自先前调用的活动张量。

假设我们正在使用以下代码基准测试推理运行

import torch

@torch.compile(mode="reduce-overhead")
def my_model(x):
    y = torch.matmul(x, x)
    return y

x = torch.randn(10, 10, device="cuda")
y1 = my_model(x)
y2 = my_model(x)
print(y1)
# RuntimeError: Error: accessing tensor output of CUDAGraphs that has been overwritten by a subsequent run.

在单独的 CUDA Graph 实现中，第一次调用的输出将被第二次调用覆盖。在 CUDAGraph Trees 中，我们不希望在迭代之间添加意外的依赖关系，导致无法命中热路径，也不希望过早释放先前调用中的内存。我们的启发式方法是：在推理中，我们对 torch.compile 的每次调用都开始一个新的迭代；在训练中，只要没有待处理的后向传播未被调用，我们也这样做。如果这些启发式方法不正确，您可以使用 torch.compiler.mark_step_begin() 标记新迭代的开始，或者在开始下一次运行之前克隆先前迭代的张量（在 torch.compile 之外）。

比较¶

注意事项	独立的 CudaGraph	CUDAGraph 树
内存可能增加	在每次图编译时（新大小等）	如果你同时运行非 cudagraph 内存
录制	在任何新的图调用时	将在你通过程序采用的任何新的、唯一的路径上重新录制
注意事项	调用一个图将覆盖先前的调用	无法在模型的独立运行之间持久化内存 - 例如一个训练循环或一次推理运行