aoti-debug

AOTI 除錯指南

本 Skill 協助診斷與修復常見的 AOTInductor 問題。

錯誤樣式路由

請檢查錯誤訊息並轉至對應的子指南：

Triton 索引越界 (Triton Index Out of Bounds)

若錯誤符合以下樣式：

Assertion `index out of bounds: 0 <= tmpN < ksM` failed

→ 請參閱 triton-index-out-of-bounds.md 指南

其他所有錯誤

請繼續閱讀下方章節。

---

第一步：務必先確認裝置與 Shape 是否匹配

針對任何 AOTI 錯誤（segfault、例外狀況、當機、輸出不正確），請務必優先檢查以下幾點：

1. 編譯裝置 == 載入裝置：模型載入的裝置類型，必須與編譯時所使用的裝置類型完全一致
2. 輸入裝置一致：執行階段 (runtime) 的輸入資料必須與已編譯模型處於相同的裝置上
3. 輸入 Shape 一致：執行階段輸入資料的 Shape，必須符合編譯時採用的 Shape（或滿足動態 Shape 限制）

# 編譯期間 - 請留意裝置與 shape
model = MyModel().eval()           # 使用什麼裝置？CPU 還是 .cuda()？
inp = torch.randn(2, 10)           # 使用什麼裝置？Shape 是什麼？
compiled_so = torch._inductor.aot_compile(model, (inp,))

# 載入期間 - 裝置類型必須與編譯時一致
loaded = torch._export.aot_load(compiled_so, "???")  # 必須與上方的模型/輸入裝置一致

# 推論期間 - 裝置與 shape 必須匹配
out = loaded(inp.to("???"))  # 必須匹配編譯裝置，shape 也必須一致

若有任何一項不匹配，將會引發從 segfault、例外狀況到輸出錯誤等各種問題。

核心限制：裝置類型匹配 (Device Type Matching)

AOTI 要求編譯與載入必須使用相同的裝置類型。

• 若在 CUDA 上編譯，就必須在 CUDA 上載入（裝置編號/device index 可以不同）
• 若在 CPU 上編譯，就必須在 CPU 上載入
• 不支援跨裝置載入（例如：在 GPU 上編譯，卻在 CPU 上載入）

常見錯誤樣式

1. 裝置不符導致 Segfault (Device Mismatch Segfault)

症狀：在呼叫 aot_load() 或執行模型時發生 segfault、例外狀況或當機。

常見錯誤訊息範例：

• The specified pointer resides on host memory and is not registered with any CUDA device
• 在 AOTInductorModelBase 載入常數期間當機
• Expected out tensor to have device cuda:0, but got cpu instead

原因：編譯與載入的裝置類型不一致（請參閱前述的「第一步」）。

解法：確保編譯與載入使用相同的裝置類型。若在 CPU 編譯，請在 CPU 載入；若在 CUDA 編譯，請在 CUDA 載入。

2. 執行階段輸入裝置不一致 (Input Device Mismatch at Runtime)

症狀：模型執行時拋出 RuntimeError。

原因：輸入裝置與編譯裝置不匹配（請參閱前述的「第一步」）。

更佳的排查方式：設定 AOTI_RUNTIME_CHECK_INPUTS=1 來執行程式，以獲得更清晰的錯誤訊息。此標誌會驗證所有輸入屬性，包含裝置類型、dtype、尺寸 (sizes) 與步長 (strides)：

AOTI_RUNTIME_CHECK_INPUTS=1 python your_script.py

這會產生明確且可據以操作的錯誤訊息，例如：

Error: input_handles[0]: unmatched device type, expected: 0(cpu), but got: 1(cuda)

排查 CUDA 非法記憶體存取 (IMA) 錯誤

若遇到 CUDA 非法記憶體存取 (CUDA illegal memory access) 錯誤，請遵循以下系統化步驟排查：

步驟 1：基本檢查 (Sanity Checks)

在深入分析前，可先嘗試啟用以下排查標誌：

AOTI_RUNTIME_CHECK_INPUTS=1
TORCHINDUCTOR_NAN_ASSERTS=1

這些標誌會在編譯階段（程式碼生成/codegen 階段）生效：

• AOTI_RUNTIME_CHECK_INPUTS=1 檢查輸入是否符合編譯時所使用的 guard
• TORCHINDUCTOR_NAN_ASSERTS=1 在每個 kernel 執行前後加入檢查 NaN 的程式碼

步驟 2：精準定位 CUDA IMA 錯誤

CUDA IMA 錯誤有時具有非確定性 (non-deterministic)。可使用以下標誌來確定性地重現錯誤：

PYTORCH_NO_CUDA_MEMORY_CACHING=1
CUDA_LAUNCH_BLOCKING=1

這些標誌會在執行階段生效：

• PYTORCH_NO_CUDA_MEMORY_CACHING=1 停用 PyTorch 的快取分配器 (Caching Allocator)，該分配器平時會預先分配超出即時所需的較大緩衝區，這通常是導致 CUDA IMA 錯誤不固定的主因。
• CUDA_LAUNCH_BLOCKING=1 強制 kernel 逐一同步發射。由於 kernel 預設是非同步發射的，若不啟用此標誌，只會收到警告訊息 "CUDA kernel errors might be asynchronously reported"。

步驟 3：使用中間值除錯器識別問題 Kernel

使用 AOTI 中間值除錯器 (Intermediate Value Debugger) 來精確定位出問題的 kernel：

AOT_INDUCTOR_DEBUG_INTERMEDIATE_VALUE_PRINTER=3

這會在執行階段逐一印出 kernel。搭配前面的環境變數標誌，即可看出是在哪一個 kernel 發射後立刻引發了錯誤。

若要檢視特定 kernel 的輸入內容：

AOT_INDUCTOR_FILTERED_KERNELS_TO_PRINT="triton_poi_fused_add_ge_logical_and_logical_or_lt_231,_add_position_embeddings_kernel_5" AOT_INDUCTOR_DEBUG_INTERMEDIATE_VALUE_PRINTER=2

如果傳入該 kernel 的輸入資料異常，請進一步檢查產出該不正常輸入資料的前一個 kernel。

其他除錯工具

日誌紀錄與追蹤 (Logging and Tracing)

• tlparse / TORCH_TRACE：提供完整的輸出程式碼並記錄使用的 guard
• TORCH_LOGS：設定 TORCH_LOGS="+inductor,output_code" 可查看更多 PT2 的內部日誌
• TORCH_SHOW_CPP_STACKTRACES：設定為 1 可顯示更詳細的 C++ 呼叫堆疊 (stack traces)

常見問題來源

• 動態 Shape (Dynamic shapes)：歷史經驗中，動態 Shape 是許多 IMA 錯誤的根源。在排查動態 Shape 相關情境時請特別注意。
• 自訂運算子 (Custom ops)：特別是使用 C++ 實作且搭配動態 Shape 的情況。其中 meta 函式可能需要支援 SymInt (Symint'ified)。

API 變更說明

已廢棄 API (Deprecated)

torch._export.aot_compile()  # Deprecated
torch._export.aot_load()     # Deprecated

現行 API

torch._inductor.aoti_compile_and_package()
torch._inductor.aoti_load_package()

新版 API 會將裝置元資料 (metadata) 儲存在 package 中，因此 aoti_load_package() 會自動套用正確的裝置類型。使用者只能變更裝置編號 (device index，例如 cuda:0 換成 cuda:1)，無法變更裝置類型 (device type)。

環境變數總覽

環境變數	生效時機	用途
AOTI_RUNTIME_CHECK_INPUTS=1	編譯階段	驗證輸入是否符合編譯時的 guard
TORCHINDUCTOR_NAN_ASSERTS=1	編譯階段	在 kernel 執行前後檢查是否有 NaN
PYTORCH_NO_CUDA_MEMORY_CACHING=1	執行階段	使 IMA 錯誤能確定性地重複發生
CUDA_LAUNCH_BLOCKING=1	執行階段	強制 kernel 以同步方式發射
AOT_INDUCTOR_DEBUG_INTERMEDIATE_VALUE_PRINTER=3	編譯階段	在執行階段印出 kernel 資訊
TORCH_LOGS="+inductor,output_code"	執行階段	查看 PT2 內部日誌
TORCH_SHOW_CPP_STACKTRACES=1	執行階段	顯示 C++ 呼叫堆疊 (stack traces)