CUDAGraph 树¶

背景¶

CUDAGraph¶

有关 CUDAGraphs 的更长背景介绍，请阅读使用 CUDAGraphs 加速 pytorch。

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

CUDA Graphs 可以显著加速，特别是对于 CPU 开销高或计算量小的模型。它存在许多限制，例如需要运行相同的内核，使用相同的参数和依赖项以及内存地址。

控制流是不可能的
触发主机到设备同步的内核（例如 .item()）会出错
内核的所有输入参数都固定为记录时的状态
CUDA 内存地址是固定的，但是这些地址的内存值可以更改
没有必要的 CPU 操作或 CPU 侧效应

PyTorch CUDAGraph 集成¶

PyTorch 提供了一个便捷的包装器，围绕 CUDAGraphs，它可以处理与 PyTorch 的缓存分配器的一些棘手交互。

CachingAllocator 对所有新分配使用单独的内存池。在 CUDAGraph 记录期间，内存的核算、分配和释放与在 eager 运行时完全相同。在重放时，仅调用内核，并且分配器没有变化。在初始记录之后，分配器不知道用户程序中正在主动使用哪些内存。

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

制作图化可调用对象¶

制作图化可调用对象是 PyTorch 抽象，用于在一系列可调用对象之间共享单个内存池。图化可调用对象利用了 CUDA Graph 记录的特性，即缓存分配器会精确核算内存，从而安全地在单独的 CUDA Graph 记录之间共享内存。在每次调用中，输出都作为活动内存保留，防止一个可调用对象覆盖另一个可调用对象的活动内存。图化可调用对象只能以单一顺序调用；第一次运行的内存地址会被刻录到第二次运行中，依此类推。

TorchDynamo 之前的 CUDA Graphs 集成¶

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

CUDAGraph 树集成¶

与图可调用对象类似，CUDA Graph 树在所有图捕获之间使用单个内存池。但是，CUDA Graph 树不是要求单一的调用序列，而是创建单独的 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 Graphs 的相同约束适用：必须使用相同的参数（静态大小、地址等）调用相同的内核
在记录和重放之间必须观察到相同的内存模式：如果一个图的张量输出在记录期间在另一个图之后死亡，那么它也必须在重放期间这样做。
CUDA 池中的活动内存强制两个记录之间存在依赖关系
这些记录只能以单一顺序调用 1 - > 2 - > 4

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

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

首先，我们将命中我们已经在图 1 中记录的优化路径 CUDAGraph.replay()。然后我们将命中图 3。与之前一样，我们需要在记录之前预热图一次。在预热运行时，内存地址不固定，因此图 4 也将回退到 inductor，非 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 和 x' 驻留在不同的 CUDA 内存地址上。

仔细查看输入突变函数会发现有三种类型的输入

来自 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 树会为输入张量的每个唯一形状重新记录 CUDAGraph。这会导致单个 inductor 图有多个 CUDAGraph。当形状有限时（例如，推理中的批大小），重新记录 CUDAGraph 是有利可图的。但是，如果输入张量形状频繁更改甚至在每次调用时都更改，则重新记录 CUDAGraph 可能无利可图。Nvidia 在 CUDA Graph 中每个内核启动使用 64 KB 的设备内存，直到 CUDA 12.4 和驱动程序版本 550+。如果重新记录许多 CUDAGraph，则此内存成本可能非常可观。

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

NCCL 支持¶

CUDAGraph 树支持带有 nccl 运算符的函数。虽然 CUDAGraph 树对 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 树会检查函数是否满足这些要求，并在必要时跳过 CUDAGraph。在这里，我们列出跳过 CUDAGraph 的常见原因。

输入突变：CUDAGraph 树会跳过原地突变 eager 输入的函数。仍然支持原地突变参数和缓冲区，或来自 CUDAGraph 树管理函数的输出张量。有关更多详细信息，请参阅输入突变支持部分。
CPU 运算符：包含 CPU 运算符的函数将被跳过。请将该函数拆分为多个函数，并将 CUDAGraph 树应用于仅包含 GPU 运算符的函数。
多设备运算符：如果函数包含多个设备上的运算符，则会跳过该函数。目前，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 Graphs 没有很好的方法来处理来自先前调用的活动张量。

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

import torch

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

x = torch.randn(10, 10)
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 树中，我们不想在迭代之间添加意外的依赖关系，这会导致我们无法命中热路径，也不希望我们过早释放先前调用的内存。我们的启发式方法是在推理中，我们在每次 torch.compile 调用时开始新的迭代，在训练中，只要没有尚未调用的挂起后向，我们就执行相同的操作。如果这些启发式方法是错误的，您可以使用 torch.compiler.mark_step_begin() 标记新迭代的开始，或者在开始下一次运行之前克隆先前迭代的张量（在 torch.compile 之外）。

比较¶

潜在的陷阱	单独的 CudaGraph	CUDAGraph 树
内存可能会增加	在每次图编译时（新大小等）	如果您还在运行非 cudagraph 内存
记录	在图的任何新调用时	将在您程序中进行的任何新的、唯一的路径上重新记录
潜在的陷阱	一个图的调用将覆盖先前的调用	无法在模型的单独运行之间持久化内存 - 一个训练循环训练，或一次推理运行