Eager Autograd on Leaves

LeanderK: <p>I am working on a huge model where I am very much memory constrained. I continuously generate terms for regularisation, my loss looks something like:</p>
<p><code>loss = accuracy + sum_ penalty alpha*penalty</code></p>
<p>my computational graph looks like this:</p>
<pre><code class="lang-auto">computation -&gt; computation -&gt; computation -&gt; computation -&gt; computation -&gt; (...)
                    |                          |                      
                    -&gt; penalty                 -&gt; penalty                                                       (...)
</code></pre>
<p>since any further computation is independent of the penalty, i could calculate the gradients for the penalty eagerly. So let’s say my function is called <code>calc_penalty(tensor)</code>, what I would like is to immediately calculate the gradient wrt the input so that every tensor used in the computation of <code>calc_penalty</code> can be freed (even the input-tensor).</p>

KFrank: <p>Hi Leander!</p>
<p>Let me tell you what I <em>think.</em>  (I haven’t tested or experimented with any<br>
of this, but I think it will work.)</p>
<p>First a comment on autograd terminology:  The “leaves” of the computation<br>
graph are the tensors that have <code>requires_grad = True</code>, aren’t derived from<br>
anything else (at least not from anything that has <code>requires_grad = True</code>),<br>
and are the tensors with respect to which you wish to compute the gradient.<br>
The “root” (or possibly roots) of the computation graph is the final loss value<br>
whose gradient you wish to compute and on which you call <code>.backward()</code>.</p>
<p>Using this terminology, I think you might have meant “Roots,” rather than<br>
“Leaves,” in your title.</p>
<aside class="quote no-group" data-username="LeanderK" data-post="1" data-topic="154885" data-full="true">
<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/leanderk/48/8099_2.png" class="avatar"> LeanderK:</div>
<blockquote>
<p>my computational graph looks like this:</p>
<pre><code class="lang-auto">computation -&gt; computation -&gt; computation -&gt; computation -&gt; computation -&gt; (...)
                    |                          |                      
                    -&gt; penalty                 -&gt; penalty                                                       (...)
</code></pre>
<p>since any further computation is independent of the penalty, i could calculate the gradients for the penalty eagerly.</p>
</blockquote>
</aside>
<p>If I understand correctly, I believe that you can do this.  I think that the following<br>
pseudo-code sketch should work (but to repeat, I haven’t tried this):</p>
<pre data-code-wrap="python"><code class="lang-python"># comp1
comp1 = f1 (stuff)

# penalty1, with some explicit intermediate results
p1a = g1a (comp1)
p1b = g1b (p1a)
penalty1 = g1c (p1b)

# accumulate penalty1 grad
penalty1.backward (retain_graph = True)   # preserve graph through f1 (as well as g1c, g1b, and g1a)
penalty1 = None   # release penalty1 (and then release intermediate results)
p1b = None        # set references holding on to penalty1-specific graph to None
p1a = None        # should release section of graph from penalty1 back through g1c, g1b, and g1a

# comp2
comp2 = f2 (comp1)

# penalty2, with some explicit intermediate results
p2a = g2a (comp2)
p2b = g2b (p2a)
penalty2 = g2c (p2b)

# accumulate penalty2 grad
penalty2.backward (retain_graph = True)   # preserve graph back through f2 and f1 (as well as g2c, g2b, and g2a)
penalty2 = None   # release penalty2 (and then release intermediate results)
p2b = None        # set references holding on to penalty2-specific graph to None
p2a = None        # should release section of graph from penalty2 back through g2c, g2b, and g2a

# comp3 and final loss
comp3 = f3 (comp2)
loss = loss_fn (comp3)

# accumulate final loss grad (and release remaining graph)
loss.backward()   # backpropagates through loss_fn, f3, f2, and f1 (and further back through stuff)
</code></pre>
<p>The idea is two-fold:  First, you call <code>.backward (retain_graph = True)</code><br>
“eagerly” on your penalty terms, using <code>retain_graph = True</code> to preserve<br>
the computation graph through the “backbone” of the graph.  Second, as I<br>
understand it, the references in the computation graph that keep tensors in<br>
memory start in the downstream roots (e.g., final loss value and your penalties)<br>
and refer to nodes further upstream towards the leaves (e.g., the parameters<br>
you are training).  Therefore, if you set certain of your own root references to<br>
<code>None</code>, everything in the section of the graph between that deleted root<br>
reference and the still-referred-to backbone can have its memory freed.<br>
Furthermore, that “penalty-section” of the graph will be available to have<br>
its memory freed even though you specified <code>retain_graph = True</code> when<br>
you called <code>penalty.backward()</code>.</p>
<p>Anyway, that’s my theory …</p>
<aside class="quote no-group">
<blockquote>
<p>what I would like is to immediately calculate the gradient wrt the input so that every tensor used in the computation of <code>calc_penalty</code> can be freed (even the input-tensor).</p>
</blockquote>
</aside>
<p>Yes, I believe that this will work, with the proviso that if your “input-tensor”<br>
is needed to perform some part of the backbone backward pass, the<br>
backbone will automatically maintain a reference to it, so it won’t be freed<br>
(until your final <code>loss.backward()</code> is called).</p>
<p>Best.</p>
<p>K. Frank</p>

LeanderK: <p>I am sorry for taking so long, I didn’t come to implemented it until now because other things were more complicated than I thought.</p>
<p>I need a bit more control over the gradients. As I understood it, <code>backward</code> computes the whole tree. How can I only compute it until the input of the penalty (<code>comp1</code> and <code>comp2</code>?)? I need to manipulate the gradient again after computing it, there’s a scaling for each sample I now compute at the end of the network.</p>

KFrank: <p>Hi Leander!</p>
<aside class="quote no-group" data-username="LeanderK" data-post="3" data-topic="154885" data-full="true">
<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/leanderk/48/8099_2.png" class="avatar"> LeanderK:</div>
<blockquote>
<p>How can I only compute it until the input of the penalty (<code>comp1</code> and <code>comp2</code>?)? I need to manipulate the gradient again after computing it,</p>
</blockquote>
</aside>
<p>I don’t understand your use case, but let me make a couple of comments:</p>
<p>If you need to backpropagate the gradient of <code>penalty1</code> back through <em>only</em><br>
the computation of <code>penalty1</code> (and not through the computation of <code>comp1</code>)<br>
then simply, <code>.clone().detach()</code> <code>comp1</code>, set <code>requires_grad = True</code> on<br>
the <code>detach()</code>ed copy, compute <code>penalty1</code>, and call <code>penalty1.backward()</code>.</p>
<p>But I imagine that you need to backpropagate <code>penalty1</code>’s gradient,<br>
perhaps in modified form, back through the computation of <code>comp1</code> as<br>
well, presumably in order to update model parameters that are used<br>
to compute <code>comp1</code>.</p>
<p>In this case, you could <code>.clone().detach()</code>, backpropagate just back to<br>
the end of <code>comp1</code>, and save the resulting gradient, at a cost in memory,<br>
for modification and further processing at some later time.</p>
<p>But it seems to me that you would be better off performing the full<br>
backpropagation of <code>penalty1</code> all at once, modifying its gradient “on<br>
the fly.”</p>
<p>To do this, you can write a <a href="https://pytorch.org/tutorials/beginner/examples_autograd/two_layer_net_custom_function.html" rel="noopener nofollow ugc">custom autograd function</a> whose <code>forward()</code><br>
method just passes <code>comp1</code> through unchanged, but whose <code>backward()</code><br>
method modifies the gradient of <code>penalty1</code> according to your desire before<br>
passing it on back upstream through the computation of <code>comp1</code>.</p>
<p>Best.</p>
<p>K. Frank</p>

LeanderK: <p>I have not implemented it yet, since I am still struggling with other parts, but to me this seems to be the correct idea! It might be hard to understand my use-case, it’s quite special since the tensors my network operates on are in the gigabytes, so very very large. I currently use gradient-checkpointing heavily, but that’s just super slow.</p>

LeanderK: <p>Update: I have implemented it and it works! If somebody wants to do something similar, I would advise to look at how checkpoints are implemented</p>

soulitzer: <blockquote>
<p>As I understood it, <code>backward</code> computes the whole tree. How can I only compute it until the input of the penalty (<code>comp1</code> and <code>comp2</code> ?)?</p>
</blockquote>
<p>FYI you can specify the inputs you want to backward to with <code>.backward(inputs=)</code> and it will only backprop through the parts of your graph that are needed to compute gradients for those specified inputs.</p>