Free some saved tensors after partial backward

Neirda: <p>Hello,</p>
<p>I am trying to compute a partial backward with respect to the inputs, and then a second one with respect to the parameters. The easy solution is to do:</p>
<pre data-code-wrap="py"><code class="lang-py">loss.backward(inputs=[inputs], retain_graph=True)
...
loss.backward(inputs=list(model.parameters()), retain_graph=False)
</code></pre>
<p>However, as I understand, this solution has 2 issues.</p>
<ul>
<li>1 - It recomputes all intermediates gradients for the second backward pass.  For instance, in this graph, [Sum, Pow, Sqrt, Cos] nodes will be recomputed.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/b/bbf6e39f3ea595f48225fb02741a0d7b481ac843.png" data-download-href="https://discuss.pytorch.org/uploads/default/bbf6e39f3ea595f48225fb02741a0d7b481ac843" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/b/bbf6e39f3ea595f48225fb02741a0d7b481ac843_2_211x500.png" alt="image" data-base62-sha1="qOOf9ZD5asuvxDD3f168TS9r41Z" width="211" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/b/bbf6e39f3ea595f48225fb02741a0d7b481ac843_2_211x500.png, https://discuss.pytorch.org/uploads/default/original/3X/b/b/bbf6e39f3ea595f48225fb02741a0d7b481ac843.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/b/b/bbf6e39f3ea595f48225fb02741a0d7b481ac843.png 2x" data-dominant-color="ECEEEC"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">285×674 16.6 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
I was able to find a workaround for this using hooks on grad_fn by saving the grads for some nodes and starting the second backward from there. However, doing this still necessitates having <code>retain_graph=True</code> otherwise some necessary saved_tensors are freed. This leads to the second problem,</li>
<li>2 - The entire graph and all corresponding saved tensors are kept in memory between the first and second backward pass. In this slightly more complex graph, the only nodes we need to compute the second backward are the ones circled in red, from MmBackward nodes to the weights. The tensor <code>result</code> saved for ReluBackward could be deleted.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/3/d3c413ce51561eadd1462858656442f81417db1e.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/d3c413ce51561eadd1462858656442f81417db1e" title="annotated-bgraph"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/3/d3c413ce51561eadd1462858656442f81417db1e_2_359x500.jpeg" alt="annotated-bgraph" data-base62-sha1="udmQ5hZwz2dSulDZ6RrlWP8poTk" width="359" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/3/d3c413ce51561eadd1462858656442f81417db1e_2_359x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/d/3/d3c413ce51561eadd1462858656442f81417db1e_2_538x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/d/3/d3c413ce51561eadd1462858656442f81417db1e_2_718x1000.jpeg 2x" data-dominant-color="F1ECE4"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">annotated-bgraph</span><span class="informations">1328×1847 199 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></li>
</ul>
<p>On bigger examples the amount of memory saved between the 2 backward passes could be huge, so I would like to know if this is possible. The problem is, since we keep the entire graph and it is not editable from the torch API, there is no way (that I know of) to manually delete some autograd nodes or saved tensors.</p>
<p>Does anyone know if there is a way to partly delete the backward graph while only keeping the nodes we want, with their associated saved tensors?</p>

soulitzer: <p>Instead of using hooks, a slightly different approach is:</p>
<ol>
<li>
<p>You can use the inputs= argument to .grad or .backward to get autograd to return the intermediate grad_outputs.</p>
</li>
<li>
<p>You can use GradientEdge <a href="https://github.com/pytorch/pytorch/pull/127766" class="inline-onebox" rel="noopener nofollow ugc">[autograd] Support GradientEdge as output for torch.autograd.grad by soulitzer · Pull Request #127766 · pytorch/pytorch · GitHub</a> to backward starting from an intermediate node without needing to keep the corresponding tensor alive.</p>
</li>
</ol>
<blockquote>
<p>The entire graph and all corresponding saved tensors are kept in memory between the first and second backward pass.</p>
</blockquote>
<p>The issue here is that in such cases where the forward and backward are disjoint is to just be able to run with retains_graph=False. <a href="https://github.com/pytorch/pytorch/issues/138118" class="inline-onebox" rel="noopener nofollow ugc">[autograd] do not clear all variables when retain_graph=False · Issue #138118 · pytorch/pytorch · GitHub</a>.<br>
Otherwise as you mention, with a workaround like retain_graph=True the full graph would be kept  alive even though you only need part of it.</p>
<p>There’s no way to partially clear a graph in a fine-grained way today, but you can use saved_tensor_hooks <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> to manage the saved tensors yourself. When you no longer need the particular saved tensor, you can manually clear it. It  requires you to manually define for which nodes you’d like to clear. <code>torch._C._current_autograd_node()</code> which isn’t very easy.</p>

Neirda: <p>Thanks for your answer. The solution I’m using is to replace layers by a custom autograd node that computes only dy/dx and saves <code>grad_output</code> and <code>x</code> somewhere, then having a separate function to compute dy/dw. The downside is that I need to re-implement dy/dw for every layer, but in my case I only care for Linears so that’s fine</p>

soulitzer: <p>I’m curious what you are using this for by the way. (we’ve been thinking about some of these problems for zero-bubble pipelining, but the we’re building things that are specific to that use case)</p>

Neirda: <p>For exactly the same problem, zero-bubble pipeline haha<br>
I’m doing some experiments on distributed training and I needed to have an implementation of this. I looked at torch pipelining’s way of handling it but it seems a bit odd<br>
Since you keep some nodes for the <code>stage_backward_weight</code> call,  aren’t all nodes kept between the 2 calls due to the graph structure? There are still some references to the other nodes, which have references to their saved tensors. And because the backward is called with <code>retain_graph=True</code>, they’re not deleted by autograd. Is there something I’m missing?</p>
<pre data-code-wrap="py"><code class="lang-py">x = torch.randn(batch_size, seq_len, dim, device = device)
target = torch.randn(batch_size, seq_len, dim, device = device)
with Measure("forward"):
    loss = F.mse_loss(model(x), target)

with Measure("backward input"):
    dinputs, param_groups = stage_backward_input([loss], None, [x], model.parameters())

with Measure("backward parameters"):
    wgrads = stage_backward_weight(model.parameters(), param_groups)
</code></pre>
<p>Running that small example on a Transformer model gives this, where memory does not seem freed after <code>stage_backward_inputs</code><br>
(Measure prints the memory difference after the computation)<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/a/e/ae91ffaa71fa749f88c7bdae6f2a3e3953562898.png" data-download-href="https://discuss.pytorch.org/uploads/default/ae91ffaa71fa749f88c7bdae6f2a3e3953562898" title="mem_torch_zb"><img src="https://discuss.pytorch.org/uploads/default/original/3X/a/e/ae91ffaa71fa749f88c7bdae6f2a3e3953562898.png" alt="mem_torch_zb" data-base62-sha1="oUjSB63JZ3RNrWFGDoPjQHKK8Na" width="523" height="178"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">mem_torch_zb</span><span class="informations">523×178 6.58 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>On the same model, using my implementation, I was able to get this:<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/0/60b6c4839c203525266273911ec66ae462efde29.png" data-download-href="https://discuss.pytorch.org/uploads/default/60b6c4839c203525266273911ec66ae462efde29" title="mem_elf_zb"><img src="https://discuss.pytorch.org/uploads/default/original/3X/6/0/60b6c4839c203525266273911ec66ae462efde29.png" alt="mem_elf_zb" data-base62-sha1="dNzpJiOFUweZUvCM7v2NBlYMxGp" width="469" height="188"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">mem_elf_zb</span><span class="informations">469×188 4.78 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>So, maybe I got something wrong on how those functions work? They showed <a href="https://discuss.pytorch.org/t/distributed-w-torchtitan-training-with-zero-bubble-pipeline-parallelism/214420">here</a> that they measured the memory reduction during training, but I’m wondering what was the model used</p>

soulitzer: <blockquote>
<p>Is there something I’m missing?</p>
</blockquote>
<p>You’re not missing too much I think haha - Yes, I’d expect the peak memory to be higher than what is theoretically possible due to retains_graph=True on the entire graph.</p>
<p>But in my head, part of why this isn’t so bad is because (1 )the first stage was previously (1F1B) the bottleneck on peak activation memory and this change does not make the peak memory on the first stage worse and  (2) its most impactful on the last stage but all this does is just bring the peak memory of the last stage close to the first stage.</p>

Neirda: <p>Yes that’s true, it’s not that important for zero-bubble schedules since the first stage is the bottleneck. It only matters when the peak on the last device is higher somehow<br>
Thanks again for the information. If you ever implement a generic way to do that in autograd I’d be interested in using it!</p>