CUDA memory leak when interop with Pytorch

[woodknight]

I am utilizing Warp to implement a custom operator for my PyTorch neural network and encountered a CUDA memory leaking issue. I found the cause of that is the input I send to warp.launch function is from a PyTorch tensor with requires_grad=True. Is this considered a bug? I have attached the piece of the code to reproduce the error. If you run it, you will see the CUDA memory usage continue to grow until out of memory.

import warp as wp
import torch

@wp.kernel
def op_kernel(x: wp.array(dtype=wp.float32), y: wp.array(dtype=wp.float32)):
    i = wp.tid()
    y[i] = 2.0 * x[i]

wp.init()

i = 0
while True:
    i += 1
    print(i)

    x = torch.rand((4096), dtype=torch.float32, device='cuda', requires_grad=True)
    wp_x = wp.from_torch(x)
    wp_y = wp.empty_like(wp_x)
    wp.launch(kernel=op_kernel, dim=x.shape, inputs=[wp_x], outputs=[wp_y])

[eric-heiden]

Hi @woodknight,

Thank you for reporting this issue! We will have a fix for this in the next version. For a quick fix, you could ensure that the gradient tensor of x is allocated before calling wp.form_torch, i.e. something like

...
x = torch.rand((4096), dtype=torch.float32, device='cuda', requires_grad=True)
x.grad = torch.zeros_like(x, requires_grad=False)
wp_x = wp.from_torch(x)
...

But this doesn't resolve the memory leak entirely (just slows it down), we will take a further look into this.

[nvlukasz]

We have a change that should fix the gradient-related leaks and will try to get it out shortly.

[woodknight]

I have tested the newest release and I see the problem has been fixed. Thanks!

[wangpinzhi]

import warp as wp
import torch

@wp.kernel
def op_kernel(x: wp.array(dtype=wp.float32), y: wp.array(dtype=wp.float32)):
    i = wp.tid()
    y[i] = 2.0 * x[i]

wp.init()

i = 0
tape = wp.Tape()
while True:
    i += 1
    print(i)
    x = torch.rand((4096), dtype=torch.float32, device='cuda', requires_grad=True)
    wp_x = wp.from_torch(x)
    wp_y = wp.empty_like(wp_x)
    
    with tape:
        wp.launch(kernel=op_kernel, dim=x.shape, inputs=[wp_x], outputs=[wp_y])
    tape.backward()
    tape.zero()

I found that if use the 'Tape', the gpu memory usage still increases. Is it a BUG?

[nvlukasz]

I found that if use the 'Tape', the gpu memory usage still increases. Is it a BUG?

Hi @wangpinzhi, I don't think that's a memory leak. The memory usage keeps going up because the same Tape object is being used in each iteration and it hangs on to all the gradients. Tape.zero() zeros out the gradients but doesn't delete them.

If a new Tape object is used for each iteration, the memory usage does not increase:

i = 0
while True:
    i += 1
    print(i)
    x = torch.rand((4096), dtype=torch.float32, device='cuda', requires_grad=True)
    wp_x = wp.from_torch(x)
    wp_y = wp.empty_like(wp_x)
    
    tape = wp.Tape()    # <--- here

    with tape:
        wp.launch(kernel=op_kernel, dim=x.shape, inputs=[wp_x], outputs=[wp_y])
    tape.backward()
    tape.zero()