Tracemalloc giving wrong numbers for memory?

username: Imahn
cooked: <p>Let’s take a look at this code:</p>
<pre><code class="lang-python">import tracemalloc
import numpy as np

tracemalloc.start()
a = np.ones((3, 1024, 1024, 1024), dtype=np.uint8)
current, peak = tracemalloc.get_traced_memory()
print(f'Current: {current / 1024**2:.2f} [GB], peak: {peak / 1024**2:.2f} [GB] / 1024**2')
</code></pre>
<p>The output I get is</p>
<pre><code class="lang-bash">Current: 3072.00 [MB], peak: 3072.00 [MB]
</code></pre>
<p>Obviously, we expect an array with a size of about <code>3</code> GB, so this looks good. However, if we do the same thing in PyTorch, i.e.</p>
<pre><code class="lang-python">import torch
import tracemalloc

tracemalloc.start()
x = torch.ones((3, 1024, 1024, 1024), dtype=torch.uint8)
current, peak = tracemalloc.get_traced_memory()
print(f'Current: {current / 1024**2:.2f} [MB], peak: {peak / 1024**2:.2f} [MB]')
</code></pre>
<p>I get the following output:</p>
<pre><code class="lang-bash">0.00 [MB], 0.00 [MB]
</code></pre>
<p>Something seems to be off. I tested PyTorch version <code>1.9.0</code> and Python version <code>3.8.10</code>.</p>
<hr>
<p>My question is: Is this expected behavior? I myself am surprised.</p>

username: tiramisuNcustard
cooked: <p>I don’t know anything about how tensor.ones is implemented (source code). My guess would be that it is similar to lazy evaluation/generator. PyTorch will not allocate memory and fill it with 1’s until the tensor is actually needed for operation on the hardware. This way, PyTorch doesn’t waste precious resources (there tends to be lots of data to be processed in ML/data science) with data that is not needed immediately.</p>

username: Imahn
cooked: <p><a class="mention" href="/u/tiramisuncustard">@tiramisuNcustard</a> I tried the same thing with <code>torch.rand()</code> instead of <code>torch.ones()</code> - though I had to use <code>dtype = torch.float32</code> in that case - and basically observed the same behavior. Do you think PyTorch will still not fill the tensor up with random numbers before we do some operations on it?</p>
<p>Also, any way we can explicitly test your hypothesis? For example by doing operations on the tensor that force PyTorch to write the tensor onto the hardware, but that are not memory-intensive?</p>

username: tiramisuNcustard
cooked: <p><a class="mention" href="/u/Imahn">@Imahn</a>, I would create two 1 GB random tensors. Check the memory allocation. Then multiply the two tensors and store it in a variable. Then check the memory allocation again. If there is a difference in the two memory allocation checks then I would be satisfied with my answer. The reason I am suggesting 1GB is because, again, I don’t know how PyTorch implements these functions (will they allocate memory if the size is not too large - let’s say for example &lt; 200 MB?).</p>

username: Imahn
cooked: <p><a class="mention" href="/u/tiramisuncustard">@tiramisuNcustard</a> Sorry for late reply. I tested the following code (<code>x</code> and <code>y</code> should roughly occupy <code>1</code> GB of memory):</p>
<pre><code class="lang-python">import torch 
import tracemalloc

tracemalloc.start()
x = torch.rand(256, 1024, 1024)
y = torch.rand(256, 1024, 1024)
current, peak = tracemalloc.get_traced_memory()
out = x * y
current_mult, peak_mult = tracemalloc.get_traced_memory()
print(f'Current/peak before multiplication [GB]: {current / 1024**3:.2f}, '
      f'peak before multiplication: {peak / 1024**3:.2f} [GB]'
      f'\nCurrent/peak after multiplication [GB]: {current_mult}, {peak_mult}'
      f'\nTensor dtype: {x.dtype}, output shape: {out.shape}')
</code></pre>
<p>The output is:</p>
<pre><code class="lang-auto">Current/peak before multiplication [GB]: 0.00, 0.00 
Current/peak after multiplication [GB]: 192.00, 192.00
Tensor dtype: torch.float32, output shape: torch.Size([256, 1024, 1024])
</code></pre>
<p>Tbh, I’m now totally puzzled about why the memory consumption is so huge after the multiplication (<code>192</code> GB)? I would have expected <code>2</code> GB for <code>x</code> and <code>y</code>, and maybe a few GB for the multiplication, but <code>O(100)</code> GB for a multiplication sounds like a lot too me?</p>

username: tiramisuNcustard
cooked: <p><a class="mention" href="/u/Imahn">@Imahn</a>, Your first print statement is the following <code>Current/peak before multiplication [GB]: {current / 1024**3:.2f}</code>, if we copy it over to the print statement after the multiplication i.e. <code>Current/peak after multiplication [GB]: {current_mult / 1024**3:.2f}</code> what do we get?</p>
<p>Maybe, the missing <code>/ 1024**3:.2f</code> is causing the trouble?</p>

username: Imahn
cooked: <p><a class="mention" href="/u/tiramisuncustard">@tiramisuNcustard</a> Sorry, totally right. Corrected output:</p>
<pre><code class="lang-python">Current/peak before multiplication [GB]: 0.00, 0.00 
Current/peak after multiplication [GB]: 0.00, 0.00
</code></pre>

username: tiramisuNcustard
cooked: <p>What does <code>out.element_size() * out.nelement()</code> give you? Can you use it before <code>current_mult, peak_mult = tracemalloc.get_traced_memory()</code> and see what it shows?</p>

username: Imahn
cooked: <p>Sure, before</p>
<pre><code class="lang-python">current_mult, peak_mult = tracemalloc.get_traced_memory()
</code></pre>
<p>I added</p>
<pre><code class="lang-auto">print(f'Expected memory consump [GB]: {out.nelement() * out.element_size() / 1024**3:.2f}')
</code></pre>
<p>(as you suggested), the output is</p>
<pre><code class="lang-auto">Expected memory consump [GB]: 1.00
</code></pre>
<p>I did the same for <code>x</code> and <code>y</code> before</p>
<pre><code class="lang-auto">current, peak = tracemalloc.get_traced_memory()
</code></pre>
<p>and I get for both <code>x</code> and <code>y</code> 1 GB each as the expected memory consumption…</p>