FSDP Backward pre-fetch documentation

david-macleod: <p>The documentation which related to <code>torch.distributed.fsdp.BackwardPrefetch</code> strategy seems to imply that the memory allocated for gradients is freed before the ReduceScatter phase (or I am just not understanding)</p>
<p>I have drawn what I believe to be an example of the strategies. CUDA stream 1 is responsible for the Backward gradient computation kernels, and CUDA stream 2 runs the AllGather and ReduceScatter collectives. I have also drawn what I believe to be the duration of the P (parameter) allocations and G (gradient) allocations, so illustrate when they are freed.</p>
<p><code>BACKWARD_PRE</code><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/f/bf46d7b46be930c7b4f981f65c1487acd288c57c.png" data-download-href="https://discuss.pytorch.org/uploads/default/bf46d7b46be930c7b4f981f65c1487acd288c57c" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/f/bf46d7b46be930c7b4f981f65c1487acd288c57c_2_690x300.png" alt="image" data-base62-sha1="ri6Yxm4bFJ9hOw3BnSPI9bEOk1u" width="690" height="300" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/f/bf46d7b46be930c7b4f981f65c1487acd288c57c_2_690x300.png, https://discuss.pytorch.org/uploads/default/optimized/3X/b/f/bf46d7b46be930c7b4f981f65c1487acd288c57c_2_1035x450.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/b/f/bf46d7b46be930c7b4f981f65c1487acd288c57c.png 2x" data-dominant-color="1B1B1C"><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">1101×480 24.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><code>BACKWARD_POST</code><br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/8/3/83550cb8da54a089dbee1860fb7c3f19030ff84a.png" data-download-href="https://discuss.pytorch.org/uploads/default/83550cb8da54a089dbee1860fb7c3f19030ff84a" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/8/3/83550cb8da54a089dbee1860fb7c3f19030ff84a_2_690x249.png" alt="image" data-base62-sha1="iJOLxqH66P0NYRQxy3c5bQaoFBg" width="690" height="249" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/8/3/83550cb8da54a089dbee1860fb7c3f19030ff84a_2_690x249.png, https://discuss.pytorch.org/uploads/default/optimized/3X/8/3/83550cb8da54a089dbee1860fb7c3f19030ff84a_2_1035x373.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/8/3/83550cb8da54a089dbee1860fb7c3f19030ff84a.png 2x" data-dominant-color="1A1B1B"><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">1273×461 24 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>The documentation claims that <code>BACKWARD_POST</code> saves memory because peak memory has resident the</p>
<ul>
<li>The next set of parameters</li>
<li>The current set of gradients in memory</li>
</ul>
<p>But as we are overlapping ReduceScatter with next gradient calculation my intuition would be that at some point we must have</p>
<ul>
<li>The current parameters</li>
<li>The current gradients (mid-computation)</li>
<li>The previous gradients (mid-reduction)</li>
</ul>
<p>As can be seen on the diagram. I would appreciate anyone letting me know what I am missing here!</p>

agu: <p>I think actually <code>BACKWARD_POST</code> does not really save peak memory in a principled way. I think <code>BACKWARD_POST</code> is actually just prefetching incorrectly when transitioning from one module depth to another (depth when viewing the model as a tree of modules). <a href="https://github.com/pytorch/pytorch/issues/108190" class="inline-onebox" rel="noopener nofollow ugc">[FSDP] incorrect backward prefetch order when using BackwardPrefetch.BACKWARD_POST · Issue #108190 · pytorch/pytorch · GitHub</a></p>