Performance Benchmarking: Unstable Timing for JIT-compiled Matrix Multiplication

[juliusuberall]

While doing microbenchmarking in JAX, we ran into the issue of measuring unstable wall-clock time.

We ended up creating a minimal case to evaluate this behaviour further - the code is below. We have a function doing a matrix multiplication, trace and JIT compile the function and then do 100K independent runs for benchmarking.

We tested this both on CPU and GPU using Google Collab. In both cases the measured time between runs is fluctuating heavily, which makes it difficult to do reliable benchmarking. Find attached the time plots for GPU and CPU.

Is this an expected behaviour ? Are we missing something essential ? Did people encounter this before while benchmarking with JAX ?
Thankful for any ideas or suggestions!

import jax
import time
import numpy as np

@jax.jit
def forward(a, b):
    return a @ b

x = jax.random.normal(jax.random.PRNGKey(0), (128,128))  # input
forward(x,x).block_until_ready()  # warmup

times = []
for _ in range(100000):
    start = time.perf_counter()
    forward(x,x).block_until_ready() 
    end = time.perf_counter()
    times.append((end - start) * 1e6)

print(f"Mean: {np.mean(times):.2f} µs, Std Dev: {np.std(times):.2f} µs")

[dfm]

Never mind my comment about asynchronous dispatch: I somehow didn't see your blocking in my initial read of the code! I'll take a look later today to see if I have more relevant suggestions.