CUDA 图形树

背景

CUDAGraph

了解更多关于 CUDAGraphs 的背景信息，请阅读使用 CUDAGraphs 加速 PyTorch。

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

CUDA 图可以显著提高速度，尤其是在模型具有高 CPU 开销或少量计算任务时。然而，由于要求在相同条件下重复执行相同的内核、传递相同的参数和处理相同的依赖关系及内存地址，这也会带来一些限制。

无法实现控制流
当内核触发主机到设备的同步（例如 .item()）时出现错误
所有内核的输入参数都被固定为记录时的值
CUDA 内存地址是固定的，但这些地址处的内存值可以改变。
无必要的CPU操作和副作用

PyTorch CUDAGraph 集成

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

CachingAllocator 使用单独的内存池进行所有新的分配。在 CUDAGraph 记录期间，内存的会计、分配和释放与急切运行时完全相同。重播时，仅调用内核，并且分配器没有更改。初始记录之后，分配器无法确定用户程序中正在使用的内存。

在急切分配和 cuGraph 分配之间使用独立的内存池，如果两个过程都分配了大量内存，则可能增加程序的内存占用。

创建可调用的图形函数

Make Graphed Callables 是 PyTorch 中的一个抽象概念，用于在一系列可调用对象之间共享一个内存池。Graphed Callables 利用了 CUDA 图录制过程中缓存分配器精确记录内存这一特性，在不同的 CUDA 图录制之间安全地共享内存。每次调用时，输出会被保留为活动内存，防止一个可调用对象覆盖另一个的活动内存。此外，Graphed Callables 只能以固定的顺序进行调用；第一次运行中的内存地址会固定到后续的运行中。

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 Graph 记录带（例如，1 -> 2 -> 4）来共享单个内存池中的所有内存。我们添加不变量以确保内存位置与记录时一致，并且用户程序中不存在可能会被覆盖的活动张量。

同样适用于 CUDA 图形的约束：必须使用相同的参数（如静态大小和地址）来调用相同的内核。
在记录和回放过程中，必须保持相同的内存模式：如果一个图中的张量输出在另一个图之后被销毁，则在回放时也应如此。
CUDA池中的实时内存会迫使两个记录之间存在依赖关系
这些录音只能按固定的顺序调用：1 -> 2 -> 4

所有内存都在一个共享的内存池中，因此不会像即时模式那样产生额外的内存开销。那么，如果我们遇到新的路径并运行图3会发生什么?

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

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

当第二次遇到图3时，我们已经准备就绪并开始记录。由于输入内存地址发生了变化，我们先记录图3，再重新记录图4。这样就形成了一棵CUDA 图录制树。即CUDA 图树。

  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 上。

仔细查看输入变异函数，可以发现有三种类型的输入。

来自eager模式的输入：这些张量的输入地址会在每次执行时发生变化。因为cudagraphs会冻结内存地址，所以在记录图和执行前，我们需要将这些输入复制到一个静态地址张量中。
参数和缓冲区: 我们假设并进行运行时检查，这些张量在每次执行中的内存地址保持不变。因此，由于记录的内存地址与实际执行时的内存地址一致，所以无需复制它们的内容。
来自 CUDAGraph Trees 的先前输出张量：由于 cudagraph 的输出张量地址是固定的，如果我们先运行 CUDAGraph1，然后运行 CUDAGraph2，那么从 CUDAGraph1 输入到 CUDAGraph2 的数据将具有固定内存地址。这些输入（如参数和缓冲区）不需要复制到静态地址张量中。我们在运行时检查这些输入是否稳定，如果不稳定则需要重新记录。

CUDAGraph Trees支持对参数、缓冲区以及先前由CUDAGraph Trees生成的张量进行输入修改。对于来自即时模式的输入修改，CUDAGraph Trees将不使用CUDAGraph运行函数，并记录由于输入被修改而跳过的日志信息。以下示例展示了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。这会导致单个生成器图中有多个CUDAGraph。当存在有限的形状（例如推理中的批处理大小）时，重新记录CUDAGraph是有利的。然而，如果输入张量的形状频繁变化甚至每次调用都不同，则重新记录CUDAGraph可能不会带来好处。在CUDA 12.4及Driver Version 550+之前的版本中，Nvidia每个内核启动使用64 KB的设备内存用于CUDAGraph。因此，如果进行多次CUDAGraph重新记录，会显著增加内存成本。

对于输入张量形状经常变化的函数，我们建议将输入张量填充为几个固定的形状，以便继续从 CUDAGraph 中获益。此外，可以通过设置 torch._inductor.config.triton.cudagraph_skip_dynamic_graphs=True 来跳过对具有动态形状输入的函数进行 CUDAGraph 处理，仅对具有静态输入张量形状的函数进行 CUDAGraph。

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 需要静态输入张量地址和支持 GPU 操作等要求，CUDAGraph Trees 会检查函数是否满足这些条件，并在必要时跳过 CUDAGraph。在这里，我们列出了一些常见的跳过 CUDAGraph 的原因。

输入变异: CUDAGraph Trees 会跳过那些就地修改即时输入的函数。对于就地修改参数、缓冲区，或来自 CUDAGraph Tree 管理的函数的输出张量仍然提供支持。请参阅 输入变异支持 部分以获取更多详细信息。
CPU 操作符：包含 CPU 操作符的函数会被跳过。请将函数拆分成多个部分，仅对完全由 GPU 操作符组成的函数应用 CUDAGraph 树。
多设备操作符: 如果一个函数包含多个设备上的操作符，则该函数将被跳过。目前，CUDAGraph 是针对每个设备单独应用的。请使用支持跨设备通信的库，如 NCCL。更多详情，请参见 NCCL 支持 部分。
未支持的符号: 通常在动态形状期间会出现未支持的符号。当前，CUDAGraph Trees 会为每个唯一的输入张量形状记录一个 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 图固定了内存地址，因此它没有很好的方法来处理来自前一次调用的活动张量。

假设我们现在使用以下代码来进行推理运行的性能测试：

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图实现中，第一次调用的输出会被第二次调用覆盖。在CUDAGraph树中，我们不希望在不同迭代之间引入意外依赖关系，导致无法命中热路径，也不希望过早释放前一次调用的内存。我们的启发式方法是，在推理时每次调用都开始一个新的迭代（对于torch.compile），而在训练时只要没有未被调用的pending backward就继续这样做。如果这些启发式方法不正确，你可以使用torch.compiler.mark_step_begin() 标记新的迭代开始，或者在开始下一次运行之前（在torch.compile之外）克隆前一个迭代的张量。

比较

安全隐患	分离Cuda图	CUDAGraph 树
内存可以扩展	在每次图表编译（例如新的尺寸等）时	如果你也在使用非cudagraph内存
录音	在每次调用新的图时	当你在程序中采用新的独特路径时，将会进行重新录制
安全隐患	一个图的调用会覆盖之前的调用结果	无法在模型的不同运行之间保存内存，例如在一次训练循环或一次推理过程中