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.