Get the warning: W external/xla/xla/service/hlo_rematerialization.cc:2194] Can't reduce memory use below 10.96GiB (11774361600 bytes) by remateri alization; only reduced to 11.62GiB (12478088656 bytes)

[JyChang012]

I got the following warning when benchmarking Hugging Face transformers' Bert model.

W external/xla/xla/service/hlo_rematerialization.cc:2194] Can't reduce memory use below 10.96GiB (11774361600 bytes) by rematerialization; only reduced to 11.62GiB (12478088656 bytes)                                                            

It seems that XLA is trying to reduce memory usage to 10.96GB by rematerialization, but my GPU has 16GB VRAM. I also set XLA_PYTHON_CLIENT_MEM_FRACTION=.97 and nvidia-smi shows that JAX pre-allocate ~15GB VRAM. So why does it need to rematerialize? Would it hurt performance?

What jax/jaxlib version are you using?

jax                      0.4.8
jaxlib                   0.4.7+cuda11.cudnn86

Which accelerator(s) are you using?

Tesla T4

Additional system info
Python 3.9.16 on Linux

NVIDIA GPU info

Thu Apr 20 08:36:42 2023
+-----------------------------------------------------------------------------+
| NVIDIA-SMI 515.65.01    Driver Version: 515.65.01    CUDA Version: 11.7     |
|-------------------------------+----------------------+----------------------+
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |
|                               |                      |               MIG M. |
|===============================+======================+======================|
|   0  Tesla T4            On   | 00000000:00:1E.0 Off |                    0 |
| N/A   36C    P0    31W /  70W |  14947MiB / 15360MiB |      0%      Default |
|                               |                      |                  N/A |
+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+
| Processes:                                                                  |
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |
|        ID   ID                                                   Usage      |
|=============================================================================|
|    0   N/A  N/A    117742      C   python                          14943MiB |
+-----------------------------------------------------------------------------+

[Kelvinlby]

I met exactly the same problem on my RTX3090.
And I'm using the latest version of JAX.

[IrishWhiskey]

I have the same issue, using jax v0.4.33 with cuda12. I didn't have this issue when using jax v0.4.20 with cuda11.

However I also get the following error message:

external/xla/xla/service/gpu/nvptx_compiler.cc:893] The NVIDIA driver's CUDA version is 12.4 which is older than the PTX compiler version 12.6.77. Because the driver is older than the PTX compiler version, XLA is disabling parallel compilation, which may slow down compilation. You should update your NVIDIA driver or use the NVIDIA-provided CUDA forward compatibility packages.

I wonder whether having an old cuda version may be causing the re-materialization issue

[olive004]

I have the same problem. My GPU is always using 18607MiB / 24564MiB even if I reduce the size of my simulation (eg number of samples).

I am using the following versions of things:

jax==0.4.29
jaxlib==0.4.29+cuda12.cudnn91
jaxtyping==0.2.34

CUDA Version = 12.2.0
Driver = 535.183.01

Ubuntu = 20.04
Python = 3.10

+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 535.183.01             Driver Version: 535.183.01   CUDA Version: 12.2     |
|-----------------------------------------+----------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id        Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |         Memory-Usage | GPU-Util  Compute M. |
|                                         |                      |               MIG M. |
|=========================================+======================+======================|
|   0  NVIDIA GeForce RTX 4090        Off | 00000000:17:00.0 Off |                  Off |
| 31%   41C    P2              77W / 450W |  18607MiB / 24564MiB |     72%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
|   1  NVIDIA GeForce RTX 4090        Off | 00000000:4E:00.0 Off |                  Off |
| 31%   31C    P8              20W / 450W |    403MiB / 24564MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+
                                                                                         
+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
+---------------------------------------------------------------------------------------+

I also have the following environmental variables set:

os.environ['EQX_ON_ERROR']='breakpoint'
os.environ['JAX_DISABLE_JIT']='1'
os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"

I never get this problem on my other PC, which runs Ubuntu 22.04.