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 记录期间，内存会像在急切运行期间一样被精确地计算、分配和释放。在重播时，仅调用内核，分配器不会发生任何变化。在初始记录之后，分配器不知道用户程序中哪些内存正在被积极使用。

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

创建图形化可调用对象¶

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

TorchDynamo 以前的 CUDA Graphs 集成¶

使用 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 Graphs 的相同约束适用：必须使用相同的参数（静态大小、地址等）调用相同的内核
记录和重播之间必须观察到相同的内存模式：如果一个图形的张量输出在记录期间在另一个图形之后消失，则它在重播期间也必须消失。
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 之外）。

比较¶

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