Implementing multiple recomputations on top of `torch.utils.checkpoint`

username: shirazb
cooked: <p>Hi,</p>
<p>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>I am attempting to implement memory checkpointing as done in <code>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>
<li>
<p><strong>In <code>CheckpointFunction</code>’s backward, why detach the input before recomputing the forward?</strong></p>
<p>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><strong>Why does <code>CheckpointFunction</code>’s backward return <code>(None, None) + grads</code></strong>? Where <code>grads</code> are the grads wrt to the inputs. I really am confused here.</p>
</li>
<li>
<p><strong>CheckpointFunction’s forward produces an output that does not require grad, so surely the next layer’s backward will not be set to propagate to it’s backward?</strong></p>
<p>Say I have a model whose forward performs <code>h(checkpoint(g, f(x)))</code>, which I’ll equivalently write as <code>f -&gt; checkpoint(g) -&gt; h</code>. Say <code>x</code> has <code>requires_grad=True</code>. Then, my understanding of autograd is that something along the lines of the following will occur:</p>
<ol>
<li>
<p><code>f</code>’s forward will see that its input requires grad and thus produce an output that requires grad too (and whose gradient function is <code>f</code>’s backward function). It will save the relevant tensors for the backward and set the backward to propagate the gradient back to its input’s gradient function.</p>
</li>
<li>
<p>The output will be fed into <code>CheckpointFunction</code>’s forward. It will behave similar to as described above, except for this (from the source code):</p>
<blockquote>
<p>with torch.no_grad():<br>
&nbsp; outputs = run_function(*args)<br>
return outputs</p>
</blockquote>
<p>This code means that <code>output</code> will have <code>requires_grad=False</code>, and that no backward gradient graph will have been created.</p>
</li>
<li>
<p>Thus, when this <code>output</code> is fed into <code>h</code>’s forward - even if <code>h</code> contains parameters that require grad meaning the output will require grad too - it will not set up the backward graph to propagate gradient to <code>output</code>’s backward function, which is <code>CheckpointFunction</code>’s backward. <strong>This would mean <code>CheckpointFunction</code>’s backward will not get invoked on a call to <code>h.backward()</code>, and so neither will <code>g.backward()</code> or <code>f.backward()</code></strong>.</p>
</li>
</ol>
</li>
<li>
<p><strong>Implementing multiple recomputations by checkpoint()-ing a model whose child model is checkpoint()-ed too will result in the child saving its input in the first forward pass, not the second (first recomputation pass) as intended.</strong></p>
<p>I was trying to implement multiple recomputations by trivially building on the existing <code>checkpoint</code> function, for example by creating a module that performs <code>1 -&gt; drop(2) -&gt; 3</code>, where the module <code>2</code> itself performs <code>2a -&gt; drop(2b)</code>, where <code>drop</code> is like a higher-order model whose forward simply performs <code>checkpoint(child_model, x)</code>. Thus, running <code>2</code> should drop <code>2a</code> and <code>2b</code> in the forward pass; and in the backward pass, recompute <code>2a</code>, recompute <code>2b</code> and drop it, then recompute <code>2b</code> a second time, this time actually saving it. Obviously, that’s not so smart but its a simple example of multiple recomputations: <em>all of <code>2</code> is dropped in the first forward pass, and when we are recomputing it, <code>2a</code> is checkpointed and <code>2b</code> is dropped again</em>.</p>
<p>Or that is the intent, but I believe the following chain of events will occur in practice:</p>
<ol>
<li>
<p><code>1</code>’s forward is performed, whose output is propagated to <code>drop(2)</code>'s forward function.</p>
</li>
<li>
<p><code>drop(2)</code>'s forward invokes <code>CheckpointFunction</code>’s forward, which saves (checkpoints) the input, and then invokes <code>2</code>’s forward without tracking gradient. Again, the intent is that both the outputs of <code>2a</code> and <code>2b</code> will be dropped at this stage.</p>
</li>
<li>
<p><code>2a</code> performs its forward and the output is passed to <code>drop(2b)</code>'s forward.</p>
</li>
<li>
<p>This invokes <code>CheckpointFunction</code>’s forward <strong>which will save the input to <code>2b</code> thus checkpointing it, which, as mentioned in step 2, I do not want to happen!</strong></p>
<p>On the other hand, maybe it will be freed because of how everything is set up and autograd’s reference counting?</p>
<p>Maybe, as <code>drop(2b)</code> saves its input, it propagates forward with no grad, and <code>2a</code> itself was run with no grad; you get an unreachable reference cycle between them that will get garbage collected?</p>
<p>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><strong>If my anaylsis is right, I will have to implement this <code>drop</code> operator from scratch such that it avoids this behaviour, correct?</strong></p>
</li>
</ol>
<p>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 <em>greatly</em> appreciated.</p>

username: shirazb
cooked: <p>Note it’s been a while since I last looked at this. Also, I think the source for <code>torch.utils.CheckpointFunction</code> might be slightly different now, though hopefully not in a way that affects this question.</p>
<ol>
<li>
<p>Say we didn’t detach the input, when we recompute the output and run the backward, the backward propagation would propagate past the input, as it is still attached to the gradient graph from the prior computations that created it. However, the function we are implementing (<code>CheckpointFunction</code>'s backward) is supposed to be just one operator in that gradient graph, and return just its backward, so we do not want to cause that side effect.</p>
<p>With respect to whether detaching duplicates the underlying buffers, I don’t remember what I found but it should be simple to verify by experimentation (the tensor objects have a pointer to the underlying buffer; make a detached copy and check them for equality).</p>
</li>
<li>
<p>This was actually really simply and took far too long for me to figure out. The API of Autograd is assuming that the backward function returns one gradient for each argument given to the forward function, excluding the first context argument. If we look at the corresponding forward of <code>CheckpointFunction</code>, we see there are three inputs, the first two being ‘pseudo inputs’ that configure the checkpointing rather than being actual tensors, but we still need to account for them when returning the gradients from the backward function, hence <code>(None, None) + grads</code>.</p>
</li>
<li>
<p>I’m afraid I can’t recall how I figured this out. Sorry.</p>
</li>
<li>
<p>I <em>think</em> my analysis here was correct; the tensors to be checkpointed only on the first recomputation were being checkpointing immediately.</p>
<p>To get around this, I added a <code>recomputation_depth</code> parameter to the forward function of <code>CheckpointFunction</code>, and only checkpointed the input if the depth was 0. The forward of the <code>Drop</code> module carries its own <code>recomputation_depth</code> parameter and, in its forward, decrements it before invoking the checkpoint function.</p>
<p>Note that, as explained in <strong>2.</strong>, we now must return <code>(None, None, None) + grads</code> in the backward, due to the additional argument given to the forward.</p>
<p>This successfully allows you to encode multiple recomputations using these <code>Drop</code> modules, for example by doing:</p>
<pre><code class="lang-python">torch.nn.Sequential(
    layer1,
    Drop(torch.nn.Sequential(
        layer2,
        Drop(torch.nn.Sequential(
            layer3,
            layer4,
        2)),          # &lt;----- recursion depth of inner drop
        layer5,
    1                 # &lt;----- recursion depth of outer drop
)
</code></pre>
<p>where layer 1 is not recomputed; 2 and 5 are recomputed once; and 3 and 4 are recomputed twice.</p>
</li>
</ol>
<p>Hope this helps.</p>

username: Dale_Song
cooked: <p>Is there a method to swap in&amp;out data using pytorch? As far as I know, recomputation will result in a bit of computation efficiency loss, so, is it possible using swap in&amp;out instead in pytorch?</p>

username: shirazb
cooked: <p>I haven’t looked much into PyTorch implementations for swapping, but I do recall there may have been some work on implementing something similar to TensorFlow Large Model Support (TFLMS), which was a swapping solution for TensorFlow.</p>
<p>A quick google search gave me this thread: <a href="https://discuss.pytorch.org/t/thoughts-on-use-of-cpu-ram-as-a-swap-for-gpu/13193" class="inline-onebox">Thoughts on use of CPU ram as a swap for GPU</a></p>
<p>Original TFLMS Paper: <a href="https://arxiv.org/pdf/1807.02037.pdf" rel="nofollow noopener">https://arxiv.org/pdf/1807.02037.pdf</a></p>