CUDAGraph 树

背景

CUDAGraph

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

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

CUDA 图可以带来很大的加速效果，尤其是在 CPU 开销很大或计算量很小的模型中。但是，它也有一些限制，例如需要使用相同的参数和依赖项以及内存地址运行相同的内核。

• 不支持控制流

• 触发主机到设备同步的内核（例如 .item()）会导致错误

• 内核的所有输入参数都固定为记录时的值

• CUDA 内存地址是固定的，但这些地址处的内存值可以更改

• 没有必要的 CPU 操作或 CPU 副作用

PyTorch CUDAGraph 集成

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

缓存分配器为所有新的分配使用一个单独的内存池。在 CUDAGraph 记录期间，内存的分配和释放与急切运行期间完全相同。在回放时，只调用内核，分配器不会发生任何变化。在初始记录之后，分配器不知道用户程序中哪些内存正在被积极使用。

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

创建图形化可调用对象

创建图形化可调用对象 是一个 PyTorch 抽象，用于在一系列可调用对象之间共享一个内存池。图形化可调用对象利用了在 CUDA 图记录期间，缓存分配器会精确地记录内存这一事实，以便在单独的 CUDA 图记录之间安全地共享内存。在每次调用中，输出都作为活动内存保留，防止一个可调用对象覆盖另一个可调用对象的活动内存。图形化可调用对象只能按单个顺序调用；第一次运行的内存地址会被写入第二次运行，依此类推。

TorchDynamo 之前的 CUDA 图集成

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

CUDAGraph 树集成

与图形化可调用对象类似，CUDA 图树在所有图形捕获之间使用单个内存池。但是，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 图记录的记录带，例如 1 -> 2 -> 4，我们在单个内存池中共享所有内存。我们添加不变量以确保内存始终位于记录时的相同位置，并且用户程序中不存在可能被覆盖的活动张量。

• CUDA 图的相同约束适用：必须使用相同的参数（静态大小、地址等）调用相同的内核。

• 在记录和回放之间必须观察到相同的内存模式：如果一个图的一个张量输出在记录期间在另一个图之后消失，那么它在回放期间也必须消失。

• CUDA 池中的活动内存强制了两个记录之间的依赖关系

• 这些记录只能按单个顺序调用 1 - > 2 -> 4

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

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

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

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

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

输入变异支持

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

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

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

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

• 来自 Eager 的输入：我们假设这些张量的输入张量地址在每次执行时都会变化。因为 cudagraph 会冻结内存地址，所以在记录和执行图之前，我们需要将这些输入复制到一个静态地址张量。

• 参数和缓冲区：我们假设（并在运行时检查）这些张量在每次执行时都具有相同的张量地址。我们不需要复制它们的内容，因为记录的内存地址将与执行的内存地址相同。

• 来自 CUDA 图树的先前输出的张量：因为 cudagraph 的输出张量地址是固定的，如果我们运行 CUDAGraph1，然后运行 CUDAGraph2，那么来自 CUDAGraph1 输入到 CUDAGraph2 的输入将具有固定的内存地址。这些输入，就像参数和缓冲区一样，不需要复制到静态地址张量。我们在运行时检查以确保这些输入是稳定的，如果不是，我们将重新记录。

CUDA 图树支持对参数、缓冲区以及来自 CUDA 图树的先前输出的张量的输入变异。对于来自 Eager 的输入的变异，CUDA 图树将运行该函数而不使用 CUDA 图，并发出由于输入变异而跳过日志。以下示例显示了 CUDA 图树对来自 CUDA 图树的先前输出的张量的支持。

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 的输入的函数启用 CUDA 图树，请重写该函数以避免输入变异。

注意

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

动态形状支持

动态形状 表示输入张量在不同函数调用之间具有不同的形状。由于 CUDA 图需要固定的张量地址，因此 CUDA 图树为每个输入张量的唯一形状重新记录 CUDA 图。这导致单个 inductor 图有多个 CUDA 图。当形状有限时（例如，推理中的批次大小），重新记录 CUDA 图是有益的。但是，如果输入张量形状频繁变化甚至在每次调用时都变化，重新记录 CUDA 图可能不划算。Nvidia 在 CUDA 图中每个内核启动使用 64 KB 的设备内存，直到 CUDA 12.4 和驱动程序版本 550+。如果重新记录 CUDA 图很多，此内存成本可能很大。

对于输入张量形状频繁变化的函数，我们建议将输入张量填充到几个固定的张量形状，以继续享受 CUDA 图带来的好处。此外，设置 torch._inductor.config.triton.cudagraph_skip_dynamic_graphs=True 可以跳过对具有动态形状输入的函数进行 cudagraph，而只对具有静态输入张量形状的函数进行 cudagraph。

NCCL 支持

CUDA 图树支持包含 nccl 运算符的函数。虽然 CUDA 图树对 CUDA 图执行每个设备记录，但 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

跳过 CUDA 图的原因

由于 CUDA 图有一些要求，例如静态输入张量地址和不支持 CPU 运算符，因此 CUDA 图树会检查函数是否满足这些要求，并在必要时跳过 CUDA 图。这里，我们列出了跳过 CUDA 图的常见原因。

• 输入变异：CUDA 图树跳过就地变异 Eager 输入的函数。就地变异参数和缓冲区或来自 CUDA 图树管理函数的输出张量仍然受支持。有关更多详细信息，请参阅输入变异支持部分。

• CPU 运算符：包含 CPU 运算符的函数将被跳过。请将函数拆分为多个函数，并在仅包含 GPU 运算符的函数上应用 CUDA 图树。

• 多设备运算符：如果函数包含多个设备上的运算符，则会跳过该函数。目前，CUDA 图是在每个设备的基础上应用的。请使用受支持的库（如 NCCL）进行跨设备通信。有关更多详细信息，请参阅NCCL 支持部分。

• 未备份的自由符号：未备份的自由符号通常发生在 动态形状 期间。CUDA 图树目前为每个唯一的输入张量形状记录一个 CUDA 图。有关更多详细信息，请参阅动态形状支持。

• 不兼容的运算符：如果函数包含不兼容的运算符，则 CUDA 图树会跳过该函数。请在函数中用受支持的运算符替换这些运算符。我们展示了不兼容运算符的完整列表

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 图会固定内存地址，所以 CUDA 图没有很好的方法来处理来自先前调用的活动张量。

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

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 图实现中，第一次调用的输出将被第二次调用覆盖。在 CUDA 图树中，我们不想在迭代之间添加意外的依赖关系，这会导致我们无法命中热点路径，也不想过早地释放先前调用的内存。在推理中，我们的启发式方法是在每次调用 torch.compile 时启动一个新的迭代，在训练中也是如此，只要没有尚未调用的挂起的反向传播。如果这些启发式方法不正确，可以使用 torch.compiler.mark_step_begin() 标记新迭代的开始，或者在开始下一次运行之前克隆先前迭代的张量（在 torch.compile 之外）。

比较

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