Numpyro: SVI leaking memory?

[nh2liu]

I'm having issues with numpyro, specifically SVI using GPU memory when running the same model multiple times.

I've set XLA_PYTHON_CLIENT_ALLOCATOR=platform as advised in jax docs to debug this further.

Here's a sample gist of the code I am using to debug this.

https://gist.github.com/nh2liu/2570812c5913a8e0889e746274188aba

The loop at the bottom runs a SVI for each iteration. Even though nothing is returned (python should dereference and GC), I'm seeing the GPU memory usage increasing on nvidia-smi and this script will eventually throw a jax out of memory when trying to allocate exception.

I've also tested this with a sample normal function with jax and did not see this behavior.

This is CUDA 10.0, jax=0.1.67, numpyro=0.2.4, jaxlib=0.1.47.

[neerajprad]

Thanks for the sample gist, @nh2liu. Do you also see this with the CPU? The only thing that I can think of is that JAX has a compilation cache which can grow without bound. I don't have a GPU available currently, but could you check if the number of jax.interpreters.xla.DeviceArray objects increases as the loop runs (using gc.get_objects)?

[fehiepsi]

JAX has a compilation cache which can grow without bound

Agreed! I think JAX users usually solve this issue by doing e.g. jit(train)(X, y, key).

[neerajprad]

I think JAX users usually solve this issue by doing e.g. jit(train)(X, y, key).

@fehiepsi - how does using the jit affect the issue?

[fehiepsi]

I think each time we run the train function, a new object svi (in particular, svi.update - which is used in scan) is created hence the new cache is created. By jit(train), only a cache for train is created. Other intermediate objects will not be re-cached. Does this make sense to you?

[nh2liu]

Thanks everyone!

@jit(train)(X, y, key) actually did solve the issue here, it is the JAX compilation cache causing the trouble. I never thought to jit the train function but yeah, there turns out to be a big speed difference as well.

I thought that by setting XLA_PYTHON_CLIENT_ALLOCATOR=platform, jax would automatically gc as advertised on the docs but apparently this cache is not affected by it.

[neerajprad]

Thanks for explaining, @fehiepsi. One thing to note is that this will rarely affect users who stick with the functional paradigm, since each function will have a single unique id unlike this case where we end up creating new objects each time the loop is called. I think using classes otherwise simplifies many more things for us, but this is something to be aware of. I think a simple solution would be to have an LRU cache eviction policy in JAX, which we can discuss with the JAX dev team.

[fehiepsi]

Yeah! Actually, it will rarely affect functional programmers who don't use lambda (e.g. lax.scan(lambda ...) will be a bad practice). :D