Memory leak

[rztz]

Describe the bug
The CudaJit Backend leaks memory when creating and dropping Tensors.

To Reproduce

fn test_memory_leak() {
    unsafe { libloading::Library::new("libtorch.so").unwrap() };

    // number of loops to overflow GPU memory are dependend on your hardware
    let loops = 10_000;

    // this works and memory does not grow
    test_device_memory_leak::<burn::backend::Candle>(
        &burn::backend::candle::CandleDevice::cuda(0),
        loops,
    );

    // this works and memory does not grow
    test_device_memory_leak::<burn::backend::LibTorch>(
        &burn::backend::libtorch::LibTorchDevice::Cuda(0),
        loops,
    );

    // this overflows my 10 gig memory at aboout loop 2600
    test_device_memory_leak::<burn::backend::CudaJit>(
        &burn::backend::cuda_jit::CudaDevice::new(0),
        loops,
    );
}

fn test_device_memory_leak<B: Backend>(device: &<B as Backend>::Device, loops: usize) {
    const DIM: usize = 4;
    let shape = [32, 4, 84, 84];
    let l = shape.iter().product();
    let data_v = vec![0; l];

    for i in 0..loops {
        if i % 100 == 0 {
            println!("i: {}", i);
        }
        let td = TensorData::new(data_v.clone(), shape);
        let t: Tensor<B, DIM> = Tensor::from_data(td, device);
        drop(t);
    }
}

Expected behavior
drop(t) should free the memory

Screenshots

Desktop (please complete the following information):
Fedora 41, Nvidia cuda 12.6

[nathanielsimard]

Thanks for the bug report. I've look at it and the problem is created by the fusion feature flag. The problem is that there are unused variables in the capture computation stream that never get cleanup, because the cleanup is done on tensors created by computation, but not by data. So if I use the tensor created, there is no longer a memory leak:

fn test_device_memory_leak<B: Backend>(device: &<B as Backend>::Device, loops: usize) {
    const DIM: usize = 4;
    let shape = [32, 4, 84, 84];
    let l = shape.iter().product();
    let data_v = vec![0; l];

    for i in 0..loops {
        if i % 100 == 0 {
            println!("i: {}", i);
        }
        let td = TensorData::new(data_v.clone(), shape);
        let t: Tensor<B, DIM> = Tensor::from_data(td, device);
        let s: Tensor<B, DIM> = t.sum();
        drop(s);
    }
}

The code in question:

burn/crates/burn-fusion/src/ops/float.rs

Line 18 in 140ea75

fn float_from_data(data: TensorData, device: &Device<Self>) -> FloatTensor<Self> {

A good fix would be to register an operation instead of actually doing the copy, so our orphan detector would work, and we could even skip the copy altogether if the tensor was never used.

Also not releated, but if we use a fusable operation (like +), we generate invalid compute shaders, so we should set a limit in our fuser.

[rztz]

My use case is DQN recurrent learning. I wanted to preload data into GPU memory, or send tensors with .todevice from CPU to GPU. (on the other backends, since cuda_jit has no CPU backend).
This data is later randomly sampled multiple times (cloned or maybe never used)