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 抽象，用于在多个可调用对象之间共享一个内存池。Graphed Callables 利用了这样一个事实：在 CUDA 图形记录中，内存由缓存分配器精确计算，以便在单独的 CUDA 图形记录之间安全地共享内存。在每次调用中，输出都保留为活动内存，防止一个可调用对象覆盖另一个可调用对象的活动内存。Graphed Callables 只能按单个顺序调用；第一次运行的内存地址会刻录到第二次运行中，依此类推。

TorchDynamo 之前的 CUDA 图形集成¶

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

CUDAGraph 树集成¶

与 Graph Callables 类似，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 也会回退到感应器，非 cudagraph 调用。

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

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

限制¶

由于 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 之外）。

比较¶

Footguns	单独的 CudaGraph	CUDAGraph 树
内存可能会增加	在每次图形编译（新大小等）时	如果你也在运行非 cudagraph 内存
记录	在图形的任何新调用上	将在你通过程序采取的任何新的、唯一的路径上重新记录
Footguns	一个图形的调用将覆盖先前的调用	无法在通过模型的单独运行之间保留内存 - 一个训练循环训练或一次推理运行