how do I interpret these traces?

[jerry-kobold]

Hello JAX communit! I'm fairly new to JAX, and am trying to use it to do some physics simulations by connecting an existing model to the blackjax package. I have a function that looks like this:

import jax
import jax.numpy as jnp

def log_likelihood(*args):
    ... # compute a physical model and the log density
    return logdensity

log_ll_jax = jax.jit(log_likelihood)

I've omitted the internals of the log_likelihood function because that is based on our internal code. But the function computes the right result (as defined by our internal tests) and is jit-able. This function implements one forward step of a model that I want to now hook into blackjax to do HMC.

When I trace one run of the step like so:

jax.profiler.start_trace('/tmp/jax-profile/')
jax.block_until_ready(log_ll_jax(initial_position))
jax.profiler.stop_trace()

I get the following trace:

I've highlighted a bit from the trace that centers around the ThunkExecutor::Execute method in the client thread, which helpfully tells me that it's waiting for the computation to complete. One thing that jumps out at me is that the ThunkExecutor::Execute methods in the worker threads seem to be a bit sparse, but ok, fine. Seems generally reasonable.

Now I hook this up to the blackjax framework (using their example) and try to run the warmup step to determine the inverse mass matrix and step size:

rng_key, warmup_key, sample_key = jax.random.split(rng_key, 3)
warmup = blackjax.window_adaptation(blackjax.nuts, log_ll_jax)
# do only a few steps
(state, parameters), _ = warmup.run(warmup_key, initial_position, num_steps=5)
# start benchmark here to avoid capturing compilation
jax.profiler.start_trace('/tmp/jax-profile/')
jax.block_until_ready(state)
jax.profiler.stop_trace()

This runs just fine but seems to be very slow. When I look at the profiling output, I see the following:

I am using a 32-core CPU as my device, so this is a subset but what I am struggling to understand here is why the ThunkExecutor::Execute that is shown there in brown is waiting so long, when it looks naively like the computation is finished. What am I missing here and how can I interpret what's happening? I realize that this references blackjax but I think my question is more about how I should understand the difference between these traces rather than the specifics of the library.

[kamalraj23]

hey,
ur timeline is misleading b'cause u start the profiler after ur already enqueued the heavy BlackJAX warmup work. With JAX’s async/ dispatch, warmup.run(...) kicks off a big compiled computation and returns immediately,when u later call block_until_ready(state) u mostly see the host-side ThunkExecutor::Execute span “waiting for computation to complete” while the actual kernels ran before tracing began so the sparse worker threads,

check this and let me know

jax.profiler.start_trace('/tmp/jax-profile')
(state, params), _ = warmup.run(key, init_pos, num_steps=5)
jax.tree_util.tree_map(lambda x: x.block_until_ready(), (state, params))
jax.profiler.stop_trace()

##For diagnosis, -

import jax
jax.config.update("jax_async_dispatch", False)  # turn back on for performance later

##& to separate compile 4m execute on the 1st :

compiled = warmup.run.lower(key, init_pos, num_steps=5).compile()

jax.profiler.start_trace("/tmp/jax-profile")
(state, params), _ = compiled(key, init_pos, num_steps=5)
jax.tree_util.tree_map(lambda x: x.block_until_ready(), (state, params))
jax.profiler.stop_trace()

[jerry-kobold]

@kamalraj23 Thanks, appreciate your explanation. I was trying to avoid capturing the compilation step with my benchmarks but evidently I blundered into a different mistake. I will try what you suggest and let you know.