CUDAGraph 樹狀結構¶

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 配置都配置了大量的記憶體，那麼在它們之間使用獨立的記憶體池可能會增加程式的記憶體使用量。

建立圖形化可呼叫物件¶

建立圖形化可呼叫物件是一種 PyTorch 抽象概念，用於在一系列可呼叫物件之間共用單一記憶體池。圖形化可呼叫物件利用了在 CUDA Graph 記錄時，快取配置器會精確地計算記憶體，以便在獨立的 CUDA Graph 記錄之間安全地共用記憶體。在每次叫用中，輸出都會保留為活動記憶體，防止一個可呼叫物件覆寫另一個可呼叫物件的活動記憶體。圖形化可呼叫物件只能以單一順序叫用；第一次執行的記憶體位址會燒錄到第二次執行，依此類推。

TorchDynamo 先前的 CUDA Graphs 整合¶

使用 cudagraph_trees=False 執行時，不會在獨立的圖形擷取之間重複使用記憶體，這可能會導致大量的記憶體退化。即使對於沒有圖形中斷的模型，這也會產生問題。正向和反向是獨立的圖形擷取，因此正向和反向的記憶體池不會共用。特別是，在正向中儲存的活動記憶體無法在反向中回收。

CUDAGraph 樹狀結構整合¶

與圖形化可呼叫物件一樣，CUDA Graph 樹狀結構在所有圖形擷取中使用單一記憶體池。但是，CUDA Graph 樹狀結構不會要求單一叫用序列，而是建立獨立的 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 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 內存
錄製	在圖形的任何新調用時	將在您通過程序採取的任何新的唯一路徑上重新記錄
陷阱	調用一個圖形將覆蓋先前的調用	無法在通過模型的單獨運行之間持久化內存 - 一個訓練循環訓練或一次推理運行