Performance Degradation When Upgrading JAX from 0.4.31 to 0.4.35 on CPU

[Routhleck]

After upgrading JAX from version 0.4.31 to 0.4.35, I observed a performance degradation when running the same program on a CPU. I write a test program includes two main computations: one that updates a progress bar using jax.debug.callback within a jax.jit function, and another that does not update the progress bar. Both computations showed slower execution times in JAX 0.4.35 compared to JAX 0.4.31.

Below is the test program I used to measure the performance:

import tqdm.auto
import time
import jax
import jax.numpy as jnp
import os

print(f'jax version: {jax.__version__}')

jax.config.update('jax_cpu_enable_async_dispatch', False)

TOTAL_STEP = 1000000

progress_bar = tqdm.auto.tqdm(range(TOTAL_STEP))
progress_bar.set_description("Processing")

@jax.jit
def compute_with_progress_bar():
    def body_fun(i, ret):
        ret += jnp.sum(jnp.square(jnp.arange(1000000))) + i
        # update the progress bar
        jax.debug.callback(lambda *args: progress_bar.update(), ())
        return ret, ret

    return jax.lax.scan(body_fun, 0, jnp.arange(TOTAL_STEP))[0]

start = time.time()
result = jax.block_until_ready(compute_with_progress_bar())
end = time.time()

progress_bar.close()

print()
print(f'time(with progress bar): {end - start} s')

@jax.jit
def compute():
    def body_fun(i, ret):
        ret += jnp.sum(jnp.square(jnp.arange(1000000))) + i
        # update the progress bar
        # jax.debug.callback(lambda *args: progress_bar.update(), ())
        return ret, ret

    return jax.lax.scan(body_fun, 0, jnp.arange(TOTAL_STEP))[0]

start = time.time()
result = jax.block_until_ready(compute())
end = time.time()

print()
print(f'time(without progress bar): {end - start} s')

Here are the execution times for both versions of JAX:

jax version: 0.4.31
Processing: 100%|██████████| 1000000/1000000 [00:13<00:00, 75304.32it/s]

time(with progress bar): 13.278408288955688 s

time(without progress bar): 0.045131683349609375 s

jax version: 0.4.35
Processing: 100%|██████████| 1000000/1000000 [00:24<00:00, 40211.97it/s]

time(with progress bar): 24.86821460723877 s

time(without progress bar): 0.06705331802368164 s

[hawkinsp]

Try with XLA_FLAGS=--xla_cpu_use_thunk_runtime=false?

We migrated to a newer version of the XLA:CPU runtime, which has different performance characteristics.

@ezhulenev @penpornk

[Routhleck]

Thank you for the suggestion. Indeed, there is some performance improvement compared to the previous results, but it is still slightly slower than JAX 0.4.31.

[patrick-kidger]

We've also had a report of an 8x slowdown here:

patrick-kidger/diffrax#518

In particular this identifies the offending diff as being between jaxlib==0.4.32.dev20240807 and jaxlib==0.4.32.dev20240812 .

[jdehning]

The flag XLA_FLAGS=--xla_cpu_use_thunk_runtime=false worked for me, for the slowdown I observed with diffrax: patrick-kidger/diffrax#518