Does PyTorch compute the gradients for all nodes at the same level in parallel?

OldAlbumin: <p>Assume we have a simple network represented by the computation graph below:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/5/6/56df45f9fc398cf9324ab72fe93212e8706449c1.png" data-download-href="https://discuss.pytorch.org/uploads/default/56df45f9fc398cf9324ab72fe93212e8706449c1" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/5/6/56df45f9fc398cf9324ab72fe93212e8706449c1_2_467x500.png" alt="image" data-base62-sha1="covqfx2pg5Wp4Yfq7OY2Lf8do5P" width="467" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/5/6/56df45f9fc398cf9324ab72fe93212e8706449c1_2_467x500.png, https://discuss.pytorch.org/uploads/default/original/3X/5/6/56df45f9fc398cf9324ab72fe93212e8706449c1.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/5/6/56df45f9fc398cf9324ab72fe93212e8706449c1.png 2x" data-dominant-color="FAFDF6"><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">697×745 38.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>Here, M1 through M4 are all <code>nn.Module</code>, say each being</p>
<pre><code class="lang-auto">nn.Sequential(
  nn.Linear(dim1, dim2),
  nn.ReLU(),
  nn.Linear(dim2, dim3)
)
</code></pre>
<p><strong>Q1.</strong> Since M1 through M4 are independent and at the same level in computation graph, <strong>will <code>backward()</code> compute the gradients for them in parallel?</strong> (Anything else will involve sequential operations and be suboptimal.) Please share a reference from original docs if they mention this.</p>
<p><strong>Q2.</strong> What happens if there are too many nodes at the same level and for some compute / memory constraints a parallel <code>backward()</code> cannot be executed?</p>

ptrblck: <p>Yes, I believe it will run in parallel introduced in <a href="https://github.com/pytorch/pytorch/commit/34ce58c909543632a0ba41791c9a83f29580b23b">this PR</a> by <a class="mention" href="/u/colesbury">@colesbury</a> in 2017 but he can correct me if that’s not the case.</p>

colesbury: <p>I’m not sure of the current state of things. Probably <a class="mention" href="/u/ezyang">@ezyang</a> or <a class="mention" href="/u/alband">@albanD</a> know.</p>
<p>In 2017, the parallel execution was only across devices, not for individual modules on the same device.</p>

ezyang: <p>It is still single threaded per device.</p>

OldAlbumin: <p>Thanks.</p>
<p>I’m wondering why does this not make PyTorch very slow (i.e. processing modules at the same level sequentially)? Shouldn’t parallelizing this give a big speedup?</p>

albanD: <p>It depends a lot on how big the modules are. The most common use cases have very large modules that already saturate a single device fully when they run. So there is no opportunity to run multiple modules at the same time on the same device.</p>
<p>This is why we parallelize on different devices but not for a single device and this is not slowing down the common use case.</p>