Memory used by `autograd` when `torch.scatter` is involved

mseeger: <p>I am trying to implement gradient computation (fine-tuning) for large context LLMs in the presence of sparse inference and KV caching. In this context, the following arises.</p>
<p>Denote KV cache buffer content by <code>cache_t</code> which has (simplified) shape <code>(batch, cache_len)</code>, where <code>cache_len</code> is the number of slots in the cache, and <code>batch</code> is everything else, e.g. batch size, embedding dimension. A large context is processed sequentially, which means that during multi-head attention (MHA), <code>cache_t</code> is mapped to <code>cache_(t+1)</code> by writing KV content for new tokens into the cache, overwriting previous content. This is done as</p>
<pre><code class="lang-auto">cache_(t+1) = torch.scatter(cache_t, -1, index, new_kv)
</code></pre>
<p>where <code>index</code> maps to positions to be overwritten, and <code>new_kv</code> is the new content.</p>
<p>The new content <code>new_kv</code> has shape <code>(batch, new_len)</code>, where <code>new_len</code> is much smaller than <code>cache_len</code>. And then, <code>cache_t</code> takes part in MLA computations being the key and value tensors the query tensor for the t-th sequential step is interacting with. If you know token-by-token generation, this would correspond to <code>new_len = 1</code>, but when the whole input is given (during training), we can choose <code>new_len &gt; 1</code>.</p>
<p>My question is: If the computation graph looks like this, what memory does <code>autograd</code> store during the forward pass in training mode? Say there are <code>T</code> sequential steps.</p>
<ul>
<li>It could be simple, but wasteful: <code>autograd</code> stores all <code>cache_t</code>, so <code>T</code> times <code>batch * cache_len</code>.</li>
<li>It could be a little more difficult, but use less memory: <code>autograd</code> stores the first <code>cache_0</code>, but then only two tensors of size <code>batch * new_len</code> for each <code>t</code>, namely <code>new_kv</code> and the content that is overwritten, so <code>torch.gather(cache_t, -1, index)</code>. Obviously, I can obtain <code>cache_(t+1)</code> from <code>cache_t</code> if I know these two. This would be one time <code>batch * cache_len</code> plus <code>T</code> times <code>2 * batch * new_len</code>.</li>
</ul>
<p>Can somebody help me?</p>

soulitzer: <p>scatter is linear wrt cache_t so that does not need to be saved for backward by autograd</p>
<p>unrelated question: isn’t KV caching something you’d typically use inference. Interesting that you need autograd here?</p>

mseeger: <p>Thanks. So, even if the graph uses these <code>cache_t</code> in further linear operations, they’ll not be stored? With further linear, I mean that <code>cache_t</code> contains keys and values, as input into MHA, so keys is linearly combined from some queries (much smaller), and values is combined with attention weights (also much smaller).</p>
<p>This is great.</p>

mseeger: <p>I am trying to figure out how to differentiate through a certain version (with some blocked gradients) of the computation graph, where KV caches are used inside. The main issue is memory usage, due to forward storing intermediate activations. This is why I am asking.</p>
<p>In the end, you need to fine-tune models with large context widths, right? There should be solutions to this already, but I am trying to find something more generic.</p>

soulitzer: <blockquote>
<p>So, even if the graph uses these <code>cache_t</code> in further linear operations,</p>
</blockquote>
<p>ahh not exactly, for these bilinear types of operations like linear/conv in order to compute grads wrt the weight, you still need to save the activations.</p>
<blockquote>
<p>There should be solutions to this already, but I am trying to find something more generic.</p>
</blockquote>
<p>Generally for activation memory, there’s <a href="https://pytorch.org/docs/stable/checkpoint.html" class="inline-onebox" rel="noopener nofollow ugc">torch.utils.checkpoint — PyTorch 2.6 documentation</a> btw</p>

mseeger: <p>Thanks. I am already planning to use activation checkpointing. But if all the <code>cache_t</code> are going to be stored, this may still be too much.</p>

mseeger: <p>Can you think of a way to figure this out? I have the code. Is there some hook mechanism which allows me to find out which auxiliary tensors are created in the training mode forward pass?</p>

mseeger: <p>Hello, I found <a href="https://pytorch.org/tutorials/intermediate/autograd_saved_tensors_hooks_tutorial.html" class="inline-onebox" rel="noopener nofollow ugc">Hooks for autograd saved tensors — PyTorch Tutorials 2.6.0+cu124 documentation</a>. With this, I see a (low-level) way to do what I want to do.</p>
<p>But I am really struggling understanding how <code>torch.utils.checkpoint</code> works. I am facing an instance of nested checkpointing.</p>
<p>Do you happen to know of any complete worked-through non-trivial example of nested checkpointing done with <code>torch.utils.checkpoint</code>? Thanks!</p>

mseeger: <p>I guess I’ll just do this not using <code>torch.utils.checkpoint</code>, but chopping the graph into pieces by hand. I don’t quite get how people can spend a lot of work implementing something complex like this and then don’t bother to link to a non-trivial example.</p>

soulitzer: <p>Hi, sorry for the late reply - there’s a comment block here  <a href="https://github.com/pytorch/pytorch/blob/99c9a31386c5da2fe47c3755f71ade5ebb9615eb/torch/utils/checkpoint.py#L603" class="inline-onebox" rel="noopener nofollow ugc">pytorch/torch/utils/checkpoint.py at 99c9a31386c5da2fe47c3755f71ade5ebb9615eb · pytorch/pytorch · GitHub</a></p>
<p>Happy to answer any questions on that. Also if you share more details about your case, I can provide a potentially more directly relevant example.</p>