How to optimize memory allocation with intermediate tensor operations

mslyon: <p>I’m creating a custom <code>nn.Module</code> and I’ve been running to GPU OOM problems because at times the tensors within the module can get quite large. From a memory management perspective I know that using tensor methods like <code>Tensor.view</code>, <code>Tensor.expand</code> etc are cheap as they don’t reassign the underlying data structure. To better optimise my code I have been using <code>pytorch_memlab</code> to inspect the memory footprint.</p>
<p>I have noticed that the following operations do not assign the intermediate tensor <code>c</code> in memory, i.e. the pointwise multiplication result is not stored in memory, and is instead immediately passed to the <code>torch.sum</code> function.</p>
<pre><code class="lang-python">def some_module_func(self, a, b):
    # First expand a &amp; b using cheap operation expand
    a = a.expand(-1, 50, 50, 50, -1, -1)
    b = b.expand(-1, 50, 50, 50, -1, -1)
    # Then do elementwise multiplication, if this tensor was saved in memory it would be large
    c = torch.mul(a, b)
    # Then sum over the expanded dimensions
    d = torch.sum(c, dim=(1, 2, 3))
    return d
</code></pre>
<p>However, an operation like <code>torch.gather</code> seemingly will assign a new tensor in memory as output. Here the expanded tensor <code>a</code> will be created in memory after the <code>a.gather</code> operation.</p>
<pre><code class="lang-python">def some_other_module_func(self, a, b, indexer):
    a = a.expand(-1, 50, 50, 50, -1, -1)
    # Here we rearrange dimension 1 to some predefined order
    a = a.gather(dim=1, index=indexer) # this large tensor is written to memory

    b = b.expand(-1, 50, 50, 50, -1, -1)
    # Then do elementwise multiplication
    c = torch.mul(a, b)
    # Then sum over the expanded dimensions
    d = torch.sum(c, dim=(1, 2, 3))
    return d
</code></pre>
<p>My question is: in order to better optimise my code, how can I tell which operations/functions are likely to defer memory assignment? Is there any way that I can apply <code>Tensor.gather</code> for example without the output being explicitly assigned in memory, like with <code>torch.mul</code>?</p>

ptrblck: <aside class="quote no-group" data-username="mslyon" data-post="1" data-topic="149053">
<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/mslyon/48/56032_2.png" class="avatar"> mslyon:</div>
<blockquote>
<p>I have noticed that the following operations do not assign the intermediate tensor <code>c</code> in memory, i.e. the pointwise multiplication result is not stored in memory, and is instead immediately passed to the <code>torch.sum</code> function.</p>
</blockquote>
</aside>
<p>I don’t think that’s true as I’m seeing the expected increase in memory:</p>
<pre data-code-wrap="python"><code class="lang-python">a = torch.randn(1, 1, 1, 1, 1, 1, device='cuda')
b = torch.randn(1, 1, 1, 1, 1, 1, device='cuda')

print(torch.cuda.memory_allocated())
# 2048

a = a.expand(-1, 50, 50, 50, -1, -1)
b = b.expand(-1, 50, 50, 50, -1, -1)
print(torch.cuda.memory_allocated())
# 2048

c = torch.mul(a, b)
print(torch.cuda.memory_allocated())
# 502272

d = torch.sum(c, dim=(1, 2, 3))
print(torch.cuda.memory_allocated())
# 502784
</code></pre>

mslyon: <blockquote>
<p>I don’t think that’s true as I’m seeing the expected increase in memory</p>
</blockquote>
<p>Ah, so I did not realise that memory was deallocated when leaving the namespace scope. Thus the memory checks I was performing previously were obfuscating the true memory requirements as I was only checking memory allocation outside the module method.</p>
<pre data-code-wrap="python"><code class="lang-python">class Module(torch.nn.Module):
    
    def forward(self, a, b):
        a = a.expand(-1, 50, 50, 50, -1, -1)
        b = b.expand(-1, 50, 50, 50, -1, -1)
        print(torch.cuda.memory_allocated())
        c = torch.mul(a, b)
        print(torch.cuda.memory_allocated())
        d = torch.sum(c, dim=(1, 2, 3))
        print(torch.cuda.memory_allocated())

        return d

module = Module()

d = module(a, b)
# 1024
# 501248
# 501760

print(torch.cuda.memory_allocated())
# 1536
</code></pre>
<p>So I guess my main question is: is there any way of not writing the intermediate elementwise multiplication to memory? I trued using <code>torch.utils.checkpoint.checkpoint</code> but still ran into OOM problems.</p>

Abhishek_Tyagi: <p><a class="mention" href="/u/mslyon">@mslyon</a> I was wondering if you were able to figure out how to optimize the memory management for your case?</p>

mslyon: <p><a class="mention" href="/u/abhishek_tyagi">@Abhishek_Tyagi</a> No unfortunately I did not.</p>