Leveraging NVTX Structure

Jonathan1: <p>I use NVTX and Nsight when training with PyTorch and the level of detail from traces is just outstanding, those system details have been <em>critical</em> to my research. Kudos to the team on such great work!</p>
<p>My question is, how can I reuse some of this infrastructure for a Python/CUDA project? That is, how do I get up and running to leverage, maybe a minimal subset, torch’s NVTX scaffolding?</p>
<p>For example,</p>
<ol>
<li>I would love to have complete function tracing, where any Python/C++ function is traced <em>all the way down</em> including all sub-function calls.</li>
<li>Also, <strong>seq</strong> and <strong>op_id</strong> fields on ranges are great to have.</li>
</ol>
<p>Representative examples are shown below. That level of detail, I would love to have <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/0/4/04438475b906bc190b3c458136e335379160a1c4.png" data-download-href="https://discuss.pytorch.org/uploads/default/04438475b906bc190b3c458136e335379160a1c4" title="Screenshot 2024-06-24 at 12.51.17 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/0/4/04438475b906bc190b3c458136e335379160a1c4_2_690x234.png" alt="Screenshot 2024-06-24 at 12.51.17 PM" data-base62-sha1="BIz2Yo3vVAFvOOysfhbIXJWTfC" width="690" height="234" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/0/4/04438475b906bc190b3c458136e335379160a1c4_2_690x234.png, https://discuss.pytorch.org/uploads/default/optimized/3X/0/4/04438475b906bc190b3c458136e335379160a1c4_2_1035x351.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/0/4/04438475b906bc190b3c458136e335379160a1c4_2_1380x468.png 2x" data-dominant-color="DCDCDC"><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">Screenshot 2024-06-24 at 12.51.17 PM</span><span class="informations">1650×560 66.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><br>
<em>All the way down</em></p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/5/7/57a289c4bc9c5bb22f34bfe7f459a6365c106aa1.png" data-download-href="https://discuss.pytorch.org/uploads/default/57a289c4bc9c5bb22f34bfe7f459a6365c106aa1" title="Screenshot 2024-06-24 at 12.53.42 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/5/7/57a289c4bc9c5bb22f34bfe7f459a6365c106aa1_2_690x173.png" alt="Screenshot 2024-06-24 at 12.53.42 PM" data-base62-sha1="cvfM8csveJIA2D8mlgXL20Ez6Wl" width="690" height="173" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/5/7/57a289c4bc9c5bb22f34bfe7f459a6365c106aa1_2_690x173.png, https://discuss.pytorch.org/uploads/default/optimized/3X/5/7/57a289c4bc9c5bb22f34bfe7f459a6365c106aa1_2_1035x259.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/5/7/57a289c4bc9c5bb22f34bfe7f459a6365c106aa1_2_1380x346.png 2x" data-dominant-color="E3E4E2"><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">Screenshot 2024-06-24 at 12.53.42 PM</span><span class="informations">2242×564 88.4 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><br>
<em>Note the bytes field from the NCCL call</em></p>

ptrblck: <p>CUDA kernels will be recorded by default of course, but you can add markers around calls to add more details. Your first example shows how <a href="https://github.com/NVIDIA/Fuser"><code>nvFuser</code></a> adds nvtx markers and you can check the code to see how the backend handles it.</p>

Jonathan1: <p>thanks <a class="mention" href="/u/ptrblck">@ptrblck</a>. I guess what I was looking for was something of an NVTX 101 from existing code that would allow me to reimplement even a subset of host-side traces that PyTorch provides, the <code>nvFuser</code> screenshot I included was from a PyTorch training run. Still, I agree digging through the code of both PyTorch and NCCL, their NVTX instrumentation is also very good, would be the best next step!</p>

ptrblck: <p>You can also just search the linked nvFuser repository for <code>nvtx</code> and see how the markers were implemented.</p>