Memory leak with GL and Cuda rasterizer contexts

[kirilllzaitsev]

I am trying to debug the problem of ever-growing RAM after initializing a rasterizer context. Although no renderer functionality is used, the memory steadily grows in both single-GPU and nn.DataParallel environments:

• single GPU

• multiple GPUs

Following a proposed solution to a similar issue at #23, the rendering code is executed by explicitly created threads:

            self.glctx = None
            gpu_ids = list(range(torch.cuda.device_count()))
            self.rasterizer = Rasterizer(Dispatcher(gpu_ids=gpu_ids), device=self.device)
            if self.use_ddp:
                self.rasterizer = torch.nn.DataParallel(self.rasterizer, gpu_ids)

and

class Dispatcher:
    # https://github.com/NVlabs/nvdiffrast/issues/23
    def __init__(self, gpu_ids):
        self.threads = {}
        self.events = {}
        self.funcs = {}
        self.return_events = {}
        self.return_values = {}

        for gpu_id in gpu_ids:
            device = torch.device(gpu_id)
            self.events[device] = threading.Event()
            self.return_events[device] = threading.Event()
            self.threads[device] = threading.Thread(
                target=Dispatcher.worker,
                args=(
                    self,
                    device,
                ),
                daemon=True,
            )
            self.threads[device].start()

    @staticmethod
    def worker(self, device):
        torch.cuda.set_device(device)  <- has to be set to prevent device mistmatches
        ctx = dr.RasterizeCudaContext(device=device)  <- same problem for GL/Cuda context
        while True:
            self.events[device].wait()
            assert device not in self.return_values
            self.return_values[device] = self.funcs[device](ctx)
            del self.funcs[device]
            self.events[device].clear()
            self.return_events[device].set()

    def __call__(self, device, func):
        assert device not in self.funcs
        self.funcs[device] = func
        self.events[device].set()
        self.return_events[device].wait()
        ret_val = self.return_values[device]
        del self.return_values[device]
        self.return_events[device].clear()
        return ret_val


class Rasterizer(torch.nn.Module):
    def __init__(self, dispatcher, device):
        super().__init__()
        self.dispatcher = dispatcher
        self.device = device

    def forward(self, pos, tri, resolution, mesh_tensors, pts_cam):
        def func(ctx):
            has_tex = "tex" in mesh_tensors
            rast_out, _ = dr.rasterize(ctx, pos=pos, tri=tri, resolution=resolution)
            xyz_map, _ = dr.interpolate(pts_cam, rast_out, tri)
            depth = xyz_map[..., 2]
            device = pos.device
            if has_tex:
                texc, _ = dr.interpolate(mesh_tensors["uv"].to(device), rast_out, mesh_tensors["uv_idx"].to(device))
                color = dr.texture(mesh_tensors["tex"].to(device), texc, filter_mode="linear")
            else:
                color, _ = dr.interpolate(mesh_tensors["vertex_color"].to(device), rast_out, tri)
            return {
                "depth": depth,
                "color": color,
                "rast_out": rast_out,
                "xyz_map": xyz_map,
            }

        return self.dispatcher(pos.device, func)

At the same time, trying to get more logs with dr.set_log_level(0) in the single-GPU setup results in a deadlock.

The ultimate goal is to render in a distributed training setup, not necessarily with differentiation. Could you give some guidance on which things to check?

[s-laine]

I don't immediately see a problem with the code. I'd be surprised if this were a memory leak on the C++ side of nvdiffrast, so I'm thinking that the memory must be consumed by Python objects. There are three scenarios I can think of: 1) Everything is well but Python just isn't returning memory to OS as soon as you'd expect, 2) some Python objects aren't freed when dropping out of scope, even though they are unreachable (i.e., there's a reference loop), or 3) there is a genuine memory leak with more and more reachable objects being accumulated.

There's no bulletproof way to check if it's case 1 because you cannot force Python to return allocated memory to the OS, but in my experience Python reuses memory quite aggressively, so this feels like an unlikely scenario. To check if it's case 2, you can try calling gc.collect() periodically, which resolves potential reference loops and frees the memory internally. If memory usage still climbs, you can try to check if objects are being accumulated in memory by calling gc.get_objects() and analyzing the difference between iterations. Just be careful to avoid creating and keeping around additional references while doing this.

[kirilllzaitsev]

Thank you. Haven't looked at the object lists yet, but adding gc.collect() after processing each batch did not solve it, the memory being consumed more slowly but steadily:

Aside from the RAM issue, one thing that could be useful for someone trying nvdiffrast with multiple GPUs is building the kernel beforehand to avoid deadlocks later:

        if torch.distributed.get_rank() == 0:
            import nvdiffrast.torch as dr
            ctx = dr.RasterizeGLContext(device=torch.device("cuda:0")) 
            dummy_pos = torch.zeros(1, 3, 4, device="cuda:0")
            dummy_tri = torch.zeros(1, 3, dtype=torch.int32, device="cuda:0")
            dummy_rast, _ = dr.rasterize(ctx, dummy_pos, dummy_tri, [1, 1])
            dummy_attr = torch.zeros(1, 3, 4, device="cuda:0")
            _ = dr.interpolate(dummy_attr, dummy_rast, dummy_tri)
        torch.distributed.barrier()