<p dir="auto">Hi,</p>
<p dir="auto">My masters thesis is on making neural nets use less memory. One technique I am looking at is memory checkpointing. I can solve for the optimal policy (including multiple recomputations), given the memory budget and per-operator compute/memory costs.</p>
<p dir="auto">I am attempting to implement memory checkpointing as done in <code class="notranslate">torch.utils.checkpoint</code>, except allowing for multiple recomputations. However, there a couple things from the implementation that I'm not quite sure I understand. Apologies if anything is obvious, I have been using PyTorch for &lt;2 days.</p>
<ol dir="auto">
<li>
<p dir="auto"><strong>Why detach the inputs in the backward before recomputing the forward?</strong> Would you get something like the gradient of the forward being accumulated twice into the input's gradient? Also, does this detaching not duplicate the input resulting in higher than necessary memory cost?</p>
</li>
<li>
<p dir="auto">I was trying to implement multiple recomputations by trivially building on the existing checkpoint function, for example by creating a module that performs <code class="notranslate">1 -&gt; drop(2) -&gt; 3</code>, where the module <code class="notranslate">2</code> itself performs <code class="notranslate">2a -&gt; drop(2b)</code>. <code class="notranslate">drop</code> is like a higher-order model whose forward simply performs <code class="notranslate">checkpoint(child_model, x)</code>. Thus, running <code class="notranslate">2</code> should drop <code class="notranslate">2a</code> and <code class="notranslate">2b</code> in the forward pass; and in the backward pass, recompute <code class="notranslate">2a</code>, recompute <code class="notranslate">2b</code> and drop it, then recompute <code class="notranslate">2b</code> a second time, this time actually saving it. Obviously, that's not so smart but its a simple example.</p>
<p dir="auto">I believe the following chain of events will occur in practice:</p>
<ol dir="auto">
<li>
<p dir="auto"><code class="notranslate">1</code>'s forward is performed, whose output is propagated to <code class="notranslate">drop(2)</code>'s forward function.</p>
</li>
<li>
<p dir="auto"><code class="notranslate">drop(2)</code>'s forward invokes <code class="notranslate">CheckpointFunction</code>'s forward, which saves (checkpoints) the input, and performs <code class="notranslate">2</code>'s forward without tracking gradient. The intent is that both the outputs of <code class="notranslate">2a</code> and <code class="notranslate">2b</code> will be dropped at this stage.</p>
</li>
<li>
<p dir="auto"><code class="notranslate">2a</code> performs its forward and the output is passed to <code class="notranslate">drop(2b)</code>'s forward.</p>
</li>
<li>
<p dir="auto">a. If I set up this 'drop' wrapper to not detach its input before passing to the child model, the <code class="notranslate">CheckpointFunction</code> forward will throw an error about none of the inputs have <code class="notranslate">requires_grad=True</code>.</p>
<p dir="auto">b. If I do detach the input, I <em>suspect</em> this will not work either: <strong>it will save the input to <code class="notranslate">2b</code> thus checkpointing it, which, as mentioned in step 3, I do not want to happen!</strong></p>
<p dir="auto">On the other hand, maybe it will be freed because of how everything is set up and autograd's reference counting?</p>
<p dir="auto">Maybe, as <code class="notranslate">drop(2b)</code> saves its input, and <code class="notranslate">2a</code> was run with no grad, you get an unreachable cycle between them that will get garbage collected?</p>
<p dir="auto">I really lack the autograd understanding to know (see the &lt;2 days of PyTorch). I also do not know how to profile this to observe if it drops the tensor or not.</p>
</li>
</ol>
<p dir="auto"><strong>If my anaylsis is right, I will have to implement this <code class="notranslate">drop</code> operator from scratch such that it avoids this behaviour, correct?</strong></p>
</li>
</ol>
<p dir="auto">Thank you for making it this far and sorry if the above explanations are not great, it would be easier with diagrams. Any help on this would be greatly appreciated.</p>
<p dir="auto">It looks like THPVariable_invert should be implemented in C++ to avoid duplicate work from a JIT team.</p>
<p dir="auto">For more details, please, see <a href="https://github.com/pytorch/pytorch/pull/21113" data-hovercard-type="pull_request" data-hovercard-url="/pytorch/pytorch/pull/21113/hovercard">PR</a> comments.</p>