Gradient checkpointing

AlbertZeyer: <p>I wonder how <a href="https://pytorch.org/docs/stable/checkpoint.html" rel="noopener nofollow ugc"><code>torch.utils.checkpoint</code></a> is supposed to be used in practice. The outputs of the checkpointing function are potentially anyway stored elsewhere for backprop, as those are not under the <code>checkpoint</code> anymore. E.g. consider the example:</p>
<pre data-code-wrap="python"><code class="lang-python">y = (a + b) * c
</code></pre>
<p>And now <code>checkpoint</code> around the <code>a + b</code>:</p>
<pre data-code-wrap="python"><code class="lang-python">a = ...
b = ...
c = ...
x = checkpoint(lambda (_a, _b): _a + _b, a, b)
y = x * c
</code></pre>
<p>My intention is that the result of <code>a + b</code> is not stored in memory but recomputed. However, in this example, it doesn’t really make sense to recompute <code>x = a + b</code> here, because the result is anyway stored for the backprop of <code>x * c</code>. Or is the intended usage of <code>checkpoint</code> for this example actually like:</p>
<pre data-code-wrap="python"><code class="lang-python">a = ...
b = ...
c = ...
y = checkpoint(lambda (_a, _b, _c): (_a + _b) * c, a, b, c)
</code></pre>
<p>?</p>
<p>Another question: Are the explicit args to <code>checkpoint</code> actually necessary? Or can I just use this instead:</p>
<pre data-code-wrap="python"><code class="lang-python">a = ...
b = ...
c = ...
y = checkpoint(lambda: (a + b) * c)
</code></pre>
<p>?<br>
In real world, e.g. any model parameters are probably not expressed as explicit inputs but they are used within the function.</p>

soulitzer: <aside class="quote no-group" data-username="AlbertZeyer" data-post="1" data-topic="205416">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/albertzeyer/48/45362_2.png" class="avatar"> AlbertZeyer:</div>
<blockquote>
<p>The outputs of the checkpointing function are potentially anyway stored elsewhere for backprop</p>
</blockquote>
</aside>
<p>This is correct, the idea is that you’d use checkpointing on large enough regions such that the activations saved at the boundaries are relatively a small proportion of overall activations.</p>
<blockquote>
<p>Are the explicit args to <code>checkpoint</code> actually necessary? Or can I just use this instead:</p>
</blockquote>
<p>Yes, that is supported. There’s a subtle difference in when closed over buffers are cleared, but that shouldn’t matter for basic usage, i.e., forward then backward once.</p>

AlbertZeyer: <aside class="quote no-group" data-username="soulitzer" data-post="2" data-topic="205416">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/839c29/48.png" class="avatar"> soulitzer:</div>
<blockquote>
<p>the idea is that you’d use checkpointing on large enough regions such that the activations saved at the boundaries are relatively a small proportion of overall activations.</p>
</blockquote>
</aside>
<p>Ok, so for my example, if I want that <code>a + b</code> is not stored but recomputed, I would need to use it like this:</p>
<pre data-code-wrap="python"><code class="lang-python">a = ...
b = ...
c = ...
y = checkpoint(lambda: (a + b) * c)
</code></pre>
<p>In this example, is the multiplication <code>... * c</code> also recomputed then or not? Because I actually do <em>not</em> want to recompute that. In practice, this multiplication is likely a big costly matmul. I know exactly that I only want to recompute the <code>a + b</code> but nothing else. I fear that the whole API of <code>checkpoint</code> might do more recomputations than what I actually want, and/or not those recomputations that I actually want.</p>
<p>As I understand the code and documentation, <code>use_reentrant=False</code> is already more like what I want. For my simple example, I think it would not recompute the multiplication <code>...  * c</code>, right? But how would it do that actually? Looking at the code, I see that it calls <code>recompute_fn</code> once it needs some of the activations in the backward pass. But <code>recompute_fn</code> will just call the original function <code>fn</code> at some point, which includes the computation of the multiplication? <strong>Edit</strong> Ah, now I see the logic with <code>_StopRecomputationError</code>. This looks all quite hacky to me…</p>

soulitzer: <p>Yes, it would not be recomputed with use_reentrant=False via _StopRecomputationError. use_reentrant=True does not have this logic so the entire forward is always recomputed in that path.</p>

AlbertZeyer: <p>So I think my original questions are answered then. But I wonder: In my case, and in this example, I know, I want to recompute <code>x = a + b</code>, but nothing else, whatever other computations are coming afterwards (which I maybe don’t even have control over, some other unrelated code). The current <code>checkpoint</code> API does not really allow me to specify that. Any code which comes after it must also be put into the <code>checkpoint</code> function, but only the part which uses this <code>x</code> directly.</p>
<p>Is this unusual that I only want to recompute <code>a + b</code> and nothing else? Basically, with the current <code>checkpoint</code> API, I specify the points (tensors) which are stored in memory. However, I want an API where I specify the tensors which are not stored.</p>

soulitzer: <p>That’s sounds like something reasonable to do. However, unfortunately there’s no automatic way to apply checkpoint to the consumers of this output.</p>
<p>The closest API there is to this today is selective activation checkpoint <a href="https://pytorch.org/docs/main/checkpoint.html#torch.utils.checkpoint.create_selective_checkpoint_contexts" class="inline-onebox" rel="noopener nofollow ugc">torch.utils.checkpoint — PyTorch main documentation</a> (landed very recently, available in nightlies or if you build from source).</p>
<p>With that you can pass in a policy of “if my op, recompute, otherwise, save”. However it will save the output of ALL ops in the region except the op that you specified, not just the the buffers that are needed for backward computation.</p>

AlbertZeyer: <p>I wonder, if I use <code>torch.autograd.graph.saved_tensors_hooks</code> around my whole model, then I can exactly decide what I want to store and what to recompute. How much performance penalty would I get by using <code>saved_tensors_hooks</code> all the time?</p>
<p>I also had the idea to delay <code>saved_tensors_hooks.__exit__</code> to some later point in <code>checkpoint</code>, specifically to register some <code>Tensor.__del__</code> hook (e.g. like <a href="https://github.com/albertz/playground/blob/master/torch-tensor-track.py" rel="noopener nofollow ugc">this</a>) on the outputs of <code>fn</code> and once this hook gets called, then do the <code>saved_tensors_hooks.__exit__</code>. But there are many further details to make this work: E.g. I need to check <code>new_frame.forward_completed</code>, and any newly created tensors should be stored, just not those which are created so far. This is probably not really possible in a clean way.</p>
<p>Also, in <code>saved_tensors_hooks</code>, I somehow need to check whether a tensor depends on some of the computation from the <code>fn</code> or not - not sure how to do this.</p>

soulitzer: <blockquote>
<p>How much performance penalty would I get by using <code>saved_tensors_hooks</code> all the time?</p>
</blockquote>
<p>Good question. You do always incur an additional python round trip. But, whether this matters really depends on whether your model is overhead bound.</p>
<blockquote>
<p>I wonder, if I use <code>torch.autograd.graph.saved_tensors_hooks</code> around my whole model</p>
</blockquote>
<p>This is a good idea, and I have explored something similar in the past - the idea is that I record the dependency relations between what exactly is needed to compute during backward via a RecomputableTensor object in order to only recompute precisely what is necessary.</p>
<p>The idea is that I define this new RecomputableTensor which doesn’t actually hold data but knows how to compute itself based on other RecomputableTensors. When I save for backward, I use saved tensor hooks to save RecomputableTensor instead of plain tensor. So let’s say during backward, I want a tensor saved for backward needed for the first gradient computed, it would recursively call recompute until it sees a RecomputableTensor that has already been materialized and the way selective checkpoint would work is that we would make sure RecomputableTensors generated from those specific ops are already materialized, and so that recursion stops early.</p>
<p>It should also be able to save precisely what is needed, the RecomputableTensors are kept alive in 3 ways (1) the original tensors (2) other recomputable tensors (3) saved onto the graph. If a saved tensor is not kept alive by a chain of recomputable tensors rooted in a object saved onto the graph, then it would have no references keeping it alive as soon as the original tensor dies.</p>

soulitzer: <aside class="onebox githubgist" data-onebox-src="https://gist.github.com/soulitzer/ec1049a947be046de7fbc2af61a4ee8c">
  <header class="source">

      <a href="https://gist.github.com/soulitzer/ec1049a947be046de7fbc2af61a4ee8c" target="_blank" rel="noopener nofollow ugc">gist.github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://gist.github.com/soulitzer/ec1049a947be046de7fbc2af61a4ee8c" target="_blank" rel="noopener nofollow ugc">https://gist.github.com/soulitzer/ec1049a947be046de7fbc2af61a4ee8c</a></h4>

  <h5>sac2.py</h5>
  <pre><code class="Python">import torch
import functools
from torch.utils._python_dispatch import TorchDispatchMode
import torch.utils._pytree as pytree
from torch.utils.weak import WeakTensorKeyDictionary


class RecomputableTensor(torch.Tensor):
    @staticmethod
    def __new__(cls, t, func, args):</code></pre>
   This file has been truncated. <a href="https://gist.github.com/soulitzer/ec1049a947be046de7fbc2af61a4ee8c" target="_blank" rel="noopener nofollow ugc">show original</a>

<p>
</p>

  </article>

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

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


AlbertZeyer: <p>Oh right, I forgot about <code>__torch_dispatch__</code>. (I lack a bit the experience on what I can hock/dispatch.)</p>
<p>Yea, via this mechanism, I could exactly do what I need. That seems to be great!</p>
<p>Note, in <a href="https://github.com/rwth-i6/returnn" rel="noopener nofollow ugc">RETURNN</a>, for the TensorFlow backend, I have implemented <a href="https://github.com/rwth-i6/returnn/blob/master/returnn/tf/util/gradient_checkpoint.py" rel="noopener nofollow ugc">this API for gradient checkpointing</a>, where you can do sth like this:</p>
<pre><code class="lang-auto">with gradient_checkpoint_scope():
    x = a + b
y = x * c
</code></pre>
<p>All tensors in <code>gradient_checkpoint_scope</code> will get specially marked. This is using TF graph mode. I just need to record the first node/op id and the last. And then I can iterate through all ops from the scope like this:</p>
<pre data-code-wrap="python"><code class="lang-python">for op_id in range(first_op_id, last_op_id):
    op = graph._nodes_by_id[op_id]
    ...
</code></pre>
<p>And then I overwrite the grad func of all consumers:</p>
<pre data-code-wrap="python"><code class="lang-python">for op_out in op.outputs:
    assert isinstance(op_out, tf.Tensor)
    for op_ in op_out.consumers():
        _set_wrapped_grad_func(op_)
</code></pre>
<p>This is before I calculate the gradient (<code>tf.gradients</code>, similar as Torch <code>tensor.backward()</code>).<br>
In my wrapped gradient function, instead of passing any of the original tensors from that <code>gradient_checkpoint_scope</code> to the original gradient function, I copy all the computation graph under this scope and pass the copied tensors. I don’t really need to care about adding some unnecessary computations to the computation graph, as TensorFlow only lazily calculates whatever is needed in the end.</p>
<p>So, I wonder if I can implement the same logic in PyTorch, i.e. the same <code>gradient_checkpoint_scope</code> API. Very similar to your example, I could use <code>__torch_dispatch__</code> to record the computation graph and also record all tensors which are created there. If I would use <code>saved_tensors_hooks</code> on the whole model, I can check whenever some of the tensors come from such checkpoint scope, and if so, recompute that.</p>
<p>If I don’t want to have <code>saved_tensors_hooks</code> on the whole model, I can solve this too: I can <code>saved_tensors_hooks.__enter__</code> when I enter the checkpoint scope, but don’t immediately <code>saved_tensors_hooks.__exit__</code> afterwards. Instead, I can hook into <code>Tensor.__del__</code> (as described above) for all tensors created in that scope, and once they all get deleted, then I can do <code>saved_tensors_hooks.__exit__</code> (and also exit the <code>__torch_dispatch__</code>). So this stays very local then.</p>
<p>It gets a bit trickier to handle the RNG state correctly, and maybe also handle AMP. I guess I need to make sure to store and then later fork and reset the state correctly, and also to replay the recorded computations exactly in the same order to make sure it’s all deterministic. But this is similar to the current Torch <code>checkpoint</code> logic.</p>
<p>So, this sounds like a plan…</p>

soulitzer: <p>Yes, torch dispatch is a very useful tool. Let me know how things go, if you decide to work on it more.</p>
<p>RNG is indeed tricky! There’s quite a few things that don’t work well with the prototype, but unfortunately I don’t have too much time to work it these days.</p>