DDP implementation with overlapping communication and computation with backward hook

Robin_Ranabhat: <p>Dear community,  I am in grave help with verifying correctness of program flow of a minimal DDP training script ..</p>
<h2><a name="p-474460-environment-1" class="anchor" href="#p-474460-environment-1"></a>Environment :</h2>
<p>Single node, two gpus, PCIE connection (no nvlink)</p>
<h2><a name="p-474460-what-i-did-2" class="anchor" href="#p-474460-what-i-did-2"></a>What I did</h2>
<p>I built upon an excellent DDP tutorial ( <a href="https://docs.pytorch.org/tutorials/intermediate/ddp_tutorial.html" class="inline-onebox" rel="noopener nofollow ugc">Getting Started with Distributed Data Parallel — PyTorch Tutorials 2.8.0+cu128 documentation</a> ) where gradients are all_reduced after <code>loss.backward()</code>.<br>
As version 2 of this, I wanted to use backward hooks to communicate gradients as soon as they were available during backprop using `register_post_accumulate_grad_hook`.</p>
<p>This is my version 2 : <a href="https://github.com/robinnarsinghranabhat/pytorch-optimizations-notes/blob/main/ddp-from-scratch/ddp.py" class="inline-onebox" rel="noopener nofollow ugc">pytorch-optimizations-notes/ddp-from-scratch/ddp.py at main · robinnarsinghranabhat/pytorch-optimizations-notes · GitHub</a></p>
<h2><a name="p-474460-question-3" class="anchor" href="#p-474460-question-3"></a>Question :</h2>
<p>Both version 1 and version 2 generate same gradients and final loss, so I know logic is correct.<br>
But I can’t judge if my logic for overlapping communication and computation using streams is okay.</p>
<p>I profiled it here : <a href="https://github.com/robinnarsinghranabhat/pytorch-optimizations-notes/blob/main/ddp-from-scratch/ddp_profiled.py" class="inline-onebox" rel="noopener nofollow ugc">pytorch-optimizations-notes/ddp-from-scratch/ddp_profiled.py at main · robinnarsinghranabhat/pytorch-optimizations-notes · GitHub</a><br>
Frankly. I am  trying to understand the chrome trace results. can’t tell if even the profiling is done correctly. More than happy to share tracing results.</p>

Robin_Ranabhat: <p>I have debugged the solution. Have posted this ddp implementation here along with a brief write up.</p>
<aside class="onebox githubfolder" data-onebox-src="https://github.com/robinnarsinghranabhat/pytorch-optimizations-notes/tree/main/03.%20ddp-training-from-scratch">
  <header class="source">
      <img src="https://github.githubassets.com/favicons/favicon.svg" class="site-icon" alt="" width="32" height="32">

      <a href="https://github.com/robinnarsinghranabhat/pytorch-optimizations-notes/tree/main/03.%20ddp-training-from-scratch" target="_blank" rel="noopener nofollow ugc">github.com</a>
  </header>

  <article class="onebox-body">
    <h3><a href="https://github.com/robinnarsinghranabhat/pytorch-optimizations-notes/tree/main/03.%20ddp-training-from-scratch" target="_blank" rel="noopener nofollow ugc">pytorch-optimizations-notes/03. ddp-training-from-scratch at main ·...</a></h3>


  <p><span class="label1">Guidelines to Improve performance of custom pytorch modules - robinnarsinghranabhat/pytorch-optimizations-notes</span></p>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>