Nvtx.range_push/pop like facility in torch.profiler

Abhishek_Ghosh: <p>When running multiple iterations of some PyTorch module, to understand the situation better, one can use <code>torch.cuda.nvtx.range_push("iteration{}".format(iter))</code> and <code>torch.cuda.nvtx.range_pop()</code>, which gives a nice timing diagram on nsys as follows:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/c/7/c7d10d37bcd894f3953e6aa9bc29a273d68b09d4.png" data-download-href="https://discuss.pytorch.org/uploads/default/c7d10d37bcd894f3953e6aa9bc29a273d68b09d4" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/3X/c/7/c7d10d37bcd894f3953e6aa9bc29a273d68b09d4.png" alt="image" data-base62-sha1="svEUh6iws7PyTQo4lzpHWADHjsE" width="690" height="451" data-dominant-color="3D3E3F"><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">image</span><span class="informations">894×585 23.5 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>
<p>But the problem with nsys is that, it does not show the details of the stack traces (like the exact code regions with line numbers) during the run. In this situation, pytorch’s profiler comes handy.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/e/7e7b6f9e0e45e0e85f921610b7d8af5d5ae8beb5.png" data-download-href="https://discuss.pytorch.org/uploads/default/7e7b6f9e0e45e0e85f921610b7d8af5d5ae8beb5" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/e/7e7b6f9e0e45e0e85f921610b7d8af5d5ae8beb5_2_690x306.png" alt="image" data-base62-sha1="i2UCxOcl5NHPOBEuTDppVSY3Nyd" width="690" height="306" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/e/7e7b6f9e0e45e0e85f921610b7d8af5d5ae8beb5_2_690x306.png, https://discuss.pytorch.org/uploads/default/optimized/3X/7/e/7e7b6f9e0e45e0e85f921610b7d8af5d5ae8beb5_2_1035x459.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/7/e/7e7b6f9e0e45e0e85f921610b7d8af5d5ae8beb5.png 2x" data-dominant-color="3F3D3C"><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">image</span><span class="informations">1351×600 101 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>
<p>But the problem which I am facing is, I am unable to display in the <code>chrome://tracing</code> the iteration level separation as in nvtx.</p>
<p>I tried using <a href="https://pytorch.org/docs/stable/profiler.html#torch.profiler.itt.range_push" rel="noopener nofollow ugc">torch.profiler.itt.range_push</a> and  <a href="https://pytorch.org/docs/stable/profiler.html#torch.profiler.itt.range_pop" rel="noopener nofollow ugc">torch.profiler.itt.range_pop</a> from the <a href="https://pytorch.org/docs/stable/profiler.html#intel-instrumentation-and-tracing-technology-apis" rel="noopener nofollow ugc">Intel Instrumentation and Tracing Technology APIs</a> as follows:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
from torchvision.models import resnet18

model = resnet18().cuda()
inputs = [torch.randn((5, 3, 224, 224), device='cuda') for _ in range(10)]

model_c = torch.compile(model, mode="reduce-overhead")

def fwd_bwd(inp):
    out = model_c(inp)
    out.sum().backward()

# warm up
for i in range(5):
    fwd_bwd(inputs[0])

print(f"{torch.profiler.itt.is_available()=}")
with torch.profiler.profile(with_stack=True) as prof:
    for i in range(1, 4):
        torch.profiler.itt.range_push(f"iteration_{i}") # like torch.cuda.nvtx.range_push()
        fwd_bwd(inputs[i])
        torch.profiler.itt.range_pop() # like torch.cuda.nvtx.range_pop()
        prof.step()

prof.export_chrome_trace("trace_torch_compile_reduce-overhead_with_itt_range_push_pop.json")
</code></pre>
<p>Output:</p>
<pre><code class="lang-auto">$ python3 test_pytorch_profiler_test_torch_compile.py 
torch.profiler.itt.is_available()=True
STAGE:2024-05-09 19:05:18 27886:27886 ActivityProfilerController.cpp:314] Completed Stage: Warm Up
STAGE:2024-05-09 19:05:18 27886:27886 ActivityProfilerController.cpp:320] Completed Stage: Collection
STAGE:2024-05-09 19:05:18 27886:27886 ActivityProfilerController.cpp:324] Completed Stage: Post Processing
</code></pre>
<p>But it does not create any stack entry like the nvtx equivalent does.<br>
I just see an entry saying that the <code>itt.range_push</code> api was called:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/6/16ce88fb3f25a2b69db0b818e119894050d547b4.png" data-download-href="https://discuss.pytorch.org/uploads/default/16ce88fb3f25a2b69db0b818e119894050d547b4" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/6/16ce88fb3f25a2b69db0b818e119894050d547b4_2_690x143.png" alt="image" data-base62-sha1="3fL0ek05Oov8UebfMFeQPbD1jM0" width="690" height="143" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/6/16ce88fb3f25a2b69db0b818e119894050d547b4_2_690x143.png, https://discuss.pytorch.org/uploads/default/optimized/3X/1/6/16ce88fb3f25a2b69db0b818e119894050d547b4_2_1035x214.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/1/6/16ce88fb3f25a2b69db0b818e119894050d547b4_2_1380x286.png 2x" data-dominant-color="2D302E"><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">image</span><span class="informations">1564×326 15.2 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>
<p>How to distinctly indicate that we are moving across different iterations.</p>

ptrblck: <aside class="quote no-group" data-username="Abhishek_Ghosh" data-post="1" data-topic="202431">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/abhishek_ghosh/48/68123_2.png" class="avatar"> Abhishek_Ghosh:</div>
<blockquote>
<p>But the problem with nsys is that, it does not show the details of the stack traces (like the exact code regions with line numbers) during the run.</p>
</blockquote>
</aside>
<p>You can enable stacktraces in Nsight Systems (including pointing to Python in recent versions).</p>

Abhishek_Ghosh: <aside class="quote no-group" data-username="ptrblck">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>We can zoom into a specific iteration and check the backtrace option, which is often helpful to isolate specific regressions and see which functions were invoked.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/1/8/1816ea15e9a672d38ef5adf18451110b5d4052d0.png" data-download-href="https://discuss.pytorch.org/uploads/default/1816ea15e9a672d38ef5adf18451110b5d4052d0" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/1/8/1816ea15e9a672d38ef5adf18451110b5d4052d0_2_690x343.png" alt="image" data-base62-sha1="3r6yilqBXHQJnLdy9UU2JynQTK0" width="690" height="343" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/1/8/1816ea15e9a672d38ef5adf18451110b5d4052d0_2_690x343.png, https://discuss.pytorch.org/uploads/default/optimized/3X/1/8/1816ea15e9a672d38ef5adf18451110b5d4052d0_2_1035x514.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/1/8/1816ea15e9a672d38ef5adf18451110b5d4052d0_2_1380x686.png 2x" data-dominant-color="E7E8E2"><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">image</span><span class="informations">1918×956 107 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>
</blockquote>
</aside>
<p>In the above <a href="https://dev-discuss.pytorch.org/t/using-nsight-systems-to-profile-gpu-workload/59" rel="noopener nofollow ugc">post</a>, you talk about backtraces. They are not quite helpful in my case.</p>
<p>Can you please point me to some required documentation which deals with stacktrace in nsys?</p>