Model.to("cpu") does not release GPU memory allocated by registered buffer

ljj7975: <p>I have figured that registered_buffer does not release GPU memory when the model is moved back to CPU.</p>
<p>Here is the minimal code for reproducing the observation</p>
<pre><code class="lang-auto">import torch
from torch import nn
from subprocess import Popen, PIPE

class TestNet(nn.Module):
    def __init__(
        self
    ):
        super().__init__()
        self.register_buffer("test", torch.Tensor([1,2,3]).view(-1, 1, 1), False)

def log_nvidia_smi():
    print("--- nvidia-smi ---")
    for line in Popen('nvidia-smi', shell=True, stdout=PIPE).stdout:
        line = line.decode("utf-8")
        if "python" in line:
            print(line.rstrip())
    print("------------------")

def log_where_buffer_is(model):
    for buffer_name, buffer in model.named_buffers():
        print(f"buffer {buffer_name} is on {buffer.device}")

model = TestNet()

print("&lt; before moving the model to gpu &gt;")
log_where_buffer_is(model)
log_nvidia_smi()

# move model to gpu
model.to(torch.device("cuda"))

print("\n&lt; after moving the model to gpu &gt;")
log_where_buffer_is(model)
log_nvidia_smi()

# move model to cpu
model.to(torch.device("cpu"))

print("\n&lt; after moving the model back to cpu &gt;")
log_where_buffer_is(model)
log_nvidia_smi()

# delete model and clear cuda cache
del model
torch.cuda.empty_cache()

print("\n&lt; after deleting the model &gt;")
log_nvidia_smi()
</code></pre>
<p>The above code logs the following messages</p>
<pre><code class="lang-auto">&lt; before moving the model to gpu &gt;
buffer test is on cpu
--- nvidia-smi ---
------------------

&lt; after moving the model to gpu &gt;
buffer test is on cuda:0
--- nvidia-smi ---
|    0   N/A  N/A     2527      C   ...rtualenvs/nova/bin/python     1005MiB |
------------------

&lt; after moving the model back to cpu &gt;
buffer test is on cpu
--- nvidia-smi ---
|    0   N/A  N/A     2527      C   ...rtualenvs/nova/bin/python     1005MiB |
------------------

&lt; after deleting the model &gt;
--- nvidia-smi ---
|    0   N/A  N/A       2527      C   ...rtualenvs/nova/bin/python     1003MiB |
------------------
</code></pre>
<p>I have looked into pytorch documentations and figured that there isn’t a way to delete buffers.<br>
I attempted manually deleting the tensors but it didn’t help.</p>
<p>Is there a way to properly clear all the GPU memory after moving the model back to CPU?</p>

eqy: <p>In this case the actual memory used by the model itself (it looks like it’s just a 6 element tensor) is miniscule compared to the memory used by the CUDA context (~300MiB+), and the PyTorch CUDA kernels (~600MiB+). What is happening is that when you invoke <code>.cuda()</code> on something for the first time or initialize a device tensor, this pulls in all of PyTorch’s CUDA kernels into GPU memory and creates a CUDA context. (If you had called a library function in cuDNN or cuBLAS you would expect this usage to go even higher when those kernels are loaded!) Unfortunately, just because there are no more GPU tensors doesn’t mean that this magically goes away.</p>
<p>If you want to see the effect of releasing GPU memory actually held by the model, you might want to increase the amount of memory used by the model (e.g., have it use up 1GiB+) of GPU memory.</p>

ljj7975: <p>Thank you for your explanation.<br>
I have replaced the line</p>
<pre><code class="lang-auto">self.register_buffer("test", torch.Tensor([1,2,3]).view(-1, 1, 1), False)
</code></pre>
<p>to</p>
<pre><code class="lang-auto">self.register_buffer("test", torch.rand([10000, 10000, 10]), False)
</code></pre>
<p>so that I allocate large GPU memory.</p>
<p>The results looks as follows:</p>
<pre><code class="lang-auto">&lt; before moving the model to gpu &gt;
buffer test is on cpu
--- nvidia-smi ---
------------------

&lt; after moving the model to gpu &gt;
--- nvidia-smi ---
|    0   N/A  N/A      3401      C   ...rtualenvs/nova/bin/python     4819MiB |
------------------

&lt; after moving the model back to cpu &gt;
--- nvidia-smi ---
|    0   N/A  N/A      3401      C   ...rtualenvs/nova/bin/python     4819MiB |
------------------

&lt; after deleting the model &gt;
--- nvidia-smi ---
|    0   N/A  N/A      3401      C   ...rtualenvs/nova/bin/python     1003MiB |
------------------
</code></pre>
<p>The memory <code>1003MiB</code> seems unrelated to the model as you mention and it’s not being released properly even though I am allocating greater amount of memory</p>

ljj7975: <p>I have done little more googling based on what I figured from your explanation.</p>
<p>one of the git issue I found says that deleting pytorch cuda context is not going to be supported.</p>
<p>Is this something relevant?</p><aside class="onebox githubissue" data-onebox-src="https://github.com/pytorch/pytorch/issues/20837">
  <header class="source">

      <a href="https://github.com/pytorch/pytorch/issues/20837" target="_blank" rel="noopener nofollow ugc">github.com/pytorch/pytorch</a>
  </header>

  <article class="onebox-body">
    <div class="github-row">
  <div class="github-icon-container" title="Issue">
	  <svg width="60" height="60" class="github-icon" viewbox="0 0 14 16" aria-hidden="true"><path d="M7 2.3c3.14 0 5.7 2.56 5.7 5.7s-2.56 5.7-5.7 5.7A5.71 5.71 0 0 1 1.3 8c0-3.14 2.56-5.7 5.7-5.7zM7 1C3.14 1 0 4.14 0 8s3.14 7 7 7 7-3.14 7-7-3.14-7-7-7zm1 3H6v5h2V4zm0 6H6v2h2v-2z"></path></svg>
  </div>

  <div class="github-info-container">
    <h4>
      <a href="https://github.com/pytorch/pytorch/issues/20837" target="_blank" rel="noopener nofollow ugc">[feature request] a function to empty all the cuda buffer</a>
    </h4>

    <div class="github-info">
      <div class="date">
        opened <span class="discourse-local-date" data-format="ll" data-date="2019-05-23" data-time="00:59:17" data-timezone="UTC">12:59AM - 23 May 19 UTC</span>
      </div>

        <div class="date">
          closed <span class="discourse-local-date" data-format="ll" data-date="2019-06-24" data-time="17:21:21" data-timezone="UTC">05:21PM - 24 Jun 19 UTC</span>
        </div>

      <div class="user">
        <a href="https://github.com/Stonesjtu" target="_blank" rel="noopener nofollow ugc">
          <img alt="Stonesjtu" src="https://avatars.githubusercontent.com/u/4556044?v=4" class="onebox-avatar-inline" width="20" height="20">
          Stonesjtu
        </a>
      </div>
    </div>

    <div class="labels">
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          module: cuda
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          module: memory usage
        </span>
        <span style="display:inline-block;margin-top:2px;background-color: #B8B8B8;padding: 2px;border-radius: 4px;color: #fff;margin-left: 3px;">
          triage review
        </span>
    </div>
  </div>
</div>

  <div class="github-row">
    <p class="github-body-container">## 🚀 Feature


At present the pytorch process uses about 900MB after `lazy_in<span class="show-more-container"><a href="" rel="noopener" class="show-more">…</a></span><span class="excerpt hidden">it` cuda context. I was searching on internet whether the CUDA context can be shutdown, but seems like such API is not provided yet. Here I'm asking if we can do a `pytorch` side context clear, so the minimal CUDA memory should be allocated to pytorch runtime.

## Motivation



I'm developing an interesting function that each pytorch worker interacts with a scheduling server, dynamically moving the workload from/to GPU, so that the CUDA memory can be used for tasks with higher priority. 

## Pitch



```python
tensor = torch.Tensor(100,100).cuda()  # lazy_init happens here
del tensor
torch.cuda.de_init()  # clear all the CUDA memory maintained by pytorch runtime
```

## Alternatives



## Additional context</span></p>
  </div>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>


eqy: <p>Yes, I don’t think there are plans to do something like this in pure PyTorch.</p>

ljj7975: <p>“you cannot delete the CUDA context while the PyTorch process is still running”</p>
<aside class="quote" data-post="10" data-topic="84762">
  <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/ptrblck/48/1823_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/clearing-the-gpu-is-a-headache/84762/10">Clearing the GPU is a headache</a> <a class="badge-category__wrapper " href="/c/vision/5"><span data-category-id="5" data-drop-close="true" class="badge-category " title="Topics related to either pytorch/vision or vision research related topics"><span class="badge-category__name">vision</span></span></a>
  </div>
  <blockquote>
    No, you cannot delete the CUDA context while the PyTorch process is still running and would have to shutdown the current process and use a new one for the downstream application.
  </blockquote>
</aside>


fangyunfeng: <p>I highly recommend adding this method.That’s really necessary when deploying with PyTorch.<br>
Each time initialization a instance from the hard disk load model, it is time consuming. if you can del pytorch cuda context  without exiting the process.<br>
Thank you</p>

ptrblck: <p>CUDA 11.7+ ships with lazy module loading, which will only load the needed device kernels and will thus significantly reduce the startup time as well as CUDA context size. You could try it out and see if this works for you.</p>