A Beginner's Guide to Debugging AOT Inductor CUDA Illegal Memory Access

henrylhtsang: <p>I want to share some basic tips and tools for debugging CUDA illegal memory access (IMA) for AOT Inductor (AOTI). In the future I might fold this into pytorch docs.</p>
<h1><a name="p-473980-background-1" class="anchor" href="#p-473980-background-1"></a>Background</h1>
<p><a href="https://docs.pytorch.org/docs/stable/torch.compiler_aot_inductor.html" rel="noopener nofollow ugc">AOT Inductor</a> is part of the PT2 stack, similar to torch.compile, but it produces a compilation artifact that can work in a C++ environment. Another key difference is that errors can happen as a CUDA illegal memory error, which can happen non-deterministically and even appear transient at times. This note discusses a few ways to debug such errors.</p>
<h1><a name="p-473980-h-0-sanity-checks-2" class="anchor" href="#p-473980-h-0-sanity-checks-2"></a>0. Sanity checks</h1>
<p>Before we dive deep into reliably reproducing the error, it is good to try out some existing debugging flags:</p>
<pre><code class="lang-auto">AOTI_RUNTIME_CHECK_INPUTS=1
TORCHINDUCTOR_NAN_ASSERTS=1
</code></pre>
<p>These flags take effect at compilation time (more precisely, at codegen time). Just to explain:</p>
<ul>
<li><a href="https://docs.pytorch.org/docs/stable/torch.compiler_aot_inductor.html#troubleshooting" rel="noopener nofollow ugc"><code>AOTI_RUNTIME_CHECK_INPUTS=1</code></a> checks if the inputs satisfy the same set of guards used during compilation.</li>
<li><code>TORCHINDUCTOR_NAN_ASSERTS=1</code> adds codegen before and after each Inductor’s kernel to check for NaN.</li>
</ul>
<h1><a name="p-473980-h-1-pinpoint-the-cuda-ima-3" class="anchor" href="#p-473980-h-1-pinpoint-the-cuda-ima-3"></a>1. Pinpoint the CUDA IMA</h1>
<p>Again, one hard part is CUDA IMA error can be non-deterministic. It can happen at different locations, and sometimes not happen at all (though that just means the numerics are silently incorrect). With the following two flags, we can trigger the error deterministically.</p>
<pre><code class="lang-auto">PYTORCH_NO_CUDA_MEMORY_CACHING=1
CUDA_LAUNCH_BLOCKING=1 
</code></pre>
<p>These flags take effect at runtime. To explain:</p>
<ul>
<li><code>PYTORCH_NO_CUDA_MEMORY_CACHING=1</code> disables PyTorch’s Caching Allocator, which allocates a bigger buffer than needed immediately to reduce the number of buffer allocations. This is usually the reason why CUDA illegal memory access errors are non-deterministic.</li>
</ul>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/8/f8d2c7a502dba7a7b06ffe01237942929e62b2bc.png" data-download-href="https://discuss.pytorch.org/uploads/default/f8d2c7a502dba7a7b06ffe01237942929e62b2bc" title="cuda illegal memory access"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/f/8/f8d2c7a502dba7a7b06ffe01237942929e62b2bc_2_690x244.png" alt="cuda illegal memory access" data-base62-sha1="zvc0oLlvnx7qPMT0ZBmxor301Tm" width="690" height="244" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/f/8/f8d2c7a502dba7a7b06ffe01237942929e62b2bc_2_690x244.png, https://discuss.pytorch.org/uploads/default/optimized/3X/f/8/f8d2c7a502dba7a7b06ffe01237942929e62b2bc_2_1035x366.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/f/8/f8d2c7a502dba7a7b06ffe01237942929e62b2bc_2_1380x488.png 2x" data-dominant-color="E5EEF6"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">cuda illegal memory access</span><span class="informations">2896×1028 220 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<ul>
<li><code>CUDA_LAUNCH_BLOCKING=1</code> forces the kernels to launch one at a time. Without this, we would get the famous “CUDA kernel errors might be asynchronously reported at some other API call” warning since kernels are launched asynchronously.</li>
</ul>
<h1><a name="p-473980-h-2-identify-problematic-kernels-with-intermediate-value-debugging-4" class="anchor" href="#p-473980-h-2-identify-problematic-kernels-with-intermediate-value-debugging-4"></a>2. Identify problematic kernels with Intermediate value debugging</h1>
<p>This is a useful tool, developed by <a href="https://github.com/YUNQIUGUO" rel="noopener nofollow ugc">Rachel Guo</a> and mentioned in <a href="https://blog.ezyang.com/2024/12/ways-to-use-torch-export/" rel="noopener nofollow ugc">Ed’s blog</a>, to help pinpoint the problematic kernel and get information about the inputs and outputs of said kernel.</p>
<p>We first use</p>
<pre><code class="lang-auto">AOT_INDUCTOR_DEBUG_INTERMEDIATE_VALUE_PRINTER=3
</code></pre>
<p>which takes effect at compilation time and prints the kernel ones one by one at runtime. Together with the previous flags, this would let us know which kernel was launched right before the error happened.</p>
<p>However, it is important to note that just because the error happened in that kernel, it doesn’t mean that kernel is problematic. For example, it can happen that an earlier kernel is problematic and produces some wrong outputs. So the natural next step is to inspect the inputs to the problematic kernel, which can be done by</p>
<pre><code class="lang-auto">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
</code></pre>
<p>where the filtered kernels to print environment variable has the names of the kernels you want to inspect. If the inputs to the kernel is not as expected, you then inspect the kernel that produces the bad input.</p>
<h1><a name="p-473980-h-3-other-tools-to-make-debugging-easier-5" class="anchor" href="#p-473980-h-3-other-tools-to-make-debugging-easier-5"></a>3. Other tools to make debugging easier</h1>
<ul>
<li><a href="https://docs.pytorch.org/docs/stable/torch.compiler_troubleshooting.html#tlparse-torch-trace" rel="noopener nofollow ugc">tlparse / TORCH_TRACE</a></li>
</ul>
<p>It has complete output codes for inspection. It also records the set of guards used.</p>
<ul>
<li><a href="https://docs.pytorch.org/docs/stable/torch.compiler_troubleshooting.html#torch-logs" rel="noopener nofollow ugc">TORCH_LOGS</a>=”+inductor,output_code”</li>
</ul>
<p>To see more PT2 internal logs.</p>
<ul>
<li>TORCH_SHOW_CPP_STACKTRACES=1</li>
</ul>
<p>To potentially see more stack traces.</p>
<ul>
<li>Pay attention to <a href="https://docs.pytorch.org/docs/stable/torch.compiler_dynamic_shapes.html" rel="noopener nofollow ugc">dynamic shapes</a></li>
</ul>
<p>Historically it has been a source of many IMAs.</p>
<ul>
<li>Pay attention to Custom ops</li>
</ul>
<p>Especially when implemented in C++ and used with dynamic shapes. There is a need to Symint’ify the meta function.</p>