Implementing softmax or exp inplace

yiftach: <p>Hi,<br>
I’m using a code based on the Vision Transformer from timm and trying squeeze out the most of my available GPU memory. Often, my attempts crash at the following line:</p>
<p><a href="https://github.com/huggingface/pytorch-image-models/blob/c241081251323dfc5e8dc799d49740c48cc9096f/timm/models/vision_transformer.py#L92" class="inline-onebox" rel="noopener nofollow ugc">pytorch-image-models/timm/models/vision_transformer.py at c241081251323dfc5e8dc799d49740c48cc9096f · huggingface/pytorch-image-models · GitHub</a> which is the <strong>middle</strong> line here:</p>
<pre><code class="lang-auto">            attn = q @ k.transpose(-2, -1)
            attn = attn.softmax(dim=-1)
            attn = self.attn_drop(attn)
</code></pre>
<p>My goal in this post is not to debug that repo, and is a more general pytorch question.</p>
<p>Here’s the problem, as far as I can tell - <code>attn</code> is a very large tensor. The softmax call tries to create a second tensor of the same size and assign it to <code>attn</code>, but there exists a moment in time where both simultaneously reside in GPU memory (requiring double the GPU memory size) until the old <code>attn</code> is replaced with the new one.<br>
Assuming I’m right so far (am I?), I’m trying to see if this could be avoided. Specifically, I believe the problem could be reduced to calculating <code>exp(attn)</code> (and normalizing, later, is less demanding). I know that calling <code>torch.exp(x, out=x)</code> is not possible with autograd (as noted <a href="https://stackoverflow.com/a/53741386" rel="noopener nofollow ugc">here</a>), and as far as I understand the problem is that the input we discard is the output of the previous operator which might need it (as noted <a href="https://discuss.pytorch.org/t/torch-exp-is-modified-by-an-inplace-operation/90216">here</a>, if I understand <a class="mention" href="/u/tom">@tom</a> correctly), but as far as I know, the dot product should not need to keep its output.</p>
<p>Is there anything that can be done here?</p>
<p>Thanks,<br>
Yiftach</p>

tom: <p>This is an interesting idea for saving memory. Keep in mind, though, that the tensor will typically be deallocated after it goes out of scope, so we are indeed talking about the time when the other softmax is computed (my <a href="https://lernapparat.de/pytorch-inplace" rel="noopener nofollow ugc">inplace blogpost</a> has some more detail).</p>
<p>As slow way to do/try the inplace softmax is to compute the row-by-row softmax in a for loop out of place and then copy it into the larger tensor. Then you can wrap this inplace softmax in an autograd function and implement the backward (if you don’t want to derive the formula yourself, you could use the PyTorch code (<a href="https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/SoftMax.cpp#L291" rel="noopener nofollow ugc">https://github.com/pytorch/pytorch/blob/main/aten/src/ATen/native/SoftMax.cpp#L291</a>) or <a href="https://eli.thegreenplace.net/2016/the-softmax-function-and-its-derivative/" rel="noopener nofollow ugc">some reference</a>).</p>
<p>Best regards</p>
<p>Thomas</p>

yiftach: <p>That’s a very interesting workaround. Thanks <a class="mention" href="/u/tom">@tom</a>!</p>