Slow compile-time for jax.jit() on my server

[AustinApple]

Hello,
I am trying to run colabfold code on my server. I found the running time is much longer than I run them on Colab. Meanwhile, I always encounter the warning

E external/xla/xla/service/slow_operation_alarm.cc:65] 
********************************
[Compiling module extracted] Very slow compile? If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.

Later, I found that the culprit would be the time for jax.jit() compilation. So, I did a simple test, the snippet is following:

import time
import jax
import numpy as np
import jax.numpy as jnp

pulse = jax.device_put(np.random.rand(10000))

def f(pulse):
    return jnp.fft.fft(pulse)

t0 = time.time()
a = jax.jit(f, backend='gpu')(pulse)
t1 = time.time()
print('compile-time cost on gpu:',t1-t0)


t0 = time.time()
a = jax.jit(f, backend='gpu')(pulse)
t1 = time.time()
print('running time cost on gpu:',t1-t0)


t0 = time.time()
a = jax.jit(f, backend='cpu')(pulse)
t1 = time.time()
print('running time cost on cpu:',t1-t0)

##############################################
Result
##############################################
For colab:
compile-time cost on gpu: 0.04041171073913574
running time cost on gpu: 0.0023126602172851562
running time cost on cpu: 0.02092766761779785

For my server:
compile-time cost on gpu: 4.001883268356323
running time cost on gpu: 0.0015799999237060547
running time cost on cpu: 0.027371883392333984

I found the compile time on my server is much longer than Colab. Is there any way to accelerate the compile time or is something wrong? Thanks.

The following is the envs for my server and Colab:

For Colab (my server):
os: Ubuntu 22.04.3 LTS (CentOS Linux release 7.9.2009 (Core))
python: 3.10.12 (3.10.13)
jax: 0.4.23 (0.4.23)
jaxlib: 0.4.23 (0.4.23)
GPU: Tesla T4 (A100)
Cuda: 12.2 (11.8)

[jakevdp]

Hi, thanks for the question! I'm not sure what might lead to longer compile times on your system vs Colab. Keep in mind that compilation happens on the host rather than the device, so perhaps your local CPU is slower than the CPU in the Colab GPU runtimes?

[AustinApple]

Thanks for the prompt response. I have AMD EPYC 7352 24-Core Processor on my server, while Colab uses Intel(R) Xeon(R) CPU @ 2.30GHz. Do you think it would be the reason?

[jakevdp]

I'm not sure about the relative computational specs of those two processors.

[AustinApple]

Hello,
I have tried the code on three different machines with different CPUs (AMD EPYC 7352 24-Core Processor, Intel(R) Xeon(R) Silver 4116 CPU @ 2.10GHz, and Intel(R) Xeon(R) Platinum 8168 CPU @ 2.70GHz). I only encounter the slow compile-time issue on the machine with AMD CPU. So I think the culprit would be the AMD CPU, but still do not know the exact reason.