草稿匯出

創建於：2025 年 6 月 13 日 | 最後更新於：2025 年 7 月 16 日

警告

此功能不適用於生產環境，旨在作為除錯 `torch.export` 跟蹤錯誤的工具。

草稿匯出是匯出功能的一個新版本，旨在穩定地生成一個圖，即使存在潛在的正確性問題，並生成一個報告，列出匯出在跟蹤過程中遇到的所有問題，並提供額外的除錯資訊。對於沒有假核心的自定義運算子，它還將生成一個配置檔案，您可以註冊該檔案以自動生成假核心。

您是否曾嘗試使用 torch.export.export() 匯出模型，卻遇到了資料依賴問題？您修復了它，但又遇到了缺少假核心的問題。解決此問題後，您又遇到了另一個數據依賴問題。您不禁想，要是有一種方法可以讓我只獲取一個圖來除錯，並能在同一處檢視所有問題，這樣我就可以以後再修復它們……

現在有 draft_export 來幫忙！

draft_export 是匯出功能的一個版本，它將始終成功匯出圖，即使存在潛在的正確性問題。這些問題隨後將被編譯成報告，以便更清晰地視覺化，以後可以進行修復。

它能捕獲哪些型別的錯誤？

草稿匯出有助於捕獲和除錯以下錯誤：

• 資料依賴錯誤上的守衛

• 約束違反錯誤

• 缺少假核心

• 錯誤的假核心實現

它是如何工作的？

在正常的匯出中，我們會將示例輸入轉換為 `FakeTensors`，並使用它們來記錄操作並將程式跟蹤為圖。形狀可能改變的輸入張量（透過 `dynamic_shapes` 標記）或張量中的值（通常來自 `.item()` 呼叫）將被表示為符號形狀（`SymInt`），而不是具體的整數。然而，在跟蹤過程中可能會發生一些問題——我們可能會遇到無法評估的守衛，例如，如果我們想檢查張量中的某個項是否大於 0（`u0 >= 0`）。由於跟蹤器對 `u0` 的值一無所知，它將丟擲資料依賴錯誤。如果模型使用了自定義運算子但尚未為其定義假核心，那麼我們將因 `fake_tensor.UnsupportedOperatorException` 而出錯，因為匯出不知道如何在 `FakeTensors` 上應用它。如果自定義運算子的假核心實現不正確，匯出將默默地生成一個不匹配即時行為的不正確圖。

為了修復上述錯誤，草稿匯出使用 *真實張量跟蹤* 來指導我們在跟蹤時如何進行。當我們使用假張量跟蹤模型時，對於在假張量上發生的每個操作，草稿匯出還會對儲存的真實張量（來自傳遞給匯出的示例輸入）執行該運算子。這使我們能夠解決上述錯誤：當我們遇到無法評估的守衛（如 `u0 >= 0`）時，我們將使用儲存的真實張量值來評估該守衛。執行時斷言將被新增到圖中，以確保圖斷言與我們在跟蹤時所假設的守衛相同。如果我們遇到沒有假核心的自定義運算子，我們將使用儲存的真實張量執行該運算子的正常核心，並返回一個具有相同秩但形狀未繫結的假張量。由於我們有每個操作的真實張量輸出，我們將把這個與假核心的假張量輸出進行比較。如果假核心實現不正確，我們將捕獲此行為並生成更正確的假核心。

如何使用草稿匯出？

假設您正在嘗試匯出以下程式碼片段

class M(torch.nn.Module):
    def forward(self, x, y, z):
        res = torch.ops.mylib.foo2(x, y)

        a = res.item()
        a = -a
        a = a // 3
        a = a + 5

        z = torch.cat([z, z])

        torch._check_is_size(a)
        torch._check(a < z.shape[0])

        return z[:a]

inp = (torch.tensor(3), torch.tensor(4), torch.ones(3, 3))

ep = torch.export.export(M(), inp)

這會因為 `mylib.foo2` 的“缺少假核心”錯誤以及由於 `z` 使用 `a`（一個未繫結的 `symint`）進行切片而導致的 `GuardOnDataDependentExpression` 錯誤。

要呼叫 `draft-export`，我們可以將 `torch.export` 行替換為以下內容：

ep = torch.export.draft_export(M(), inp)

ep 是一個有效的 `ExportedProgram`，現在可以傳遞給進一步的環境！

使用草稿匯出進行除錯

在草稿匯出的終端輸出中，您應該會看到以下訊息：

#########################################################################################
WARNING: 2 issue(s) found during export, and it was not able to soundly produce a graph.
To view the report of failures in an html page, please run the command:
    `tlparse /tmp/export_angelayi/dedicated_log_torch_trace_axpofwe2.log --export`
Or, you can view the errors in python by inspecting `print(ep._report)`.
########################################################################################

草稿匯出會自動轉儲 `tlparse` 的日誌。您可以使用 `print(ep._report)` 檢視跟蹤錯誤，或者將日誌傳遞給 `tlparse` 來生成 HTML 報告。

在終端中執行 `tlparse` 命令將生成一個 tlparse HTML 報告。這是一個 `tlparse` 報告的示例：

點選“資料依賴錯誤”，我們將看到以下頁面，其中包含幫助除錯此錯誤的資訊。具體來說，它包含：

• 發生此錯誤的堆疊跟蹤

• 區域性變數及其形狀列表

• 有關此守衛如何建立的資訊

返回的 `ExportedProgram`

由於草稿匯出根據示例輸入對程式碼路徑進行專門化，因此從草稿匯出返回的 `ExportedProgram` **至少**保證對於給定的示例輸入是可執行的並返回正確的結果。其他輸入也可以工作，只要它們匹配草稿匯出時所採取的守衛。

例如，如果我們有一個基於值是否大於 5 的圖分支，如果在草稿匯出中我們的示例輸入大於 5，那麼返回的 `ExportedProgram` 將專門化該分支，並斷言該值大於 5。這意味著如果您傳入另一個大於 5 的值，程式將成功，但如果您傳入一個小於 5 的值，程式將失敗。這比 `torch.jit.trace` 更安全，後者會默默地專門化分支。`torch.export` 支援兩個分支的正確方法是使用 `torch.cond` 重寫程式碼，它將捕獲兩個分支。

由於圖中的執行時斷言，返回的 `exported-program` 也可以使用 `torch.export` 或 `torch.compile` 進行重新跟蹤，並且在自定義運算子缺少假核心的情況下需要進行少量額外操作。

生成假核心

如果自定義運算子不包含假實現，目前草稿匯出將使用真實張量傳播來獲取運算子的輸出並繼續跟蹤。然而，如果我們使用假張量執行匯出的程式或重新跟蹤匯出的模型，我們仍然會失敗，因為仍然沒有假核心實現。

為了解決這個問題，在草稿匯出之後，我們將為遇到的每個自定義運算子呼叫生成一個運算子配置檔案，並將其儲存在附加到匯出的程式的報告中：`ep._report.op_profiles`。然後，使用者可以使用上下文管理器 `torch._library.fake_profile.unsafe_generate_fake_kernels` 基於這些運算子配置檔案生成和註冊一個假實現。這樣，未來的假張量重新跟蹤就能正常工作。

工作流程可能如下所示：

class M(torch.nn.Module):
    def forward(self, a, b):
        res = torch.ops.mylib.foo(a, b)  # no fake impl
        return res

ep = draft_export(M(), (torch.ones(3, 4), torch.ones(3, 4)))

with torch._library.fake_profile.unsafe_generate_fake_kernels(ep._report.op_profiles):
    decomp = ep.run_decompositions()

new_inp = (
    torch.ones(2, 3, 4),
    torch.ones(2, 3, 4),
)

# Save the profile to a yaml and check it into a codebase
save_op_profiles(ep._report.op_profiles, "op_profile.yaml")
# Load the yaml
loaded_op_profile = load_op_profiles("op_profile.yaml")

運算子配置檔案是一個字典，將運算子名稱對映到一組配置檔案，這些配置檔案描述了運算子的輸入和輸出，並且可以手動編寫、儲存到 yaml 檔案並提交到程式碼庫。下面是一個 `mylib.foo.default` 的配置檔案的示例：

"mylib.foo.default": {
    OpProfile(
        args_profile=(
            TensorMetadata(
                rank=2,
                dtype=torch.float32,
                device=torch.device("cpu"),
                layout=torch.strided,
            ),
            TensorMetadata(
                rank=2,
                dtype=torch.float32,
                device=torch.device("cpu"),
                layout=torch.strided,
            ),
        ),
        out_profile=TensorMetadata(
            rank=2,
            dtype=torch.float32,
            device=torch.device("cpu"),
            layout=torch.strided,
        ),
    )
}

`mylib.foo.default` 的配置檔案只包含一個配置檔案，它表示對於 2 個秩為 2、`dtype` 為 `torch.float32`、裝置為 `cpu` 的輸入張量，我們將返回一個秩為 2、`dtype` 為 `torch.float32`、裝置為 `cpu` 的輸出張量。使用上下文管理器，將生成一個假核心，當給定 2 個秩為 2 的輸入張量（以及其他張量元資料）時，它將輸出一個秩為 2 的張量（以及其他張量元資料）。

如果該運算子還支援其他秩，那麼我們可以將該配置檔案新增到此配置檔案的列表中，方法是手動將其新增到現有配置檔案中，或使用新輸入重新執行草稿匯出以獲取新配置檔案，這樣生成的假核心將支援更多輸入型別。否則，它將出錯。

下一步去哪裡？

既然我們已經成功使用草稿匯出建立了一個 `ExportedProgram`，我們可以使用像 `AOTInductor` 這樣的進一步編譯器來最佳化其效能並生成可執行的製品。這個最佳化版本可以用於部署。同時，我們可以利用草稿匯出生成的報告來識別和修復遇到的 `torch.export` 錯誤，以便原始模型可以直接使用 `torch.export` 進行跟蹤。