CUDAGraph 树 ¶

背景 ¶

CUDAGraph¶

想要了解 CUDAGraph 的更详细背景，请阅读《使用 CUDAGraph 加速 PyTorch》。

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

CUDA 图可以提供很大的速度提升，尤其是对于具有高 CPU 开销或小计算量的模型。由于要求使用相同的内核、相同的参数和依赖关系以及内存地址，因此存在一些限制。

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

PyTorch CUDAGraph 集成

PyTorch 提供了一个方便的 CUDAGraph 包装器，该包装器处理了 PyTorch 缓存分配器的一些复杂交互

CachingAllocator 使用单独的内存池来处理所有新的分配。在 CUDAGraph 记录期间，内存的分配、使用和释放与急切运行期间完全一致。在回放时，仅调用内核，分配器没有任何变化。在初始记录之后，分配器不知道用户程序中哪些内存正在被积极使用。

如果两个分配器都分配了大量的内存，那么在急切分配和 cudagraph 分配之间使用单独的内存池可能会增加程序的内存。

创建图调用可调用项

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

TorchDynamo 之前 CUDA 图集成 ¶

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

CUDAGraph 树集成 ¶

与图可调用类似，CUDA 图树在所有图捕获中使用单个内存池。然而，它不需要单个调用序列，而是创建独立的 CUDA 图捕获树。让我们来看一个示例：

@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 Graph 的约束相同：必须使用相同的参数调用相同的内核（静态大小、地址等）。
记录和回放之间必须观察到相同的内存模式：如果在记录期间一个图的一个张量输出在另一个图之后死亡，那么在回放期间也必须如此。
在 CUDA 池中的实时内存强制两个录制之间产生依赖
这些录制只能按顺序 1 -> 2 -> 4 调用

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

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

首先，我们会访问已经记录在图 1 中的优化后的 CUDAGraph.replay()路径。然后我们会访问图 3。就像之前一样，我们需要在记录之前预热一次图。在预热运行中，内存地址是不固定的，所以图 4 也会回退到非 cudagraph 的感应器调用。

第二次访问图 3 时，我们已经预热并准备好记录。我们记录图 3，然后由于输入内存地址已更改，再次记录图 4。这创建了一个 CUDA Graph 记录的树。一个 CUDA Graph 树！

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

输入变异支持

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

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

输入变异函数通常会给 CUDAGraph 树带来挑战。由于 CUDAGraph 对静态 CUDA 内存地址的要求，对于每个输入张量 x，CUDAGraph 树可能会分配一个静态内存地址 x'。在执行过程中，CUDAGraph 树首先将输入张量 x 复制到静态内存地址 x'，然后回放记录的 CUDAGraph。对于输入变异函数，x'是就地更新的，但由于 x 和 x'位于不同的 CUDA 内存地址上，因此 x 上的更新不会反映在输入张量 x 上。

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

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

CUDAGraph 树支持对参数、缓冲区和先前来自 CUDAGraph 树的张量的输入突变。对于来自 eager 的输入突变，CUDAGraph 树将在没有 CUDAGraph 的情况下运行函数，并发出由于突变输入而跳过的日志。以下示例显示了 CUDAGraph 树对先前来自 CUDAGraph 树的张量的支持。

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 树，请重新编写函数以避免输入突变。

注意

通过设置 torch._inductor.config.cudagraph_support_input_mutation = True 以启用“reduce-overhead”模式的输入突变支持。

动态形状支持 ¶

动态形状意味着输入张量在函数调用过程中具有不同的形状。由于 CUDAGraph 需要固定的张量地址，CUDAGraph Trees 会为输入张量的每个唯一形状重新记录 CUDAGraph。这会导致单个感应图有多个 CUDAGraph。当形状有限（例如，推理中的批大小）时，重新记录 CUDAGraph 是有利可图的。然而，如果输入张量形状频繁变化，甚至在每个调用中，重新记录 CUDAGraph 可能并不划算。Nvidia 在 CUDAGraph 中为每个内核启动使用 64 KB 的设备内存，直到 CUDA 12.4 和驱动程序版本 550+。这种内存成本在多次 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 树。
多设备操作符：如果函数包含多个设备上的操作符，则该函数将被跳过。目前，CUDAGraph 是按设备基础应用。请使用如 NCCL 等支持的库进行跨设备通信。请参阅 NCCL 支持部分以获取更多详细信息。
释放未备份的符号：释放未备份的符号通常发生在动态形状期间。CUDAGraph 树目前为每个唯一的输入张量形状记录一个 CUDAGraph。请参阅动态形状支持部分以获取更多详细信息。
不兼容的操作符：如果函数包含不兼容的操作符，CUDAGraph 树将跳过该函数。请将这些操作符替换为函数中的支持操作符。我们展示了不兼容操作符的详尽列表：

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 内存
录音	在调用任何新图形时	在您程序中采取的任何新、独特的路径上都将重新录制
脚本陷阱	调用一个图会覆盖之前的调用	无法在您的模型的不同运行之间持久化内存 - 一个训练循环训练，或一次推理运行