Understanding the XLA_PYTHON_CLIENT_MEM_FRACTION variable related to GPU Memory Allocation

[sh0416]

Hello, I am wondering the mechanism of how the GPU memory is used in jax.

As the XLA_PYTHON_CLIENT_MEM_FRACTION is set to 0.9 as default, the about 33G for each GPU VRAM is preallocated when the script starts in 2xA100-40G.

When I load LLaMA 30b checkpoint (, which is 61GB for float16 precision), OOM occurs although it seems there is enough capacity to load the checkpoint.

So, I change the default XLA_PYTHON_CLIENT_MEM_FRACTION to 0.99, then I can successfully load the checkpoint.

There are two hypothesis that can explain this phonomena.

1. jax doesn't use the additional memory beyond XLA_PYTHON_CLIENT_MEM_FRACTION.
2. jax tried to use the additional memory, but the memory fragmentation occurs and OOM occurs while handling it.

I think the latter is true, but I ask to this community for clarity.

Thank you

[steveepreston]

same problem i stuck in.

[BDEvan5]

JAX preallocates the XLA_PYTHON_CLIENT_MEM_FRACTION when the program starts executing irrespective of what you try to load into memory. They say the reason for this is to prevent fragmentation.

If you need more memory than the default of 75%, then just increase the value to something like 0.99. Then it will preallocate more memory. It can be useful to turn of preallocation while debugging so that you can see how full the memory is. You can turn off preallocation by running export XLA_PYTHON_CLIENT_PREALLOCATE=false.

For reference see: https://jax.readthedocs.io/en/latest/gpu_memory_allocation.html