Memory usage when calculating an array by chunks

username: yiftach
cooked: <p>I want to verify I understand how things work.</p>
<p>Consider the following two alternatives (some unimportant details are left out, but I could construct a more complete example if needed):</p>
<pre><code class="lang-auto">ys = list()
for x in xs:
    l.append(f(x))
ys = torch.stack(ys)
</code></pre>
<p>vs</p>
<pre><code class="lang-auto">ys = torch.tensor(len(xs), ...)
for i, x in enumerate(xs):
    ys[i] = f(x)
</code></pre>
<p>My question - is it true that while both options are fairly similar, the peak memory usage of the 2nd is higher? I would think that if <code>c</code> is the memory needed for during computation of <code>f(x)</code>, we will need <code>c + (len(xs) + 1) * size(f(x))</code> instead of just <code>c + len(xs) * size(f(x))</code>, because in the second version the whole array exists in memory (hence <code>len(xs) * size(f(x))</code>) and we calculate <code>f(x)</code> (hence <code>c + size(f(x))</code>), and rather than directly using the memory at <code>ys[i]</code>, we calculate using a new memory chunk and finally copy to <code>ys[i]</code>.<br>
Will be happy to confirm this or learn how to verify such things.</p>

username: ptrblck
cooked: <p>The second approach should reduce the peak memory since you are not creating the intermediate tensor in <code>torch.stack</code>.<br>
Here is an example showing the allocated and max. memory usage.<br>
Note that I’ve disabled the cuBLAS workspace allocation to make the reports clearer as it would be reported otherwise.</p>
<pre><code class="lang-python">import torch
import torch.nn as nn

import os

# disallow cuBLAS to allocate a workspace to get a better understanding of the memory allocations
# alternatively you can also clear it via torch._C._cuda_clearCublasWorkspaces()
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":0:0"

###############################################################################
## setup
device = "cuda"

xs = torch.randn(1024, 16, 1024).to(device)
print(torch.cuda.memory_allocated()/1024**2)
# 64.0

f = nn.Linear(1024, 1024, bias=False).to(device)
print(torch.cuda.memory_allocated()/1024**2)
# 68.0

print(torch.cuda.max_memory_allocated()/1024**2)
# 68.0

###############################################################################
## 1st approach
ys = list()
for x in xs:
    ys.append(f(x))

print(torch.cuda.memory_allocated()/1024**2)
# 132.0 = 68.0 + 1024*16*1024*4 / 1024**2 = 68.0 + 64.0
print(torch.cuda.max_memory_allocated()/1024**2)
# 132.0

ys = torch.stack(ys)

print(torch.cuda.memory_allocated()/1024**2)
# 132.0
print(torch.cuda.max_memory_allocated()/1024**2)
# 196.0 = 132.0 + 64.0 (temp)

###############################################################################
## 2nd approach
# make sure we are at the same stage
print(torch.cuda.memory_allocated()/1024**2)
# 68.0

print(torch.cuda.max_memory_allocated()/1024**2)
# 68.0

ys = torch.empty_like(xs)
for i, x in enumerate(xs):
    ys[i] = f(x)

print(torch.cuda.memory_allocated()/1024**2)
# 132.0

print(torch.cuda.max_memory_allocated()/1024**2)
# 132.0625
</code></pre>

username: yiftach
cooked: <p>First of all, thanks <a class="mention" href="/u/ptrblck">@ptrblck</a> for the detailed answer and the elaborate code.<br>
I still have a couple of questions.</p>
<ol>
<li>I played with your example and I think the correct answer can vary. While still working with a toy example, I’m trying to imitate my real usage - <code>f</code> requires a high amount of memory, but only temporarily (<code>f</code> contains attention blocks), which is the <code>c</code> from my example.<br>
I modified your code in a slightly silly way just to describe that situation - I replace the linear block with:<pre><code class="lang-auto">def f(x):
  torch.zeros(1024 ** 3 // 4).to(device)
  return x
</code></pre>
</li>
</ol>
<p>Also, for nicer numbers and faster time, I changed to <code>xs = torch.randn(1024, 16, 1024).to(device)</code>.<br>
Now, my numbers are:<br>
1 MB in first block - we only allocate for <code>x</code><br>
1st approach:1.0, 1025.0, 2.0, 1025.0<br>
2nd approach: 2.0, 1025.0, 2.0, 1026.0<br>
So in a theoretical GPU with 1025 GPU memory, the first approach fits and the second doesn’t - right? Or did I cheat?</p>
<ol start="2">
<li>You allocate on the CPU and then use<code>.to(device)</code> instead of directly allocating on the device - is it important here or just for quick prototyping?</li>
<li>Could you please elaborate more on the effect we cancel when we do <code>os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":0:0"</code>? What might happen if we don’t do that?</li>
</ol>
<p>Much appreciated,<br>
Yiftach</p>

username: ptrblck
cooked: <ol>
<li>
<p>Your code doesn’t correspond to a real use case as you are creating intermediate tensors, which are not used or assigned at all. So it seems as if you are now profiling the caching mechanism, how fast tensors are freed (i.e. returned to the cache etc.). But yes, generally just profiling your real use case might still be a good idea.</p>
</li>
<li>
<p>No, you should directly allocate on the device for performance reasons.</p>
</li>
<li>
<p>The reported memory won’t “perfectly” match the expected one only considering the tensors since cuBLAS will allocate its workspace, too, as seen here:</p>
</li>
</ol>
<pre><code class="lang-python">## setup
device = "cuda"

xs = torch.randn(1024, 1024).to(device)
lin = nn.Linear(1024, 1024, bias=False).to(device)

print(torch.cuda.memory_allocated()/1024**2)
# 8.0

out = lin(xs)
print(torch.cuda.memory_allocated()/1024**2)
# 20.125

torch._C._cuda_clearCublasWorkspaces()
print(torch.cuda.memory_allocated()/1024**2)
# 12.0
</code></pre>
<p>The workspace allocation is irrelevant for your use case, which is why I’ve disabled it to create memory allocation values which are easier to understand.</p>

username: yiftach
cooked: <p>Fair enough, I may have over-simplified. Here is an example of a more realistic <code>f</code>, where the second approach requires more memory (at its peak):</p>
<pre><code class="lang-auto">import torch
import torch.nn as nn
import os

os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":0:0"

###############################################################################
## setup
device = "cuda"

xs = torch.randn(16, 4 * 1024, 256, device=device)
print(torch.cuda.memory_allocated()/1024**2)
# 64.0

f = nn.MultiheadAttention(256, 1, device=device)
print(torch.cuda.memory_allocated()/1024**2)
# 65.00390625

print(torch.cuda.max_memory_allocated()/1024**2)
# 65.00390625

###############################################################################
## 1st approach
ys = list()
for x in xs:
    ys.append(f(x, x, x)[0])

print(torch.cuda.memory_allocated()/1024**2)
# 1409.00390625
print(torch.cuda.max_memory_allocated()/1024**2)
# 1477.00390625

ys = torch.stack(ys)

print(torch.cuda.memory_allocated()/1024**2)
# 1473.00390625
print(torch.cuda.max_memory_allocated()/1024**2)
# 1477.00390625

###############################################################################
## 2nd approach
# make sure we are at the same stage
print(torch.cuda.memory_allocated()/1024**2)
# 65.00390625

print(torch.cuda.max_memory_allocated()/1024**2)
# 65.00390625

ys = torch.empty_like(xs)
for i, x in enumerate(xs):
    ys[i] = f(x, x, x)[0]

print(torch.cuda.memory_allocated()/1024**2)
# 1409.00390625

print(torch.cuda.max_memory_allocated()/1024**2)
# 1481.00390625
</code></pre>
<p>To be honest, even in this case of MHA, only certain choices for the sequence length and the hidden dimension make the 2nd approach worse, and the difference is small in any case. Still, I believe the answer may depend on <code>f</code>. Perhaps a general rule could be constructed about when each is favorable (or one may resort to profiling).</p>