Making autograd saved tensors hooks specific to certain arguments

mseeger: <p>Hello,</p>
<p>I am trying to save memory on a certain forward-backward computation, by realizing that a whole sequence of large tensors can be reconstructed from smaller tensors. Think <code>x_{k+1} = f(x_k, a_k), x_k = g(x_{k+1}, b_k)</code>, where <code>x_k</code> are large, <code>b_k</code> are smaller, so I really only need to store the final <code>x_K</code> and all the <code>b_k</code> to reconstruct all other <code>x_k</code>.</p>
<p>I found:</p>
<p><a href="https://pytorch.org/tutorials/intermediate/autograd_saved_tensors_hooks_tutorial.html" class="onebox" target="_blank" rel="noopener nofollow ugc">https://pytorch.org/tutorials/intermediate/autograd_saved_tensors_hooks_tutorial.html</a></p>
<p>I thought this would nail it, but I am struggling now. This is because <code>pack_hook</code> only just gets the (detached) tensor, and I see no way to identify whether it is one of the <code>x_k</code>. While every new node in the graph gets a unique name, this is not passed to <code>pack_hook</code>.</p>
<p>I can match the tensors by shape, but I have a number of these sequences, and I need to be able to tell them apart. And for many arguments of <code>pack_hook</code>, I just want to pass them through.</p>
<p>Anybody has an idea here? Can you tell me why these hooks (which are just awesome) do not also obtain information that allows me to identify the node they came from? Not doing this restricts people tyo blunt use cases such as “store all tensors to CPU / disk”, while in general, we’d want to be selective to certain tensors and not care about others.</p>

mseeger: <p>To add details. My current implementation creates certain information for nodes created in the forward pass, for example <code>b_k</code> for <code>x_k</code>. Ideally, I’d now collect these in a dictionary, where the key allows me to identify <code>x_k</code>. And <code>pack_hook</code> would use this, and (say) store <code>b_k</code> instead of <code>x_k</code> in the graph. I am missing an idea for what this key would be so that <code>pack_hook</code> can test whether its input argument is one of them and retrieve the information.</p>
<p>So, if it is <code>def pack_hook(x): ...</code>, I’d need to be able to test whether <code>x</code> corresponds to one of the <code>x_k</code> in the dictionary.</p>

soulitzer: <blockquote>
<p>Can you tell me why these hooks (which are just awesome) do not also obtain information that allows me to identify the node they came from?</p>
</blockquote>
<p>No fundamental reason, this is a reasonable request, but if I understand correctly you’re trying to use it to implement selective saving/recompute. I’m not sure it would really help you in that case. Are you planning to just rerun the user-passed function again during recompute?</p>
<p>The selective storing/saving can be done by nesting hook contexts, e.g. having a inner no-op context over the parts of forward you don’t want saved tensors to be stashed on cpu.</p>

soulitzer: <blockquote>
<p>we’d want to be selective to certain tensors and not care about others.</p>
</blockquote>
<p>There’s an API for selective activation checkpointing <a href="https://pytorch.org/blog/activation-checkpointing-techniques/" class="inline-onebox" rel="noopener nofollow ugc">Current and New Activation Checkpointing Techniques in PyTorch | PyTorch</a>.</p>
<p>It doesn’t behave optimally in certain ways, but mostly does what you want in the common case (save all matmul/compute intensive ops). It hasn’t really been tested much with nesting/recursive checkpoint however.</p>

mseeger: <p>Hello, I am very grateful for your help. I’d in particular be interested what you mean with “inner no-op context” and “user-passed function”. I am probably not aware of some tricks here.</p>
<p>Let me explain where my problem is. I’d like to implement gradient computation in the presence of KV caches. Doing this directly is totally out of the question, would require way too much memory. I’ll be using a form of nested activation checkpointing (which I coded up without <code>torch.utils.checkpoint</code>, it was not that hard). But there is another property I need to use here, namely that the KV cache buffers along the “prompt processing” axis (this is a sequential direction, because the prompt is too large to be processed in a single prefill operation – that is the whole point) have a very useful property: <code>x_k = g(x_{k+1}, b_k)</code>. Here, <code>x_k</code> are the KV cache buffers for successive chunks of the prompt being processed, and <code>b_k</code> is much smaller than <code>x_k</code>. This is because the number of tokens in each chunk is much smaller than the cache length.</p>
<p>Knowing this, I can reconstruct the whole <code>x_k</code> sequence from the <code>b_k</code> and the final <code>x_K</code>, and this saves a lot of memory. In a sense, I can “pack” <code>b_k</code> and some index info instead of <code>x_k</code>.</p>
<p>To make this work, I need to:</p>
<ul>
<li>Store the <code>b_k</code> along the forward pass. I call this “annotation”</li>
<li>Pack hook: Recognize the input argument is <code>x_k</code> (or is directly related, by reshape etc), and pack <code>b_k</code> plus some info instead of <code>x_k</code>.</li>
<li>Unpack hook: Reconstruct <code>x_k</code> from the packed information</li>
</ul>
<p>My problem is how do I detect the right annotation in the pack hook. Namely, most calls of the pack hook, I’d not do anything. I just need to identify those that correspond to the <code>x_k</code>.</p>

mseeger: <p>For now, I am implementing something pretty brittle, which identifies the <code>x_k</code> by shape and also by some entries picked at random locations (so, a kind of a hash, stored at annotation). But I’d love to do something more solid, really.</p>

soulitzer: <p>Thanks for the context.</p>
<p>Have you considered using something like</p>
<pre><code class="lang-auto">from torch.utils.weak import WeakTensorKeyDictionary
</code></pre>
<p>to identify <code>x_k</code>?</p>
<p>No-op context means a nesting a set of saved tensor hooks where the pack/unpack are <code>lambda x: x</code><br>
This means that the outer saved tensor hooks context is effectively applied selectively because autograd only sees the top-most saved tensor hooks.</p>
<p>User-passed function is the function you pass to checkpoint, e.g. the <code>fn</code> arg in <code>torch.utils.checkpoint(fn, x)</code>. Part of the issue here is that if your implementation of AC needs to recompute this function from the very beginning, selectively choosing to save not-save using hooks doesn’t actually save you any compute.</p>
<p>Tangentially related to your question, but can I ask how useful recursive/nesting checkpoint is important for your use case?</p>
<p>I’m working on a new version of AC that should make it very trivial to do a “recompute particular tensor” operation. <a href="https://gist.github.com/soulitzer/7472bc97b909474e77617f3bd474406c" class="inline-onebox" rel="noopener nofollow ugc">graph-based AC · GitHub</a><br>
but, it doesn’t supporting nesting AC (by nesting AC I mean as in possibly recomputing the same forward op more than once) and it seems like adding support may add a good amount of complexity.</p>

mseeger: <p>Hello, I need nested checkpointing because if I did one level checkpointing, the checkpoint sizes would be too large (I’d like to keep the checkpoints on the CPU).</p>

mseeger: <p>I managed to get this working, essentially by storing a small part of the tensor content and matching against that. This is not great at all, but it works.</p>
<p>I still think that the autograd saved tensor hooks feature (which is awesome) would be much more valuable if I could uniquely identify a tensor in the unpack hook that I visited in the forward loop before.</p>