Table of Contents
Shortcuts

CUDAGraph 树

CUDAGraph 背景

关于CUDAGraphs的更详细背景，请阅读使用CUDAGraphs加速PyTorch。

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

CUDA 图可以带来显著的速度提升，特别是在具有高 CPU 开销或小计算量的模型中。由于要求相同的内核以相同的参数和依赖关系运行，并且内存地址相同，因此存在一些限制。

  • 控制流程是不可能的

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

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

  • CUDA 内存地址是固定的，然而这些地址上的内存值可以改变

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

PyTorch CUDAGraph 集成

PyTorch 提供了一个 便捷的封装 用于 CUDAGraphs，处理了与 PyTorch 缓存分配器的几个棘手交互。

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

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

制作图形化可调用对象

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):
    # 图 1
    y = x * x * x
    # 图断点触发于此
    if y.sum() > 0:
        # 图 2
        z = y ** y
    else:
        # 图 3
        z = (y.abs() ** y.abs())
    torch._dynamo.graph_break()
    # 图 4
    return z * torch.rand_like(z)

# 第一次运行预热每个图，执行诸如CuBlas或Triton基准测试等操作
foo(torch.arange(0, 10, device="cuda"))
# 第二次运行进行CUDA图记录，并重放它
foo(torch.arange(0, 10, device="cuda"))
# 最后我们进入优化后的CUDA图重放路径
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相比没有额外的内存开销。现在，如果我们遇到一个新的路径并运行Graph 3会发生什么？

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

首先，我们会触发我们已经记录在图1中的优化路径，CUDAGraph.replay()。然后我们会触发图3。和之前一样，我们需要在记录之前预热图一次。在预热运行中，内存地址是不固定的，因此图4也会回退到inductor，非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)
# 运行时错误：错误：访问已被后续运行覆盖的CUDAGraphs的张量输出。

在单独的CUDA图实现中，第一次调用的输出将被第二次调用覆盖。在CUDA图树中，我们不希望在迭代之间添加意外的依赖关系，这会导致我们无法命中热路径，也不希望过早地从先前的调用中释放内存。我们的启发式方法是在推理中，我们为每个调用开始一个新的迭代，对于torch.compile也是如此，在训练中，只要没有未调用的待处理反向传播，我们也会这样做。如果这些启发式方法不正确，您可以使用torch.compiler.mark_step_begin()标记新迭代的开始，或者在开始下一次运行之前（在torch.compile之外）克隆先前迭代的张量。

比较

陷阱

分离 CudaGraph

CUDAGraph 树

内存可以增加

在每次图形编译时（新尺寸等）

如果你也在运行非cudagraph内存

录音

在任何新的图表调用时

将在您通过程序的任何新、唯一路径上重新录制

陷阱

调用一个图将会覆盖之前的调用

无法在模型的不同运行之间保持内存 - 一次训练循环训练，或一次推理运行
优云智算