How does Pytorch handle in-place operations without losing information necessary for backpropagation?

lesserfish: <p>I’ve seen several questions about in-place operations regarding their efficiency, but I’m actually more confused about the inner-workings of pytorch.</p>
<p>Let’s take a simple example like</p>
<pre><code class="lang-auto">a = torch.randn(10)
b = torch.randn(10)
c = torch.randn(10)

x1 = a * b
x2 = x1 * c 
</code></pre>
<p>In this case, things are easy. Backpropagations happens like this:</p>
<pre><code class="lang-auto">x2.grad &lt;- 1
c.grad &lt;- x2.grad * x1 = x1 = a * b
x1.grad &lt;- x2.grad * c = c
b.grad &lt;- x1.grad * a = c * a
a.grad &lt;- x1.grad * b = c * b
</code></pre>
<p>Everything works correctly. However, in this scenario we have are allocating two buffers: <code>x1</code> and <code>x2</code>. Now, what happen when we do something like this:</p>
<pre><code class="lang-auto">x = a * b
x = x * c 
</code></pre>
<p>It seems to me that the overall expression is the same. However, if we try to compute the gradients the same way we did before, we will run into the following problem:</p>
<pre><code class="lang-auto">x.grad &lt;- 1
c.grad &lt;- x.grad * x = x = a * b * c
</code></pre>
<p>Uh oh. we already got a mistake. Since we performed the multiplication with c in place, we lost the buffer containing <code>a * b</code> which was needed in order to calculate the gradient of <code>c</code>. Does Pytorch actually have buffers for every intermediate calculation?</p>
<p>I can imagine of some possible solutions to this</p>
<ol>
<li>Every intermediate node actually has its own grad. I think this is likely not the answer as it would require a shit ton of memory</li>
<li>Every intermediate step actually uses temporary buffers? If so, how is allocation / deallocation of temporary buffers handled?</li>
</ol>
<p>How is this kind of problem solved in modern frameworks?</p>

KFrank: <p>Hi Lesser!</p>
<p>The short story is that you are not performing an in-place operation.</p>
<aside class="quote no-group" data-username="lesserfish" data-post="1" data-topic="216732">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/lesserfish/48/75678_2.png" class="avatar"> lesserfish:</div>
<blockquote>
<p>Now, what happen when we do something like this:</p>
<pre><code class="lang-auto">x = a * b
x = x * c 
</code></pre>
</blockquote>
</aside>
<p><code>x = a * b</code> creates a new pytorch tensor equal to <code>a * b</code> and creates a new (or maybe<br>
reuses an existing) python reference <code>x</code> and sets it to refer to the <code>a * b</code> tensor.  Then<br>
<code>x = x * c</code> creates a <em>new</em> pytorch tensor equal to <code>x * c</code> and sets the python reference<br>
<code>x</code> to refer to that new tensor.  Although the python reference <code>x</code> now refers to what I<br>
will call the <code>x * c</code> tensor, the <code>a * b</code> tensor still exists, although it is no longer referred<br>
to by the python reference <code>x</code>.</p>
<p>If this is part of a forward pass that autograd is tracking, autograd keeps its own<br>
reference to the <code>a * b</code> tensor (if needed) for use in the backward pass.</p>
<p>(If you had wanted to modify the <code>a * b</code> tensor in place, you could use, for example,<br>
<code>x.mul_ (c)</code> at the point where <code>x</code> still refers to the <code>a * b</code> tensor.)</p>
<p>Best.</p>
<p>K. Frank</p>

lesserfish: <p>Thank you very much for your explanation, Frank!</p>
<p>So it seems that doing</p>
<pre><code class="lang-auto">x = a * b
x = x * c 
</code></pre>
<p>Is not at all different than doing</p>
<pre><code class="lang-auto">x1 = a * b
x2 = x1 * c
</code></pre>
<p>correct?</p>
<p>I just have another question, if you wouldn’t mind. In the scenario described above, both <code>x1</code> and <code>x2</code> are non-leaf tensors correct?</p>
<p>I was under the impression that non-leaf tensors do not gradient buffers associated to them. So how does backpropagation happen?  For example, previously I thought it was like this, but:</p>
<pre><code class="lang-auto">x2.grad &lt;- 1  // This can't happen as x1.grad does not exist.
c.grad &lt;- x2.grad * x1 = x1 = a * b; // x2.grad also does not exist
x1.grad &lt;- x2.grad * c = c // and so on
b.grad &lt;- x1.grad * a = c * a
a.grad &lt;- x1.grad * b = c * b
</code></pre>
<p>Am I misunderstanding the idea that non-leaf tensors do not have gradients? Do they use temporary buffers for backpropagation?</p>
<p>Thank you very much!</p>

KFrank: <p>Hi Lesser!</p>
<aside class="quote no-group" data-username="lesserfish" data-post="3" data-topic="216732">
<div class="title">
<div class="quote-controls"></div>
<img alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/lesserfish/48/75678_2.png" class="avatar"> lesserfish:</div>
<blockquote>
<p>So it seems that doing</p>
<pre><code class="lang-auto">x = a * b
x = x * c 
</code></pre>
<p>Is not at all different than doing</p>
<pre><code class="lang-auto">x1 = a * b
x2 = x1 * c
</code></pre>
</blockquote>
</aside>
<p>Yes, they are basically the same.  Note, in the first case, the <code>a * b</code> tensor no longer<br>
has a reference referring to it, so (unless something else refers to it, such as a<br>
reference used internally by autograd) it will be freed.  In the second case, in contrast,<br>
<code>x1</code> continues to refer to the <code>a * b</code> tensor, so it won’t be freed (until <code>x1</code> goes out of<br>
scope, or some such).</p>
<aside class="quote no-group">
<blockquote>
<p>both <code>x1</code> and <code>x2</code> are non-leaf tensors correct?</p>
</blockquote>
</aside>
<p>Correct.</p>
<aside class="quote no-group">
<blockquote>
<p>non-leaf tensors do not gradient buffers associated to them. So how does backpropagation happen?</p>
</blockquote>
</aside>
<p>When you backpropagate (with the default processing), autograd computes, for<br>
example, the gradient of some loss with respect to<code>x2</code> and then passes that gradient<br>
up the computation graph (the backpropagation chain) in order to compute the<br>
gradient of the loss with respect to <code>x1</code>.  When it’s done using the gradient with<br>
respect to <code>x2</code>, that gradient gets freed.</p>
<p>Non-leaf tensors can, in fact, have their gradients stored in a <code>.grad</code> property during<br>
backpropagation.  For example, calling <code>x2.retain_grad()</code> before backpropagation<br>
will cause autograd to store the gradient with respect to <code>x2</code> in <code>x2.grad</code>.</p>
<p>(This is as if <code>.retain_grad()</code> is called by default on leaf tensors and not on non-leaf<br>
tensors.)</p>
<p>The forward pass / backward pass is very memory hungry for big models, so you really<br>
do want those intermediate gradients to be freed as soon as they can be – hence this<br>
default behavior of only keeping gradients for leaf tensors.</p>
<p>Best.</p>
<p>K. Frank</p>