CUDAGraph 樹#

建立時間：2023 年 5 月 19 日 | 最後更新時間：2025 年 7 月 30 日

背景#

CUDAGraph#

有關 CUDAGraph 的更詳細背景資訊，請閱讀使用 CUDAGraph 加速 PyTorch。

CUDA Graphs 於 CUDA 10 首次亮相，它允許將一系列 CUDA 核心定義並封裝為一個單元，即操作圖，而不是一系列單獨啟動的操作。它提供了一種機制，透過一次 CPU 操作即可啟動多個 GPU 操作，從而降低了啟動開銷。

CUDA Graphs 可以帶來顯著的加速，尤其是對於 CPU 開銷高或計算量小的模型。由於要求相同的核心以相同的引數和依賴項以及記憶體地址執行，因此存在一些限制。

無法進行控制流
觸發主機到裝置同步的核心（例如 `.item()`）會報錯
核心的所有輸入引數都固定為其記錄時的值
CUDA 記憶體地址是固定的，但這些地址處的記憶體值可以更改
不包含核心 CPU 操作或 CPU 端副作用

PyTorch CUDAGraph 整合#

PyTorch 提供了 CUDAGraph 的一個方便的包裝器，該包裝器處理了與 PyTorch 快取分配器的一些棘手互動。

快取分配器使用獨立的記憶體池來處理所有新分配。在 CUDAGraph 記錄期間，記憶體的核算、分配和釋放與在即時執行期間完全相同。在回放時，只調用核心，分配器沒有任何變化。在初始記錄之後，分配器不知道哪些記憶體正在使用者程式中被主動使用。

在即時分配和 cudagraph 分配之間使用獨立的記憶體池，如果兩者都有大量記憶體分配，可能會增加程式的記憶體佔用。

建立圖表可呼叫物件#

Make Graphed Callables 是一個 PyTorch 抽象，用於在一系列可呼叫物件之間共享單個記憶體池。Graphed Callables 利用了 CUDA Graph 記錄期間快取分配器精確核算記憶體的事實，以安全地在獨立的 CUDA Graph 記錄之間共享記憶體。在每次呼叫中，輸出都會被保留為活動記憶體，防止一個可呼叫物件覆蓋另一個的可活動記憶體。Graphed Callables 只能按單一順序呼叫；第一個執行的記憶體地址將被燒錄到第二個，依此類推。

TorchDynamo 之前的 CUDA Graphs 整合#

使用 `cudagraph_trees=False` 執行不會在獨立的圖捕獲之間重用記憶體，這可能導致記憶體迴歸。即使是對於沒有圖中斷的模型，這也會有問題。前向和後向是獨立的圖捕獲，因此前向和後向的記憶體池不共享。特別是，在前向中儲存的啟用的記憶體無法在後向中回收。

CUDAGraph 樹整合#

與 Graphed 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 記錄的 tape（在此例項中是 1 -> 2 -> 4）來在獨立的記錄之間共享單個記憶體池中的所有記憶體。我們添加了不變數來確保記憶體始終位於記錄時的相同位置，並且使用者程式中沒有可能被覆蓋的活動張量。

CUDA Graphs 的相同約束適用：必須使用相同的引數（靜態大小、地址等）呼叫相同的核心
在記錄和回放之間必須觀察到相同的記憶體模式：如果在記錄期間一個圖的張量輸出在另一個圖之後死亡，那麼在回放期間也必須如此。
CUDA 池中的活動記憶體會強制兩個記錄之間產生依賴關係
這些記錄只能按單一順序呼叫 1 -> 2 -> 4

所有記憶體都共享在一個記憶體池中，因此與即時執行相比沒有額外的記憶體開銷。那麼，如果我們遇到一條新路徑並執行圖 3 會發生什麼？

圖 1 被回放，然後我們遇到圖 3，這之前我們還沒有記錄。在圖回放時，私有記憶體池不會更新，因此 y 不會反映在分配器中。如果不小心，我們會覆蓋它。為了支援在回放其他圖後重用相同的記憶體池，我們將記憶體池回滾到圖 1 結束時的狀態。現在我們的活動張量已反映在快取分配器中，我們可以安全地執行新圖了。

首先，我們將命中最佳化過的、已在圖 1 中記錄的 CUDAGraph.replay() 路徑。然後我們將命中圖 3。和以前一樣，我們需要在記錄之前預熱一次圖。在預熱執行中，記憶體地址不是固定的，因此圖 4 也將回退到 inductor 的非 cudagraph 呼叫。

第二次遇到圖 3 時，我們已經預熱並準備好記錄。我們記錄圖 3，然後再次記錄圖 4，因為輸入記憶體地址已更改。這建立了一個 CUDA Graph 記錄的樹。一個 CUDA Graph Tree！

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

輸入變異支援#

輸入變異函式是指在輸入張量上進行原地寫入的函式，如下例所示

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

輸入變異函式通常給 CUDAGraph Trees 帶來挑戰。由於 CUDAGraph 需要靜態 CUDA 記憶體地址，對於每個輸入張量 x，CUDAGraph Trees 可能會分配一個靜態記憶體地址 x'。在執行期間，CUDAGraph Trees 首先將輸入張量 x 複製到靜態記憶體地址 x'，然後回放記錄的 CUDAGraph。對於輸入變異函式，x' 是原地更新的，這不會反映在輸入張量 x 上，因為 x 和 x' 位於不同的 CUDA 記憶體地址。

仔細檢視輸入變異函式會發現有三種類型的輸入

來自即時執行的輸入：我們假設這些張量在每次執行時輸入張量地址都會發生變化。由於 cudagraph 會凍結記憶體地址，因此我們需要在圖記錄和執行之前將這些輸入複製到一個靜態地址張量。
引數和緩衝區：我們假設（並在執行時檢查）這些張量在每次執行時都有相同的張量地址。我們不需要複製它們的內容，因為記錄的記憶體地址將與執行的記憶體地址相同。
來自 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())

要為修改來自即時執行輸入的函式啟用 CUDAGraph Trees，請重寫該函式以避免輸入變異。

注意
透過設定 torch._inductor.config.cudagraph_support_input_mutation = True 來為“減少開銷”模式啟用輸入變異支援。

動態形狀支援#

動態形狀意味著輸入張量在函式呼叫之間具有不同的形狀。由於 CUDAGraph 需要固定的張量地址，CUDAGraph Trees 會為輸入張量的每種唯一形狀重新記錄 CUDAGraph。這會導致一個 inductor graph 有多個 CUDAGraph。當形狀有限時（例如，推理中的批處理大小），重新記錄 CUDAGraph 是有利的。然而，如果輸入張量形狀頻繁更改，甚至在每次呼叫時都更改，重新記錄 CUDAGraph 可能不划算。Nvidia 在 CUDA 12.4 和驅動程式版本 550+ 之前，每次啟動 CUDAGraph 時使用 64 KB 裝置記憶體。對於許多 CUDAGraph 的重新記錄，這種記憶體成本可能非常可觀。

對於輸入張量形狀頻繁更改的函式，我們建議將輸入張量填充到少數固定張量形狀，以仍然享受 CUDAGraph 的好處。此外，設定 torch._inductor.config.triton.cudagraph_skip_dynamic_graphs=True 可以跳過具有動態形狀輸入的函式進行 cudagraphing，只對具有靜態輸入張量形狀的函式進行 cudagraphing。

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 有靜態輸入張量地址的要求，並且不支援 CPU 運算子，CUDAGraph Trees 會檢查函式是否滿足這些要求，並在必要時跳過 CUDAGraph。在此，我們列出了跳過 CUDAGraph 的常見原因。

輸入變異：CUDAGraph Trees 會跳過原地變異即時輸入的函式。原地變異引數和緩衝區，或來自 CUDAGraph Tree 管理函式的輸出張量仍然受支援。請參閱“輸入變異支援”部分了解更多詳細資訊。
CPU 運算子：包含 CPU 運算子的函式將被跳過。請將函式拆分為多個函式，並將 CUDAGraph Trees 應用於僅包含 GPU 運算子的函式。
多裝置運算子：如果函式包含多個裝置上的運算子，則會跳過該函式。目前，CUDAGraph 是按裝置應用的。請使用 NCCL 等支援的庫進行跨裝置通訊。請參閱“NCCL 支援”部分了解更多詳細資訊。
釋放未支援的符號：釋放未支援的符號通常發生在動態形狀期間。CUDAGraph Trees 目前為每種唯一的輸入張量形狀記錄一個 CUDAGraph。請參閱“動態形狀支援”部分了解更多詳細資訊。
CUDAGraph 不安全的自定義運算子：某些自定義運算子可能包含 cudagraph 不安全的運算子，這會導致 cudagraph 被跳過。請參閱“CUDAGraph 不安全的自定義運算子”部分了解更多詳細資訊。
不相容的運算子：如果函式包含不相容的運算子，CUDAGraph Trees 將跳過該函式。請將函式中的這些運算子替換為支援的運算子。我們顯示了不相容運算子的詳盡列表

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

CUDAGraph 不安全的自定義運算子#

預設情況下，自定義運算子被假定為對 CUDAGraph 是安全的。但是，某些自定義運算子可能包含不支援的運算子，例如 CPU 運算子。由於編譯器將自定義運算子視為黑盒，使用者必須透過設定 `torch._C.Tag.cudagraph_unsafe` 標籤顯式將這些運算子標記為對 CUDAGraph 不安全，如下例所示。當函式包含 cudagraph 不安全的自定義運算子時，除非啟用了“CUDAGraph 分割槽”，否則 CUDAGraph 將跳過它。

@torch.library.custom_op(
    "mylib::modify",
    mutates_args=(),
    tags=(torch._C.Tag.cudagraph_unsafe,),
)
def modify(pic: torch.Tensor) -> torch.Tensor:
    pic1 = pic + 1
    pic1_cpu = (pic1.cpu() + 1) * 2
    return pic1_cpu.cuda() + pic

@modify.register_fake
def _(pic):
    return torch.empty_like(pic)

CUDAGraph 分割槽#

如前所述，CUDAGraph 不支援某些運算子（例如 CPU 運算子），這可能會限制其採用。CUDAGraph 分割槽是一種編譯器解決方案，它自動分離這些運算子，重新排序運算子以減少分割槽數量，並單獨將 CUDAGraph 應用於每個分割槽。請設定 `torch._inductor.config.graph_partition=True` 來啟用 CUDAGraph 分割槽。

考慮以下示例，其中 `x` 和 `y` 是 GPU 輸入，但 `y_cpu` 是 CPU 張量。沒有圖分割槽，此函式必須因 CPU 運算子而被跳過。透過圖分割槽，CPU 運算子被分離出來，其餘的 GPU 運算子被 cudagraphified，從而產生兩個獨立的 CUDAGraph。

def f(x, y):
    x1 = x + 1
    y1 = y + 1
    y_cpu = y1.cpu() + 1
    z = x @ y
    return x1 + y1 + z + y_cpu.cuda()

目前，CUDAGraph 分割槽支援分離以下型別的運算子

非 GPU 運算子：常見的例子包括 CPU 張量上的計算。
裝置複製運算子：裝置之間的資料傳輸，例如上面示例中的 `y1.cpu()`。
控制流運算子：控制流運算子已被分離，因為 CUDAGraph 尚不支援它們。
CUDAGraph 不安全的自定義運算子：標記有 `torch._C.Tag.cudagraph_unsafe` 的自定義運算子已被分離。請參閱“CUDAGraph 不安全的自定義運算子”部分了解詳細資訊。
未支援的 Symints：請參閱“動態形狀支援”部分了解更多資訊。

限制#

由於 CUDA Graph 固定了記憶體地址，CUDA Graphs 在處理前一次呼叫的活動張量方面沒有很好的方法。

假設我們正在使用以下程式碼執行推理進行基準測試

import torch

@torch.compile(mode="reduce-overhead")
def my_model(x):
    y = torch.matmul(x, x)
    return y

x = torch.randn(10, 10, device="cuda")
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 Graph 實現中，第一次呼叫的輸出將被第二次呼叫覆蓋。在 CUDAGraph Trees 中，我們不想在迭代之間新增意外的依賴關係，導致我們無法命中熱路徑，也不想過早地釋放先前呼叫的記憶體。我們的啟發式方法是在推理中，我們為 torch.compile 在每次呼叫時啟動一個新迭代，在訓練中，只要沒有待處理的後向尚未呼叫，我們也這樣做。如果這些啟發式方法不正確，您可以使用 `torch.compiler.mark_step_begin()` 標記新迭代的開始，或者在開始下一次執行之前克隆前一次迭代的張量（在 torch.compile 之外）。

比較#

易出錯的陷阱	單獨的 CudaGraph	CUDAGraph 樹
記憶體可能增加	每次圖編譯時（新大小等）	如果您也執行非 cudagraph 記憶體
記錄	每次呼叫圖時	將在您程式的任何新、唯一的路徑上重新記錄
易出錯的陷阱	呼叫一個圖會覆蓋前一個呼叫	無法在模型的獨立執行之間持久化記憶體 - 一個訓練迴圈訓練，或一次推理執行