CUDAGraph 樹¶

背景¶

CUDAGraph¶

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

CUDA 圖（於 CUDA 10 首次亮相）允許將一系列 CUDA 核函式定義和封裝為一個單一單元，即一個操作圖，而不是一系列單獨啟動的操作。它提供了一種機制，可以透過一個 CPU 操作來啟動多個 GPU 操作，從而減少啟動開銷。

CUDA 圖可以提供顯著的加速，特別是對於 CPU 開銷高或計算量小的模型。但也存在一些限制，例如要求使用相同的引數、依賴關係和記憶體地址來執行相同的核函式。

不支援控制流
觸發主機到裝置同步（例如 .item()）的核函式會出錯
核函式的所有輸入引數都固定為其錄製時的值
CUDA 記憶體地址是固定的，但這些地址上的記憶體值可以改變
不支援基本的 CPU 操作或 CPU 副作用

PyTorch CUDAGraph 整合¶

PyTorch 提供了對 CUDAGraphs 的便捷封裝，用於處理與 PyTorch 快取分配器的一些棘手互動。

CachingAllocator 為所有新分配使用一個獨立的記憶體池。在 CUDAGraph 錄製期間，記憶體的記賬、分配和釋放與即時執行（eager run）期間完全相同。在重放時，僅呼叫核函式，並且分配器沒有變化。初始錄製之後，分配器不知道使用者程式中哪些記憶體正在被積極使用。

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

生成圖形化可呼叫物件¶

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

TorchDynamo 先前 CUDA 圖整合¶

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

CUDAGraph Trees 整合¶

與圖形化可呼叫物件類似，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 也會回退到 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 記憶體地址。

仔細觀察輸入變異函式，會發現有三種類型的輸入

來自 eager 的輸入：我們假設這些張量的輸入地址在每次執行時都會變化。由於 cudagraphs 會凍結記憶體地址，因此在圖錄制和執行之前，我們需要將這些輸入複製到一個靜態地址張量。
引數和緩衝區：我們假設（並透過執行時檢查）這些張量在每次執行時具有相同的張量地址。我們不需要複製其內容，因為錄製的記憶體地址與執行的記憶體地址相同。
作為 CUDAGraph Trees 先前輸出的張量：由於 cudagraph 的輸出張量地址是固定的，如果我們執行 CUDAGraph1，然後執行 CUDAGraph2，則從 CUDAGraph1 進入 CUDAGraph2 的輸入將具有固定的記憶體地址。這些輸入與引數和緩衝區一樣，不需要複製到靜態地址張量。我們在執行時檢查以確保這些輸入是穩定的，如果不是，我們將重新錄製。

CUDAGraph Trees 支援對引數和緩衝區以及作為 CUDAGraph Trees 先前輸出的張量進行輸入變異。對於來自 eager 的輸入變異，CUDAGraph Trees 將在不使用 CUDAGraph 的情況下執行函式，併發出 Skipping due to mutated inputs（因變異輸入而跳過）日誌。以下示例展示了 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。這導致單個 inductor 圖對應多個 CUDAGraph。當形狀有限時（例如，推理中的批次大小），重新錄製 CUDAGraph 是有利的。然而，如果輸入張量形狀頻繁變化，甚至在每次呼叫時都變化，則重新錄製 CUDAGraph 可能不划算。在 CUDA 12.4 和驅動版本 550+ 之前，Nvidia 在 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 跳過原地變異 eager 輸入的函式。仍然支援原地變異引數和緩衝區，或由 CUDAGraph Tree 管理的函式輸出的張量。有關更多詳細資訊，請參閱“輸入變異支援”部分。
CPU 運算元：跳過包含 CPU 運算元的函式。請將函式拆分成多個函式，並僅對包含 GPU 運算元的函式應用 CUDAGraph Trees。
多裝置運算元：如果函式包含跨多個裝置的運算元，則會跳過。目前，CUDAGraph 是基於每個裝置應用的。請使用 NCCL 等支援的庫進行跨裝置通訊。有關更多詳細資訊，請參閱“NCCL 支援”部分。
未支援的自由符號：未支援的自由符號通常在動態形狀期間發生。CUDAGraph Trees 當前會為每個唯一的輸入張量形狀錄製一個 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

限制¶

由於 CUDA Graph 固定記憶體地址，CUDA Graph 沒有很好的方法處理來自先前呼叫的活動張量。

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

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 記憶體
錄製	在任何新的圖呼叫時	將在你透過程式採用的任何新的、唯一的路徑上重新錄製
注意事項	呼叫一個圖將覆蓋先前的呼叫	無法在模型的獨立執行之間持久化記憶體 - 例如一個訓練迴圈或一次推理執行