Dynamo 深入探討

TorchDynamo（或簡稱 Dynamo）是 torch.compile 中的追蹤器，而且它通常是造成那些瘋狂回溯的原因。然而，我們不能盲目地責怪 Dynamo 造成這些錯誤。為了向用戶提供其靈活性，Dynamo 被賦予了理解任何 Python 程式的艱鉅任務。特別是，Dynamo 必須在內部實現 Python 程式語言的很大一部分！

在這篇文章中，我們將從頭開始探討 Dynamo 的內部設計。我們將討論它提供的功能以及其實現方式。到本文結束時，您將更好地理解當您使用 torch.compiled 編譯 PyTorch 程式時出現錯誤或編譯成功但加速效果不佳時，究竟是哪裡出了錯。

Dynamo 簡介

在深入研究所有實現細節之前，讓我們先討論 Dynamo 的作用。

Dynamo 是一個追蹤器。這表示，給定一個函數及其輸入，它會執行該函數並將線性指令序列（沒有控制流程）記錄到圖形中。例如，請考慮以下程式

import torch

@torch.compile
def mse(x, y):
    z = (x - y) ** 2
    return z.sum()

x = torch.randn(200)
y = torch.randn(200)
mse(x, y)

如果我們將此程式儲存到檔案 example.py 中，然後執行

TORCH_LOGS=graph_code python example.py

我們會看到 Dynamo 追蹤的輸出

def forward(l_x_: torch.Tensor, l_y_: torch.Tensor):
    # File: example.py:5, code: z = (x - y) ** 2
    sub = l_x_ - l_y_
    z = sub ** 2
    # File: example.py:6, code: return z.sum()
    sum_1 = z.sum()
    return (sum_1,)

我們將其稱為給定輸入的函數圖形（或追蹤）。這是透過 FX 圖形 表示的。我們將簡單地將 FX 圖形視為儲存函數呼叫清單的容器。

我們首先要注意的是，該圖形是 PyTorch 運算的線性序列。1 Dynamo 會記錄所有 PyTorch 運算並按順序儲存它們。例如，它將 z = (x - y) ** 2 分割成其兩個組成運算，sub = l_x_ - l_y_ 和 z = sub ** 2。

當我們說追蹤是線性的時，我們的意思是沒有分支或任何控制流程。要了解這一點，請考慮

import torch

@torch.compile
def fn(x, n):
    y = x ** 2
    if n >= 0:
        return (n + 1) * y
    else:
        return y / n

x = torch.randn(200)
fn(x, 2)

使用 TORCH_LOGS=graph_code 執行時，會返回

def forward(l_x_: torch.Tensor):
    # File: example.py:5, code: y = x ** 2
    y = l_x_ ** 2
    # File: example.py:7, code: return (n + 1) * y
    mul = 3 * y
    return (mul,)

我們看到 Dynamo 從追蹤中完全刪除了 if 語句，並且只記錄了使用輸入執行的運算。

因此，應該清楚的是，函數的追蹤取決於輸入。特別是，這意味著追蹤不是在我們編寫 @torch.compile 時生成的，而是在我們使用實際參數執行 fn(x, 2) 時生成的。

這裡要注意的另一件有趣的事情是，Dynamo 刪除了函數的第二個參數。相反，它將其視為常數，並在圖形中記錄了運算 n + 1 的結果。這是 Dynamo 的另一個特性：Dynamo 會將任何非張量值視為常數……除了整數之外。現在讓我們看看整數有什麼特別之處。

Dynamo 的最後一個定義屬性是它知道如何處理動態形狀。符號形狀是指 Dynamo 追蹤形狀的能力，更一般地說，是追蹤整數，而不是將它們保留為常數。這樣可以避免重新編譯，並部署適用於生產環境中任何大小的通用模型。動態形狀出現的主要例子是批次大小，我們可能會使用固定的批次大小訓練模型，但隨後對任意批次大小執行推論，或者在處理文字或音訊時遇到的變數序列長度。

我們可以透過多次執行上面的範例來看到這一點

import torch

@torch.compile
def fn(x, n):
    y = x ** 2
    if n >= 0:
        return (n + 1) * y
    else:
        return y / n

x = torch.randn(200)
fn(x, 2)
fn(x, 3)
fn(x, -2)

在這種情況下，TORCH_LOGS=graph_code 會生成另外兩個圖形

# Graph for n==2 omitted

def forward(self, l_x_: torch.Tensor, l_n_: torch.SymInt):
    # File: a.py:5, code: y = x ** 2
    y = l_x_ ** 2

    # File: a.py:7, code: return (n + 1) * y
    add = l_n_ + 1
    mul = add * y
    return (mul,)

def forward(self, l_x_: torch.Tensor, l_n_: torch.SymInt):
    # File: a.py:5, code: y = x ** 2
    y = l_x_ ** 2

    # File: a.py:9, code: return y / n
    truediv = y / l_n_
    return (truediv,)

Dynamo 檢測到一個整數在第一次呼叫後更改了其值，並開始追蹤它。我們看到這些圖形是通用的，並且透過 SymInt 類型的物件符號化地追蹤變數 n。

如果在這些呼叫之後我們呼叫 fn(x, 4)，Dynamo 不會重新編譯，而是重複使用已經追蹤的圖形。

總之：1. Dynamo 是一個 Python 追蹤器 2. 給定一些輸入，它會返回一個包含已執行 PyTorch 函數的 FX 圖形 3. 如果它檢測到整數在呼叫之間發生了變化，它也可以追蹤它們 4. 它會特殊化任何不是張量或純量的值

當然，Dynamo 還可以做很多其他事情，比如弄清楚何時需要重新追蹤、重寫函數的位元組碼、實現圖形斷點……為了簡潔起見，我們將在後續章節中逐步討論所有這些內容。

PEP 523：向 CPython 添加框架評估 API

假設現在我們要實作 Dynamo，我們會從哪裡開始呢？對我們來說相當方便的是，PEP 523 與 Python 3.6 一起發布了。這個 PEP 旨在允許第三方為 Python 建立 JIT 編譯器。讓我們來看看怎麼做。

關於 CPython 的注意事項：CPython 內部實作為一個 堆疊機器。Python 程式會被編譯成 位元組碼，然後由這個直譯器執行。要瞭解更多關於這些位元組碼的資訊，請參閱標準庫中的 dis 模組。另請參閱 開發人員文件，以獲得 CPython 直譯器的介紹。我們假設讀者熟悉堆疊機器的概念。

PEP 523 公開了一個 API，使用者可以在其中新增自訂的每個函式的直譯器。然後，CPython 將使用這個直譯器而不是它自己的直譯器來執行函式。為了能夠執行函式，在進入時，CPython 會為自訂直譯器提供以下內容：- 函式的位元組碼 - 函式參數的值（即局部變數）及其名稱 - 全域變數的值及其名稱 - 內建函式，例如 abs 或 print

您可以在 這裡 看到所有欄位。2

總之，CPython 為使用者的直譯器提供了執行函式所需的所有資訊。3

透過這個 API，我們可以透過實作一個直譯器來實作一個追蹤器，該直譯器執行程式碼並在圖形中記錄在此執行期間發生的所有 PyTorch 操作。這正是 Dynamo 所做的。

Dynamo 使用這個 CPython API 來解析所有這些物件，並將它們打包到 一個 Python 結構 中。完成之後... 它會從 C 返回 Python。除了這段與 CPython 溝通的程式碼之外，Dynamo 完全是用 Python 實作的。

應該很清楚，裝飾器 @torch.compile 的工作是在呼叫函式時安裝必要的腳手架，將位元組碼、參數、全域變數等傳遞給 Dynamo。同樣，@torch.compile 並沒有真正編譯任何東西。

用 Python 實作 CPython

所以，我們回到了 Python 的世界。我們有函式的位元組碼，以及執行它所需的所有上下文。特別是，我們已經到達了 _convert_frame_assert。這是裝飾器 torch.compile 返回的函式！我們從 _dynamo.optimize 來到這個函式。裝飾器 torch.compile 只是 _dynamo.optimize 的一個漂亮的 API。

在開始實作 Python 直譯器之前，我們想要定義一個 IR。特別是，我們希望將所有局部和全域變數包裝在我們自己的內部類別中。這允許我們更好地追蹤這些物件，並將在 Dynamo 看來可以以相同方式處理的物件組合在一起。

內部類別結構的父類別是 VariableTracker，表示 Dynamo 理解的不同物件。例如，ListVariable 表示一個 list 物件，並在內部保留一個 VariableTrackers 列表。另一個 VariableTracker 的例子是 ConstantVariable。ConstantVariable 包裝了所有 被 Dynamo 視為常數的物件。我們還針對需要特別注意的物件提供特殊的子類別，例如 TensorVariable。所有這些內部類別都在 torch/_dynamo/variables 資料夾中定義。

Python 物件會在 VariableBuilder._wrap 中包裝到它們對應的 VariableTracker 類別中。這個函式只是一個非常長的 elif 鏈，它試圖遞迴地將 Python 輸入與適當類型的 VariableTracker 進行模式匹配。

除錯技巧。當我們從 Dynamo 中獲得意外結果時，有時是由於建構器造成的。如果建構器的邏輯錯誤，有時 Dynamo 可能會將變數包裝在錯誤的 VariableTracker 類型中，這可能會在稍後階段導致問題。當您遇到 Dynamo 錯誤時，查看錯誤中出現的 VariableTracker 類型以及引發異常的 VariableTracker 方法相當有用。特別是，有時我們會發現一個物件被追蹤為 UserDefinedObjectVariable（這是 Dynamo 的萬用類別），而它應該被追蹤為更具體的物件。在這些情況下，SourceBuilder.__call__ 邏輯通常是罪魁禍首。

除錯技巧。當使用 TORCH_LOGS=dynamo 執行程式時，其中一個列印出來的工件是以下形式的行：

TRACE LOAD_GLOBAL y [TorchInGraphFunctionVariable(<built-in method any>), TensorVariable()]

這是原始程式的位元組碼和當時堆疊的狀態。這對於找出哪個物件沒有被追蹤到正確的 VariableTracker 中非常有用。

好的，所以我們為我們的追蹤器準備了一個 IR，現在我們*只需要*重新實作 CPython 的堆疊機器。這在 symbolic_convert.py 中由 InstructorTranslatorBase 實作。

InstructionTranslatorBase 有大約 200 個方法，實作了幾乎所有的 Python 位元組碼。舉例來說，我們可以看到 BUILD_LIST 的實作

def BUILD_LIST(self, inst):
    items = self.popn(inst.argval)
    self.push(ListVariable(items, mutable_local=MutableLocal()))

這是由 l = [2, 3, 4] 之類的建構產生的位元組碼。在這種情況下，由於有三個元素，所以產生的位元組碼是 BUILD_LIST 3。這表示我們將堆疊頂部的 3 個元素彈出，並將由這三個元素形成的新列表物件推入堆疊頂部。

產生輸出圖

有了一種符號化執行 Python 程式碼的方法，我們就可以提取在給定一些輸入的情況下，在程式的符號化執行過程中發生的 PyTorch 操作。這在 Dynamo 中是透過 OutputGraph 物件實作的。OutputGraph 物件是 綁定到 `InstructionTranslator 物件的，它會追蹤建立 Dynamo 將返回的 FX 圖所需的所有資料。

FX 圖的所有輸入和中間元素都是 fx.Node。在 Dynamo 中，fx.Node 被包裝在 fx.Proxy 中。fx.Proxy 用於建構 FX 圖。特別是，它們會將對它們執行的每個 PyTorch 操作記錄到圖中。您可以透過呼叫 create_proxy 來建立要新增到圖中的新操作。然後，我們可以透過 wrap_fx_proxy 函式將其新增到圖中。

圖形會儲存張量上的操作... 以及符號整數上的操作。我們將在稍後討論符號整數，但首先我們將討論 Dynamo 如何解決一個相當重要的正確性問題。

使 Dynamo 可靠：防護

在這個時候，我們有了一種完全不考慮控制流程來追蹤程式的方法。為此，我們重新實作了整個 CPython... 如果這聽起來有點過頭，那是因為確實如此。torch.jit.trace 已經在沒有所有這些機制的情況下實作了這一點，所以這是怎麼回事？

正如其文件中所警告的，torch.jit.trace 的問題在於，只有在被追蹤的程式不依賴資料的情況下，它才能正常工作。換句話說，只有當程式本身是線性的，它才能正常工作。這意味著在編寫程式時不使用 if-else、for-while 迴圈、異常。更重要的是，我們使用的任何函式庫都不能使用任何控制流程！總之，在像 Python 這樣動態的語言中不使用控制流程實際上是一個巨大的限制。

JAX 透過始終重新追蹤並在重新追蹤後快取圖形來解決這個問題。另一方面，Dynamo 使用防護來避免每次都重新追蹤整個程式。

**防護**是一種為了針對一組範例輸入專門化框架而做出的假設（關於輸入的布林運算式）。只有當這些假設在新輸入上成立時，才能重新使用圖形。

例如，任何函式的常數輸入（例如字串）都會安裝一個防護，說明該輸入的類型應為 str 並且等於我們傳遞的字串。執行

import torch

@torch.compile
def fn(a, b):
    return a * len(b)

fn(torch.arange(10), "Hello")

使用 TORCH_LOGS=guards 會列印（以及其他防護）

___check_type_id(L['b'], 94334122025024)
L['b'] == 'Hello'

這段話的意思是「局部變數 b 應該具有特定的型別（在本例中為 str，由常數 9433... 表示），並且其值應為 'Hello'」。如果我們接著以不同的引數再次執行該函數：

import torch

@torch.compile
def fn(a, b):
    return a * len(b)

fn(torch.arange(10), "Hello")
fn(torch.arange(10), "Hi")

我們可以看到執行 TORCH_LOGS=recompiles 後失敗的守衛

Recompiling function fn in script.py:3
triggered by the following guard failure(s):
     - L['b'] == 'Hello'

守衛會在 函數的輸入被封裝在建構器中時 以及 在程式執行期間 累積。我們將在下一節中展示更多守衛的示例，但首先讓我們討論一下來源。

來源 會追蹤如何從進入當前框架時存在的原始局部變數或全域變數重建變數。具體來說，它會追蹤原始的局部和全域物件，以及它們包含的任何物件。在

def foo(x: Tensor, y: List[Tensor]):
    a = x * y[0]
    return a * x

x 和 y 中，它們的來源是 LocalSource，而 y[0] 的來源是 GetItemSource，它在內部存儲了一個 LocalSource。另一方面，a 將沒有來源，因為它是一個僅存在於 fx 圖中的中間變數。

所有這些都在 torch/_dynamo/source.py 中定義。我們可以在以下示例中看到 GetItemSource 生成的守衛

import torch

@torch.compile
def fn(x, l):
    return x * len(l[0])

fn(torch.randn(8), ["Hi", "Hello"])

生成以下守衛

___check_type_id(L['l'], 94439025877664)
len(L['l']) == 2
___check_type_id(L['l'][0], 94439025840192)
L['l'][0] == 'Hi'
___check_type_id(L['l'][1], 94439025840192)
L['l'][1] == 'Hello'

在這裡，我們看到了由 GetItemSource（[0] 和 [1]）生成的程式碼，它封裝了一個 LocalSource（L['l']）。

至此，藉助來源和守衛，我們能夠實現一個快取系統，以避免重新編譯，而無需每次都進行回溯。我們將在後續內容中更詳細地討論這個快取系統。

細心的讀者會注意到，這還沒有解釋為什麼我們需要對 Python 解釋器進行如此精細的控制，以至於必須重新實現它。我們展示的守衛示例取決於輸入物件，因此我們仍然可以在執行函數之前計算這些守衛。換句話說，我們可以在 torch.jit.trace 之上實現這個守衛系統，並以更少的努力獲得相同的功能… 讓我們來看看符號形狀。

符號形狀

我們在引言中討論的另一個觀點是，Dynamo 知道如何追蹤整數。為了實現這一點，我們使用了一個符號類 torch.SymInt，它的作用類似於 int，但它會在輸出 FX 圖中記錄對其執行的所有操作。 4 在介紹符號整數追蹤時，我們已經在引言中看到了這個類。

現在，讓我們討論定義 Dynamo 中符號形狀追蹤的三個特性，以及如何實現它們。

預設為靜態

Dynamo 預設所有整數（無論是輸入還是張量的形狀）都是靜態的。換句話說，在函數的第一次執行中，不會追蹤任何整數。然後，只有當它檢測到整數或形狀在執行過程中更改了值時，它才會追蹤它，並生成一個針對該變數的泛型圖。

我們在引言中已經看到了使用整數時的這種行為。現在讓我們來看一個使用張量形狀的示例。

import torch

@torch.compile
def fn(a, b):
    return a.shape[0] * a * b

fn(torch.randn(4, 3), torch.randn(4, 3))
fn(torch.randn(8, 3), torch.randn(8, 3))

使用 TORCH_LOGS=graph_code 運行此程式，我們可以看到這兩個呼叫被追蹤為

def forward(self, l_a_: torch.Tensor, l_b_: torch.Tensor):
    mul = 4 * l_a_
    mul_1 = mul * l_b_
    return (mul_1,)

def forward(self, s0: torch.SymInt, l_a_: torch.Tensor, l_b_: torch.Tensor):
    size = l_a_.size()
    getitem = size[0]
    mul = getitem * l_a_
    mul_1 = mul * l_b_
    return (mul_1,)

在第一個圖中，形狀被追蹤為常數，但一旦它發生變化，它就會使用 SymInts 符號地追蹤它。通常，查看中間值形狀的更簡單方法是使用 TORCH_LOGS=graph_sizes 運行程式

TRACED GRAPH TENSOR SIZES
===== __compiled_fn_1 =====
l_a_: (s0, 3)
l_a_ (concrete): (8, 3)
l_b_: (s0, 3)
l_b_ (concrete): (8, 3)
mul: (s0, 3)
mul (concrete): (8, 3)
mul_1: (s0, 3)
mul_1 (concrete): (8, 3)

在這裡，我們可以看到兩個張量引數的第一個維度是動態的，因為它由 s0 變數表示。

我們可以通過運行 TORCH_LOGS=guards 來找到 Dynamo 如何實現這一點

# Guards first call
check_tensor(L['a'], torch.float32, device=None, requires_grad=False, size=[4, 3], stride=[3, 1])
check_tensor(L['b'], torch.float32, device=None, requires_grad=False, size=[4, 3], stride=[3, 1])

# Guards second call
check_tensor(L['a'], torch.float32, device=None, requires_grad=False, size=[None, 3], stride=[3, 1])
check_tensor(L['b'], torch.float32, device=None, requires_grad=False, size=[None, 3], stride=[3, 1])

L['b'].size()[0] == L['a'].size()[0]
2 <= L['a'].size()[0]

我們看到在第一次呼叫時，守衛會檢查張量是否具有一些固定的尺寸和步幅。這些守衛在第二次執行時失敗，因此它會進行回溯。由於是一個 int 守衛失敗了，所以在第二次迭代中，它會符號地追蹤這個 int，並在這個更通用的核心上安裝更通用的守衛。

編譯效能提示。如果您知道維度的大小會發生變化，則可以在呼叫 torch.compile 之前呼叫 torch._dynamo.mark_dynamic 將其標記為動態的。這將避免第一次使用靜態形狀進行編譯。還有其他一些有用的工具函數，例如 maybe_mark_dynamic 或 mark_static。您也可以通過呼叫 torch.compile(dynamic=True) 來追蹤所有整數和形狀。這主要用於除錯目的。

0、1 總是專用的

無論我們是將維度標記為動態的，還是將整數追蹤為動態的，如果我們傳遞的輸入中該維度為 0 或 1，Dynamo 都會將其追蹤為非動態的，並為其生成一個特定的圖。這就是為什麼在上面的示例中，我們會發現 2 <= L['a'].size()[0] 形式的守衛。

做出這種選擇的原因有很多。其中有兩個特別重要：- 張量為空當且僅當其任何維度為零時 - 張量只有在其中一個步幅為一時才是連續的

鴨子類型

Dynamo 執行我們所說的「鴨子類型」。如果兩個動態整數在追蹤時具有相同的值，我們將假設它們相等，並對其進行守衛。實際上，這意味著我們不是在上面的示例中有兩個符號 s0、s1，而是將它們統一為 s0，並使用守衛 L['b'].size()[0] == L['a'].size()[0]。這使得我們能夠在編譯器內部執行融合，同時能夠生成足夠通用的核心。

符號整數上的守衛

現在，我們從高層次上了解了符號形狀是如何實現的，以及它們具有的特性。那麼，為什麼符號形狀迫使我們走上控制 CPython 解釋器的棘手路線呢？請看下面的例子

import torch

@torch.compile(dynamic=True)
def fn(a):
    if a.shape[0] * 2 < 16:
        return a
    else:
        return a + 1

fn(torch.randn(8))

這段程式碼有一個形式為 2*L['a'].size()[0] >= 16 的守衛。就函數的輸入而言，這是一個非常重要的守衛，但它是在程式執行過程中註冊的。更重要的是，在我們看到以 SymNodeVariable 引數為條件的 if 語句之前，我們無法知道需要這個守衛。這樣的條件對於 torch.jit.trace 來說是不可見的，需要對 Python 程式碼進行深入分析。

除錯提示 使用 TORCH_LOGS=dynamo 運行此程式碼會告訴我們這個守衛是在哪裡添加的

eval 2*s0 >= 16 [guard added] at script.py:5 in fn (_dynamo/variables/tensor.py:812 in evaluate_expr)

在那裡放置一個斷點並查看回溯對於理解守衛的來源非常有用。

使 Dynamo 變得完整：圖中斷

藉助我們討論的所有工具，我們有了一個追蹤器，它可以追蹤張量和整數上的 PyTorch 操作，並且擁有一個快取系統，該系統知道何時可以重用先前追蹤的圖，以及何時需要重新追蹤。所有這些都在執行任意的 Python 程式碼！

只有一個小問題。「執行任意的 Python 程式碼」這句話可能有點太籠統了。Dynamo 實現了 Python 的很大一部分，但它實現了更複雜的部分嗎，比如協程或異步？它實現了整個 Python 標準庫嗎？NumPy 也有一個 Python API。torch.compile 也能理解 NumPy 嗎？還有 Django？ 5

Python 的生態系統非常龐大，其中很大一部分是用其他更高效的語言（如 C++ 或 Rust）編寫的，並且只公開了 Python 綁定。Dynamo 不可能追蹤用 C++ 實現的 Python 物件。當追蹤器遇到它不理解的操作時，它該怎麼辦？

機器學習追蹤器處理此問題的常用方法是通知使用者它們卡住的操作，並完全放棄追蹤。這將在 PyTorch 中造成嚴重的可用性問題，因為 PyTorch 的使用者已經習慣了它提供的靈活性。舉一個真實的例子，doctr_det_predictor 模型使用 NumPy 和 cv2 庫來 對模型的結果進行後處理。

這裡是另一個可以使用 CPython 的有趣例子。Dynamo 不是報錯，而是讓 CPython 運行有問題的程式碼！為此，Dynamo 在追蹤時會生成一個圖，其中包含有問題程式碼之前的所有操作，以及一個包含之後所有操作的圖。 6 然後，在運行時，它將委託 CPython 執行第一個圖，然後是有問題的程式碼，最後是第二個圖。這種停止追蹤並生成多個圖的過程稱為圖中斷。

先承認一件事情：我在引言和第一部分完全是在說謊。Dynamo 並不是產生單一圖形，而是產生**多個圖形**！實際上，在第二個圖形之後開始回溯，可以想像成開始追蹤一個新的函式。圖形中斷後的新的圖形將會有自己的 guard、新的區域變數集合，等等。

為了討論如何實作圖形中斷，我們需要先重新審視 Dynamo 如何與 CPython 互動。透過 PEP 523，CPython 允許使用者使用自己的框架評估機制。我們沒有討論的是，CPython 也會公開自己的框架評估機制給其他人使用。Dynamo 利用這一點，讓快速的 CPython 直譯器可以執行編譯後的程式碼。對於沒有圖形中斷的函式來說，使用相同參數呼叫函式 2 次的程式的整個追蹤/執行過程如下所示：

1. 在第一次呼叫函式時

   1. Dynamo 會將函式追蹤到 FX 圖形中

      1. FX 圖形會由編譯器 (Inductor) 編譯成有效率的低階程式碼……但那是另一個故事了

   2. 它會改寫函式的 bytecode，讓它只呼叫編譯後的函式

   3. 它會將這個新的 bytecode 提供給 CPython，並要求它執行它 [這裡]

2. 在第二次呼叫函式時

   1. 它會根據新的參數檢查第一次呼叫的 guard [這裡]。因為它們與之前的參數相同，所以它們會通過

   2. 它會要求 CPython 執行與這些 guard 相關聯的 bytecode [這裡]

這個過程本身看起來過於複雜。為什麼要產生新的 bytecode 並要求 CPython 執行它，而不是簡單地建立編譯後函式的 C++ binding 並執行它？嗯，這種模式允許我們實作圖形中斷！由圖形中斷產生的 bytecode 具有以下結構：

1. 執行第一個圖形的 Bytecode

2. 將堆疊保留為如果 CPython 執行了第一個圖形時的狀態的 Bytecode。它也會重播任何對此時可見的區域變數或全域變數的修改

3. 讓 Dynamo 圖形中斷的 Bytecode

4. 執行第二個圖形的 Bytecode

讓我們在一個簡單的範例中看看這個

import torch

@torch.compile
def fn(a):
    b = a + 2
    print("Hi")
    return b + a

fn(torch.randn(4))

使用 TORCH_LOGS=bytecode 執行這個程式碼會顯示初始的 bytecode 和修改後的 bytecode

MODIFIED BYTECODE fn script.py line 3
 0 LOAD_GLOBAL              1 (__compiled_fn_0)
 2 LOAD_FAST                0 (a)
 4 CALL_FUNCTION            1
 6 STORE_FAST               3 (graph_out_0)
 8 LOAD_GLOBAL              0 (print)
10 LOAD_CONST               2 ('Hi')
12 LOAD_FAST                3 (graph_out_0)
14 LOAD_CONST               3 (0)
16 BINARY_SUBSCR
18 STORE_FAST               1 (b)

20 CALL_FUNCTION            1
22 LOAD_GLOBAL              2 (__resume_at_14_1)
24 ROT_TWO
26 LOAD_FAST                0 (a)
28 LOAD_FAST                1 (b)
30 CALL_FUNCTION            3
32 RETURN_VALUE

MODIFIED BYTECODE resume_in_fn script.py line 6
 0 LOAD_GLOBAL              1 (__compiled_fn_2)
 2 LOAD_FAST                2 (b)
 4 LOAD_FAST                1 (a)
 6 CALL_FUNCTION            2
 8 UNPACK_SEQUENCE          1
10 RETURN_VALUE

我們可以看到修改後的 bytecode 被分成兩個函式，fn 是原始函式，而另一個函式稱為 resume_in_fn。第二個函式是由 Dynamo 建立的，用於實作從圖形中斷開始的程式執行。這通常稱為 續延函式。這個續延函式只會使用正確的參數呼叫第二個編譯後的函式。初始函式的程式碼會根據我們前面描述的策略進行改寫

• L0-4. 呼叫編譯後的函式 (a + 2)。

• L6. 將其結果儲存在一個名為 graph_out_0 的區域變數中。graph_out_0 是一個 tuple

• L8-18. 將堆疊保留為圖形中斷時的狀態

• L20. 執行導致圖形中斷的程式碼

• L22-32. 呼叫編譯後的續延函式 (a + b)

Dynamo 中堆疊的程式碼產生是由 VariableTracker 子類別處理的。Dynamo 中的每個 VariableTracker 物件都有一個 reconstruct 方法，該方法會產生必要的 bytecode，以便在堆疊上建立它所代表的 Python 物件。

**除錯技巧**。圖形中斷會影響效能，因此最好避免它們。使用 TORCH_LOGS=graph_breaks 執行程式碼是一個很好的方法，可以找出程式碼遇到了多少次圖形中斷。它傳回的資訊是以 VariableTracker 物件的形式提供的，因此上述的除錯技巧有時也有助於找出導致圖形中斷的原因。

結論

Dynamo 是一個複雜的軟體。一旦你決定要實作 CPython 直譯器，你就知道你將會面臨一場硬仗。話雖如此，我們希望這篇文章能幫助你稍微了解它。

Dynamo（大部分）是用 Python 實作的。我們留下了許多連結，指向我們討論過的程式碼片段。我們希望閱讀這些程式碼片段並搜尋呼叫它們的位置，或者在它們上面設置斷點並查看呼叫堆疊，將有助於理解其餘的程式碼庫。

當然，學習一個軟體如何運作的最佳方法是擴展它。在這種情況下，最好的方法是查看 GitHub 上的開放 Dynamo 問題。其中許多問題只需要對程式碼進行非常小的更改，一旦你找到了需要進行更改的位置。

註腳

1

在文獻中，這被稱為有向無環圖 (DAG)。

2

所有這些 binding 程式碼都在 torch/csrc/dynamo/eval_frame.c 中。

3

在 CPython 術語中，所有這些物件的集合稱為 框架。

4

也有 SymBool 和 SymFloat 類別。在撰寫本文時，後者並沒有被廣泛使用。

5

有趣的是，它確實可以理解 NumPy 程式碼！請查看 這篇部落格文章 和 文件。現在，這只是因為我們使用 PyTorch 重新實作了 NumPy 才有可能。祝你好運在 PyTorch 中實作 Django……

6

假設只有一段有問題的程式碼。如果有多段，Dynamo 可以將程式碼分割成所需的圖形數量。