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 を指定して実行すると、より明確なエラーが表示されます。このフラグはデバイスタイプ、dtype、サイズ、ストライドなど、すべての入力プロパティを検証します：

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 は、カーネルを一度に 1 つずつ起動するよう強制します。これを指定しないと、カーネルは非同期に起動するため、「CUDA カーネルエラーは非同期に報告される可能性があります」という警告が表示されます。

ステップ 3：中間値デバッガーで問題のあるカーネルを特定する

AOTI 中間値デバッガーを使って、問題のあるカーネルを特定してください：

AOT_INDUCTOR_DEBUG_INTERMEDIATE_VALUE_PRINTER=3

実行時にカーネルが 1 つずつ出力されます。前述のフラグと組み合わせると、エラーの直前に起動したカーネルを確認できます。

特定のカーネルへの入力を調べるには：

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：完全な出力コードを提供し、使用されたガードを記録します
• TORCH_LOGS：TORCH_LOGS="+inductor,output_code" を使用すると、PT2 の内部ログをさらに確認できます
• TORCH_SHOW_CPP_STACKTRACES：1 に設定すると、スタックトレースをさらに確認できます

よくある問題の原因

• 動的シェイプ：これまでは、多くの IMA の原因となっていました。動的シェイプのシナリオをデバッグする際は、特に注意してください。
• カスタム op：特に動的シェイプを使用して 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	コンパイル時	入力がコンパイル時のガードに一致することを検証します
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++ のスタックトレースを表示します