Memory Issues with Forward-Mode AD?

[coult099]

Hello,

I am running into an issue with Dr.Jit which I don't fully understand. I am following the forward inverse rendering tutorial, using the cbox scene in cuda_ad_spectral mode. I am using green.reflectance.values as the variable to take the gradient with respect to, leaving the rest of the code unchanged. I am running Mitsuba inside a python 3.8.10 virtual environment.

The issue I am experiencing: when I set spp above 256 (in mi.render()), or the image size above 256x256 in the base.xml file for cbox, I get the following warning:

jit_flush_malloc_cache(): Dr.Jit exhausted the available memory and had to flush its allocation cache to free up additional memory. You may want to change your computation so that it uses less memory. This warning will only be displayed once.

The image renders and outputs to a png, so I think the issue is with the gradient step dr.grad(image). Any insight as to why this might be happening is appreciated. This does not seem like the most memory-intensive task, so I think I must be doing something incorrectly. Thank you.

My system information

• Linux Ubuntu 20.04.1, x86_64
• Intel Core i7-6850K CPU, 3.60 GHz
• 130 GiB system memory
• 2x Nvidia Titan V
• 2x Nvidia Quadro RTX 8000
• I have set the environment variables NVIDIA_VISIBLE_DEVICES and CUDA_VISIBLE_DEVICES to point to the Quadro RTX cards only, but get the same error without specifying which GPU to use

[coult099]

A small edit: I get the same issue with cuda_ad_rgb as well, for a variety of differentiable parameters (colors as well as vertex positions for different meshes).

[njroussel]

I would recommend using the prb integrator.

You can do this by changing the rendering cell as follows:

integrator = mi.load_dict({ 'type': 'prb' })
image = mi.render(scene, params, integrator=integrator, spp=4096)

This should do the trick. Take a look here if you want a better understanding of the memory consumption with this method.

[coult099]

Thank you for the reply! This fixed my issue, but only up to spp=2048; once I set spp=4096 I get the same jit_flush_malloc_cache(): error. Is this expected for high spp values?

[njroussel]

It's expected with very large wavefront sizes (spp x resolution). However, in the tutorial - without any changes - 4096 should still be quite low. (I just tried it now on a GPU with 24GB of VRAM without any issues.)

[coult099]

Thank you.