aoti-debug

AOTI 调试指南

此技能可帮助诊断和修复常见的 AOTInductor 问题。

错误模式路由

检查错误消息，并转到相应的子指南：

Triton 索引越界

如果错误符合此模式：

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

→ 按照 triton-index-out-of-bounds.md 中的指南操作

其他所有错误

继续阅读以下章节。

---

第一步：始终检查设备和形状是否匹配

对于任何 AOTI 错误（段错误、异常、崩溃、输出错误），都必须先检查以下各项：

1. 编译设备 == 加载设备：模型必须加载到与编译时相同类型的设备上
2. 输入设备匹配：运行时输入必须与编译后的模型位于同一设备上
3. 输入形状匹配：运行时输入形状必须与编译时使用的形状匹配（或满足动态形状约束）

# During compilation - note the device and shapes
model = MyModel().eval()           # What device? CPU or .cuda()?
inp = torch.randn(2, 10)           # What device? What shape?
compiled_so = torch._inductor.aot_compile(model, (inp,))

# During loading - device type MUST match compilation
loaded = torch._export.aot_load(compiled_so, "???")  # Must match model/input device above

# During inference - device and shapes MUST match
out = loaded(inp.to("???"))  # Must match compile device, shape must match

如果其中任何一项不匹配，就可能出现从段错误、异常到输出错误的各种问题。

关键约束：设备类型匹配

AOTI 要求编译和加载使用相同的设备类型。

• 如果在 CUDA 上编译，就必须在 CUDA 上加载（设备索引可以不同）
• 如果在 CPU 上编译，就必须在 CPU 上加载
• 不支持跨设备加载（例如，在 GPU 上编译、在 CPU 上加载）

常见错误模式

1. 设备不匹配导致段错误

症状：在 aot_load() 期间或模型执行时发生段错误、异常或崩溃。

错误消息示例：

• 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. 运行时输入设备不匹配

症状：模型执行期间发生 RuntimeError。

原因：输入设备与编译设备不匹配（参见上文“第一步”）。

更好的调试方法：使用 AOTI_RUNTIME_CHECK_INPUTS=1 运行，以获得更清晰的错误信息。此标志会验证所有输入属性，包括设备类型、数据类型、大小和步幅：

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 非法内存访问错误，请按以下系统化步骤操作：

步骤 1：基本检查

深入排查前，先尝试使用以下调试标志：

AOTI_RUNTIME_CHECK_INPUTS=1
TORCHINDUCTOR_NAN_ASSERTS=1

这些标志在编译时生效（代码生成时）：

• AOTI_RUNTIME_CHECK_INPUTS=1 检查输入是否满足编译时使用的相同守卫条件
• TORCHINDUCTOR_NAN_ASSERTS=1 在每个内核前后添加代码生成逻辑，以检查 NaN

步骤 2：定位 CUDA IMA

CUDA IMA 错误可能具有非确定性。使用以下标志可稳定触发错误：

PYTORCH_NO_CUDA_MEMORY_CACHING=1
CUDA_LAUNCH_BLOCKING=1

这些标志在运行时生效：

• PYTORCH_NO_CUDA_MEMORY_CACHING=1 禁用 PyTorch 的缓存分配器。该分配器会立即分配比所需容量更大的缓冲区，这通常是 CUDA IMA 错误具有非确定性的原因。
• CUDA_LAUNCH_BLOCKING=1 强制内核逐个启动。否则，由于内核会异步启动，你会看到“CUDA 内核错误可能会异步报告”的警告。

步骤 3：使用中间值调试器识别问题内核

使用 AOTI 中间值调试器定位问题内核：

AOT_INDUCTOR_DEBUG_INTERMEDIATE_VALUE_PRINTER=3

该工具会在运行时逐个打印内核。结合之前的标志，可以查看错误发生前启动的是哪个内核。

要检查特定内核的输入：

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

如果内核的输入不符合预期，请检查生成该错误输入的内核。

其他调试工具

日志记录和跟踪

• tlparse / TORCH_TRACE：提供完整的输出代码，并记录使用的 guards
• TORCH_LOGS：使用 TORCH_LOGS="+inductor,output_code" 查看更多 PT2 内部日志
• TORCH_SHOW_CPP_STACKTRACES：设为 1 以查看更多堆栈跟踪

常见问题来源

• 动态形状：过去一直是许多 IMA 的来源。调试动态形状场景时请特别留意。
• 自定义算子：尤其是使用动态形状以 C++ 实现的算子。可能需要将 meta 函数改为使用 Symint。

API 说明

已弃用的 API

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

当前 API

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

新 API 会在程序包中存储设备元数据，因此 aoti_load_package() 会自动使用正确的设备类型。你只能更改设备 索引（例如 cuda:0 与 cuda:1），不能更改设备 类型。

环境变量汇总

变量	时机	用途
AOTI_RUNTIME_CHECK_INPUTS=1	编译时	验证输入是否符合编译 guards
TORCHINDUCTOR_NAN_ASSERTS=1	编译时	检查内核执行前后是否存在 NaN
PYTORCH_NO_CUDA_MEMORY_CACHING=1	运行时	让 IMA 错误具有确定性
CUDA_LAUNCH_BLOCKING=1	运行时	强制同步启动内核
AOT_INDUCTOR_DEBUG_INTERMEDIATE_VALUE_PRINTER=3	编译时	在运行时打印内核
TORCH_LOGS="+inductor,output_code"	运行时	查看 PT2 内部日志
TORCH_SHOW_CPP_STACKTRACES=1	运行时	显示 C++ 堆栈跟踪