# Free some saved tensors after partial backward **URL:** \ **Category:** autograd\ **Created:** [January 16, 2025, 5:45pm UTC](https://discuss.pytorch.org/t/free-some-saved-tensors-after-partial-backward/215486 "2025-01-16T17:45:37Z")\ **Posts on this page:** 7\ **Page:** 1
**Author:** ![Neirda](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/neirda/32/74356_2.png) [@Neirda](https://discuss.pytorch.org/u/Neirda)\ **Post date:** [January 16, 2025, 5:45pm UTC](https://discuss.pytorch.org/t/free-some-saved-tensors-after-partial-backward/215486/1 "2025-01-16T17:45:37Z")
Hello, 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: ```py loss.backward(inputs=[inputs], retain_graph=True) ... loss.backward(inputs=list(model.parameters()), retain_graph=False) ``` However, as I understand, this solution has 2 issues. - 1 - It recomputes all intermediates gradients for the second backward pass. For instance, in this graph, [Sum, Pow, Sqrt, Cos] nodes will be recomputed. ![image](https://discuss.pytorch.org/uploads/default/original/3X/b/b/bbf6e39f3ea595f48225fb02741a0d7b481ac843.png) 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 `retain_graph=True` otherwise some necessary saved\_tensors are freed. This leads to the second problem, - 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 `result` saved for ReluBackward could be deleted. ![annotated-bgraph](https://discuss.pytorch.org/uploads/default/original/3X/d/3/d3c413ce51561eadd1462858656442f81417db1e.jpeg) 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. 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? ---
**Author:** ![soulitzer](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/839c29/32.png) [@soulitzer](https://discuss.pytorch.org/u/soulitzer)\ **Post date:** [February 11, 2025, 11:32pm UTC](https://discuss.pytorch.org/t/free-some-saved-tensors-after-partial-backward/215486/2 "2025-02-11T23:32:59Z")
Instead of using hooks, a slightly different approach is: 1. You can use the inputs= argument to .grad or .backward to get autograd to return the intermediate grad\_outputs. 2. You can use GradientEdge [[autograd] Support GradientEdge as output for torch.autograd.grad by soulitzer · Pull Request #127766 · pytorch/pytorch · GitHub](https://github.com/pytorch/pytorch/pull/127766) to backward starting from an intermediate node without needing to keep the corresponding tensor alive. > The entire graph and all corresponding saved tensors are kept in memory between the first and second backward pass. 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. [[autograd] do not clear all variables when retain\_graph=False · Issue #138118 · pytorch/pytorch · GitHub](https://github.com/pytorch/pytorch/issues/138118). 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. There’s no way to partially clear a graph in a fine-grained way today, but you can use saved\_tensor\_hooks [Hooks for autograd saved tensors — PyTorch Tutorials 2.6.0+cu124 documentation](https://pytorch.org/tutorials/intermediate/autograd_saved_tensors_hooks_tutorial.html) 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. `torch._C._current_autograd_node()` which isn’t very easy. ---
**Author:** ![Neirda](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/neirda/32/74356_2.png) [@Neirda](https://discuss.pytorch.org/u/Neirda)\ **Post date:** [February 12, 2025, 9:29am UTC](https://discuss.pytorch.org/t/free-some-saved-tensors-after-partial-backward/215486/3 "2025-02-12T09:29:19Z")
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 `grad_output` and `x` 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 ---
**Author:** ![soulitzer](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/839c29/32.png) [@soulitzer](https://discuss.pytorch.org/u/soulitzer)\ **Post date:** [February 12, 2025, 10:23pm UTC](https://discuss.pytorch.org/t/free-some-saved-tensors-after-partial-backward/215486/4 "2025-02-12T22:23:27Z")
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) ---
**Author:** ![Neirda](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/neirda/32/74356_2.png) [@Neirda](https://discuss.pytorch.org/u/Neirda)\ **Post date:** [February 13, 2025, 2:21pm UTC](https://discuss.pytorch.org/t/free-some-saved-tensors-after-partial-backward/215486/5 "2025-02-13T14:21:14Z")
For exactly the same problem, zero-bubble pipeline haha 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 Since you keep some nodes for the `stage_backward_weight` 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 `retain_graph=True`, they’re not deleted by autograd. Is there something I’m missing? ```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) ``` Running that small example on a Transformer model gives this, where memory does not seem freed after `stage_backward_inputs` (Measure prints the memory difference after the computation) ![mem_torch_zb](https://discuss.pytorch.org/uploads/default/original/3X/a/e/ae91ffaa71fa749f88c7bdae6f2a3e3953562898.png) On the same model, using my implementation, I was able to get this: ![mem_elf_zb](https://discuss.pytorch.org/uploads/default/original/3X/6/0/60b6c4839c203525266273911ec66ae462efde29.png) So, maybe I got something wrong on how those functions work? They showed [here](https://discuss.pytorch.org/t/distributed-w-torchtitan-training-with-zero-bubble-pipeline-parallelism/214420) that they measured the memory reduction during training, but I’m wondering what was the model used ---
**Author:** ![soulitzer](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/s/839c29/32.png) [@soulitzer](https://discuss.pytorch.org/u/soulitzer)\ **Post date:** [February 13, 2025, 11:41pm UTC](https://discuss.pytorch.org/t/free-some-saved-tensors-after-partial-backward/215486/6 "2025-02-13T23:41:31Z")
> Is there something I’m missing? 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. 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. ---
**Author:** ![Neirda](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/neirda/32/74356_2.png) [@Neirda](https://discuss.pytorch.org/u/Neirda)\ **Post date:** [February 14, 2025, 4:47pm UTC](https://discuss.pytorch.org/t/free-some-saved-tensors-after-partial-backward/215486/7 "2025-02-14T16:47:34Z")
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 Thanks again for the information. If you ever implement a generic way to do that in autograd I’d be interested in using it!