Performance degradation in JAX compilation on consecutive runs

[Qazalbash]

I am testing a JAX-based code (fairly complex) and have observed inconsistency in compilation time and execution behavior. On the first run, the compilation completes quickly, and both the CPU and GPU engage in intensive computations within a short period before execution proceeds smoothly. However, if I restart the exact same run without making any modifications, the compilation time increases by approximately fivefold. Additionally, even after compilation, the execution introduces pauses, seemingly shifting work to the CPU intermittently.

Could someone provide insights into the cause of this behavior and possible strategies to mitigate it?

Fast compilation run

It took around 130 seconds to compile and took around 60 seconds per iteration.

Global Tuning: 100%|██████████| 10/10 [10:58<00:00, 65.89s/it]
Global Sampling: 100%|██████████| 10/10 [09:52<00:00, 59.25s/it]

Slow compilation run

It took around 550 seconds to compile and took around the same 60 seconds per iteration.

Global Tuning: 100%|██████████| 10/10 [18:08<00:00, 108.87s/it]
Global Sampling: 100%|██████████| 10/10 [09:56<00:00, 59.69s/it]

Environment variables

#!/bin/bash

export CUDA_VISIBLE_DEVICES=0
export NPROC=4
export XLA_FLAGS="--xla_cpu_multi_thread_eigen=false --xla_gpu_enable_latency_hiding_scheduler=true"
export XLA_PYTHON_CLIENT_ALLOCATOR=platform
export JAX_COMPILATION_CACHE_DIR="$HOME/jax_cache"
export JAX_ENABLE_X64=1

nvidia-smi

Mon Feb 24 15:05:14 2025
+---------------------------------------------------------------------------------------+
| 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 TITAN RTX               Off | 00000000:01:00.0 Off |                  N/A |
| 41%   43C    P2              66W / 280W |   4696MiB / 24576MiB |      0%      Default |
|                                         |                      |                  N/A |
+-----------------------------------------+----------------------+----------------------+

+---------------------------------------------------------------------------------------+
| Processes:                                                                            |
|  GPU   GI   CI        PID   Type   Process name                            GPU Memory |
|        ID   ID                                                             Usage      |
|=======================================================================================|
|    0   N/A  N/A      1432      G   /usr/lib/xorg/Xorg                            9MiB |
|    0   N/A  N/A      1565      G   /usr/bin/gnome-shell                          4MiB |
+---------------------------------------------------------------------------------------+