How to make sure PyTorch has deallocated GPU memory?

Question

Say we have a function like this:

        def trn_l(totall_lc, totall_lw, totall_li, totall_lr):
            self.model_large.cuda()
            self.model_large.train()
            self.optimizer_large.zero_grad()

            for fb in range(self.fake_batch):
                val_x, val_y = next(self.valid_loader)
                val_x, val_y = val_x.cuda(), val_y.cuda()

                logits_main, emsemble_logits_main = self.model_large(val_x)
                cel = self.criterion(logits_main, val_y)
                loss_weight = cel / (self.fake_batch)
                loss_weight.backward(retain_graph=False)
                cel = cel.cpu().detach()
                emsemble_logits_main = emsemble_logits_main.cpu().detach()

                totall_lw += float(loss_weight.item())
                val_x = val_x.cpu().detach() 
                val_y = val_y.cpu().detach()

            loss_weight = loss_weight.cpu().detach()
            self._clip_grad_norm(self.model_large)
            self.optimizer_large.step()
            self.model_large.train(mode=False)
            self.model_large = self.model_large.cpu()
            return totall_lc, totall_lw, totall_li, totall_lr

On the first call, it allocates 8GB of GPU memory. On the next call, no new memory gets allocated, yet 8GBs are still occupied. I want to have after it is called and the produced first result to have 0 allocated GPU memory or as low as possible.

What I have tried: do retain_graph=False and .cpu().detach() everywhere - no positive effects.

Memory snapshot before

|===========================================================================|
|                  PyTorch CUDA memory summary, device ID 0                 |
|---------------------------------------------------------------------------|
|            CUDA OOMs: 0            |        cudaMalloc retries: 0         |
|===========================================================================|
|        Metric         | Cur Usage  | Peak Usage | Tot Alloc  | Tot Freed  |
|---------------------------------------------------------------------------|
| Allocated memory      |   33100 KB |   33219 KB |   40555 KB |    7455 KB |
|       from large pool |    3072 KB |    3072 KB |    3072 KB |       0 KB |
|       from small pool |   30028 KB |   30147 KB |   37483 KB |    7455 KB |
|---------------------------------------------------------------------------|
| Active memory         |   33100 KB |   33219 KB |   40555 KB |    7455 KB |
|       from large pool |    3072 KB |    3072 KB |    3072 KB |       0 KB |
|       from small pool |   30028 KB |   30147 KB |   37483 KB |    7455 KB |
|---------------------------------------------------------------------------|
| GPU reserved memory   |   51200 KB |   51200 KB |   51200 KB |       0 B  |
|       from large pool |   20480 KB |   20480 KB |   20480 KB |       0 B  |
|       from small pool |   30720 KB |   30720 KB |   30720 KB |       0 B  |
|---------------------------------------------------------------------------|
| Non-releasable memory |   18100 KB |   20926 KB |   56892 KB |   38792 KB |
|       from large pool |   17408 KB |   18944 KB |   18944 KB |    1536 KB |
|       from small pool |     692 KB |    2047 KB |   37948 KB |   37256 KB |
|---------------------------------------------------------------------------|
| Allocations           |   12281    |   12414    |   12912    |     631    |
|       from large pool |       2    |       2    |       2    |       0    |
|       from small pool |   12279    |   12412    |   12910    |     631    |
|---------------------------------------------------------------------------|
| Active allocs         |   12281    |   12414    |   12912    |     631    |
|       from large pool |       2    |       2    |       2    |       0    |
|       from small pool |   12279    |   12412    |   12910    |     631    |
|---------------------------------------------------------------------------|
| GPU reserved segments |      16    |      16    |      16    |       0    |
|       from large pool |       1    |       1    |       1    |       0    |
|       from small pool |      15    |      15    |      15    |       0    |
|---------------------------------------------------------------------------|
| Non-releasable allocs |       3    |      30    |     262    |     259    |
|       from large pool |       1    |       1    |       1    |       0    |
|       from small pool |       2    |      29    |     261    |     259    |
|===========================================================================|

And after calliing function and

torch.cuda.empty_cache()
torch.cuda.synchronize()

We get:

|===========================================================================|
|                  PyTorch CUDA memory summary, device ID 0                 |
|---------------------------------------------------------------------------|
|            CUDA OOMs: 0            |        cudaMalloc retries: 0         |
|===========================================================================|
|        Metric         | Cur Usage  | Peak Usage | Tot Alloc  | Tot Freed  |
|---------------------------------------------------------------------------|
| Allocated memory      |   10957 KB |    8626 MB |  272815 MB |  272804 MB |
|       from large pool |       0 KB |    8596 MB |  272477 MB |  272477 MB |
|       from small pool |   10957 KB |      33 MB |     337 MB |     327 MB |
|---------------------------------------------------------------------------|
| Active memory         |   10957 KB |    8626 MB |  272815 MB |  272804 MB |
|       from large pool |       0 KB |    8596 MB |  272477 MB |  272477 MB |
|       from small pool |   10957 KB |      33 MB |     337 MB |     327 MB |
|---------------------------------------------------------------------------|
| GPU reserved memory   |    8818 MB |    9906 MB |   19618 MB |   10800 MB |
|       from large pool |    8784 MB |    9874 MB |   19584 MB |   10800 MB |
|       from small pool |      34 MB |      34 MB |      34 MB |       0 MB |
|---------------------------------------------------------------------------|
| Non-releasable memory |    5427 KB |    3850 MB |  207855 MB |  207850 MB |
|       from large pool |       0 KB |    3850 MB |  207494 MB |  207494 MB |
|       from small pool |    5427 KB |       5 MB |     360 MB |     355 MB |
|---------------------------------------------------------------------------|
| Allocations           |    3853    |   13391    |   34339    |   30486    |
|       from large pool |       0    |     557    |   12392    |   12392    |
|       from small pool |    3853    |   12838    |   21947    |   18094    |
|---------------------------------------------------------------------------|
| Active allocs         |    3853    |   13391    |   34339    |   30486    |
|       from large pool |       0    |     557    |   12392    |   12392    |
|       from small pool |    3853    |   12838    |   21947    |   18094    |
|---------------------------------------------------------------------------|
| GPU reserved segments |     226    |     226    |     410    |     184    |
|       from large pool |     209    |     209    |     393    |     184    |
|       from small pool |      17    |      17    |      17    |       0    |
|---------------------------------------------------------------------------|
| Non-releasable allocs |      46    |     358    |   12284    |   12238    |
|       from large pool |       0    |     212    |    7845    |    7845    |
|       from small pool |      46    |     279    |    4439    |    4393    |
|===========================================================================|

user_7138517 · Answer

I don't think the other answer is correct. Allocation and deallocation definitely happens during runtime, the thing to note is that the CPU code runs asynchronously from the GPU code, so you need to wait for any deallocation to happen if you want to reserve more memory after it. Take a look at this:

import torch 

a = torch.zeros(100,100,100).cuda()

print(torch.cuda.memory_allocated())

del a
torch.cuda.synchronize()
print(torch.cuda.memory_allocated())

Outputs

4000256
0

So you should del the tensors you don't need and call torch.cuda.synchronize() to make sure that the deallocation goes through before your CPU code continues to run.

In your specific case, after your function trn_l returns, any variables that were local to that function, and do not have references elsewhere, will be deallocated along with the corresponding GPU tensors. All you need to do is wait for this to happen by calling torch.cuda.synchronize() after the function call.

user_4614404 · Answer

So, Pytorch does not allocate and deallocate memory from GPU in training time.

From https://pytorch.org/docs/stable/notes/faq.html#my-gpu-memory-isn-t-freed-properly:

PyTorch uses a caching memory allocator to speed up memory allocations. As a result, the values shown in nvidia-smi usually don’t reflect the true memory usage. See Memory management for more details about GPU memory management.

If your GPU memory isn’t freed even after Python quits, it is very likely that some Python subprocesses are still alive. You may find them via ps -elf | grep python and manually kill them with kill -9 [pid].

You can call torch.cuda.empty_cache() to free all unused memory (however, that is not really good practice as memory re-allocation is time consuming). Docs of empty_cace() : https://pytorch.org/docs/stable/cuda.html#torch.cuda.empty_cache

How to make sure PyTorch has deallocated GPU memory?

2 Answers