Memory leak in ppisp_cuda kernel (~18MB/step) causes OOM during training

[ehatami65]

Description

When using PPISP in a training loop, GPU memory grows linearly at approximately 18MB per training step, causing OOM errors after ~2000 steps on a 46GB GPU. The leak persists regardless of:

• Detaching the PPISP output from the computation graph
• Disabling regularization loss
• Calling torch.cuda.synchronize(), torch.cuda.empty_cache(), and gc.collect() after each forward pass

The leak appears to be inside the ppisp_cuda CUDA kernel rather than in PyTorch's autograd system.

Environment

• PPISP version: 1.0.0
• PyTorch version: 2.5.1
• CUDA version: 12.4
• GPU: NVIDIA A6000 (46GB)
• OS: Ubuntu Linux 6.8.0

Minimal Reproduction

import torch
from ppisp import PPISP, PPISPConfig

# Initialize PPISP
ppisp = PPISP(num_cameras=1, num_frames=100, config=PPISPConfig(use_controller=False))
ppisp = ppisp.cuda()
optimizers = ppisp.create_optimizers()

# Simulate training loop
height, width = 540, 960
pixel_y, pixel_x = torch.meshgrid(
    torch.arange(height, device="cuda", dtype=torch.float32) + 0.5,
    torch.arange(width, device="cuda", dtype=torch.float32) + 0.5,
    indexing="ij",
)
pixel_coords = torch.stack([pixel_x, pixel_y], dim=-1)

for step in range(3000):
    # Simulate rendered RGB from Gaussian splatting
    rgb_in = torch.rand(1, height, width, 3, device="cuda", requires_grad=True)

    # Apply PPISP
    rgb_out = ppisp(
        rgb=rgb_in,
        pixel_coords=pixel_coords,
        resolution=(width, height),
        camera_idx=0,
        frame_idx=step % 100,
    )

    # Compute loss and backward
    loss = (rgb_out - torch.rand_like(rgb_out)).pow(2).mean()
    loss.backward()

    for opt in optimizers:
        opt.step()
        opt.zero_grad(set_to_none=True)

    # Memory debug
    if step % 100 == 0:
        alloc = torch.cuda.memory_allocated() / 1024**3
        print(f"Step {step}: {alloc:.2f} GB")

    # Cleanup
    del rgb_in, rgb_out, loss

Expected: Memory stable around 1-2 GB
Actual: Memory grows ~18MB/step, OOM around step 2000

Experimental Evidence

I ran 6 controlled experiments to isolate the leak:

Experiment	Configuration	Result
No PPISP	post_processing=None	Stable at ~1.7GB (8000 steps)
With PPISP	Default	OOM at step ~2000
Detached output	rgb.detach().requires_grad_(True) after PPISP	OOM at step ~150
No reg loss	Skip get_regularization_loss()	OOM at step ~2000
Skip forward	PPISP module initialized but forward() not called	Stable at ~1.7GB (8000 steps)
Aggressive cleanup	synchronize + empty_cache + gc.collect after each call	OOM at step ~2000

Key finding: Only experiments that don't call PPISP forward have stable memory.

Memory Growth Pattern

Step 0:    0.93 GB
Step 50:   16.28 GB  (+15.35 GB)
Step 100:  28.96 GB  (+12.68 GB)
Step 150:  40.16 GB  (+11.20 GB)
Step 165:  OOM (43+ GB)

Analysis

The leak is NOT caused by:

• PyTorch autograd graph retention (detaching doesn't help)
• PPISP regularization loss computation
• Python garbage collection issues

The leak IS caused by something in _PPISPFunction.forward() or the underlying _C.ppisp_forward() CUDA kernel. The ctx.save_for_backward() tensors should be released after backward, but something prevents this.

Workaround

Currently disabling PPISP entirely as a workaround.

[nvibd]

Thank you for the detailed report. I will try to reproduce this locally and report back.

[nvibd]

BTW, have you already tried debugging inside of the forward or CUDA kernel? Since your analysis already zeroes in on that pretty closely.

[ehatami65]

Debugging Update: Leak Confirmed in CUDA Kernel

Thanks for following up. To answer your question: I've debugged as far as possible from the Python side, but I cannot debug inside the CUDA kernel itself because ppisp_cuda is a pre-compiled .so binary and I don't have access to the source code.

However, I've isolated the leak definitively to _C.ppisp_forward().

Test: Direct CUDA Kernel Call (Bypassing Autograd)

I created a minimal test that calls _C.ppisp_forward() directly, completely bypassing torch.autograd.Function:

# Direct CUDA kernel call - no autograd wrapper
rgb_out = _C.ppisp_forward(
    exposure_params, vignetting_params, color_params, crf_params,
    rgb_in, pixel_coords, width, height, camera_idx, frame_idx
)

Results:

Test	Memory Growth
Control (same tensors, no PPISP)	0 MB over 100 iterations
_C.ppisp_forward() only	+5.9 MB per call

Step    0: after_cleanup=16.1MB,  total_growth=+12.1MB
Step   40: after_cleanup=490.7MB, total_growth=+486.7MB
Step  100: after_cleanup=1202.6MB, total_growth=+1198.6MB
Step  199: after_cleanup=2377.2MB, total_growth=+2373.3MB

Avg leak per iteration: 11.87 MB

Memory is NOT freed even after:

del rgb_out
torch.cuda.synchronize()
torch.cuda.empty_cache()
gc.collect()

No Cleanup Function Available

I checked the exposed API:

>>> import torch
>>> import ppisp_cuda as _C
>>> [x for x in dir(_C) if not x.startswith('__')]
['ppisp_backward', 'ppisp_forward']

Only ppisp_forward and ppisp_backward are exposed. There's no ppisp_cleanup() or similar function to release internal state.

The leak is 100% inside the CUDA kernel, not in PyTorch autograd or Python. The test bypasses all Python layers entirely.
The leak signature suggests:

• An internal buffer/cache that grows without bounds
• Missing cudaFree() for intermediate allocations
• A memory pool that never releases back to the device

Full test script

#!/usr/bin/env python
"""Minimal reproduction of PPISP CUDA kernel memory leak."""

import gc
import torch
import ppisp_cuda as _C

def get_memory_mb():
    torch.cuda.synchronize()
    return torch.cuda.memory_allocated() / 1024**2

def test_ppisp_forward_leak(num_iterations=200, print_every=10):
    height, width = 540, 960
    num_pixels = height * width
    num_cameras, num_frames = 1, 100

    exposure_params = torch.zeros(num_frames, device="cuda", dtype=torch.float32)
    vignetting_params = torch.zeros(num_cameras, 3, 5, device="cuda", dtype=torch.float32)
    color_params = torch.zeros(num_frames, 8, device="cuda", dtype=torch.float32)
    crf_params = torch.zeros(num_cameras, 3, 4, device="cuda", dtype=torch.float32)

    pixel_y, pixel_x = torch.meshgrid(
        torch.arange(height, device="cuda", dtype=torch.float32) + 0.5,
        torch.arange(width, device="cuda", dtype=torch.float32) + 0.5,
        indexing="ij",
    )
    pixel_coords = torch.stack([pixel_x.flatten(), pixel_y.flatten()], dim=-1).contiguous()

    torch.cuda.empty_cache()
    gc.collect()
    baseline_mb = get_memory_mb()
    print(f"Baseline: {baseline_mb:.2f} MB")

    for i in range(num_iterations):
        rgb_in = torch.rand(num_pixels, 3, device="cuda", dtype=torch.float32)

        rgb_out = _C.ppisp_forward(
            exposure_params, vignetting_params, color_params, crf_params,
            rgb_in, pixel_coords, width, height, 0, i % num_frames,
        )

        del rgb_out, rgb_in
        torch.cuda.synchronize()
        torch.cuda.empty_cache()
        gc.collect()

        if i % print_every == 0:
            mem = get_memory_mb()
            print(f"Step {i}: {mem:.1f}MB (growth: {mem - baseline_mb:+.1f}MB)")

    print(f"\nTotal leaked: {get_memory_mb() - baseline_mb:.2f} MB")

if __name__ == "__main__":
    test_ppisp_forward_leak()

[nvibd]

The full source code is available here: https://github.com/nv-tlabs/ppisp/tree/main/ppisp .
The CUDA source code is under src/.

[ehatami65]

fixed now! I believe the issue was pybind11 kept internal references to the returned tensors. Python/PyTorch
couldn't garbage collect them because pybind11 was still holding onto them.

The clue:
Leaked memory per call (~6 MB) matched exactly the output tensor size:
540 × 960 × 3 × 4 bytes = 6.22 MB

The fix:
py::return_value_policy::move tells pybind11 to transfer full ownership to Python, so tensors get freed
properly.