How to Disable Preallocation of GPU Memory in JAX

[Olin1461]

I notice that JAX will preallocate 75% of the total GPU memory when the first JAX operation is run. After I set up os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false', JAX no longer preallocates GPU memory after JAX operations in my script. However, I still find that memory preallocation occurs when I use a function from another package that utilizes operations in JAX.

My question is how I can turn off memory preallocation when I use functions from other packages?

Thanks!

[Qazalbash]

You can turn off memory pre-allocation when using functions from other packages by setting the XLA_PYTHON_CLIENT_PREALLOCATE environment variable to 'false' before importing JAX or any packages that use JAX. This should prevent preallocation for all JAX operations, including those called by other packages.

Here's an example:

import os
os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false'
import jax
import package_that_uses_jax

Alternatively, you can also set this environment variable outside of your Python script, before running it. This will affect all JAX operations in your script, regardless of whether they come from JAX directly or from other packages.

export XLA_PYTHON_CLIENT_PREALLOCATE=false

Let me know if this helps!

[jakevdp]

To add to this: the key is that the XLA environment flags must be set before JAX is imported; setting them after will have no effect.