Numpyro: Add progress bar for HMC num_chains > 1

[neerajprad]

Right now, it is very hard to ascertain progress when running MCMC with num_chains > 1. Let us figure out ways to show progress bar updates. Even a progress bar without diagnostic updates will be useful.

Currently, chain methods that support progress bar:

• [x] parallel
• [x] sequential
• [x] vectorized

With this, we can make the following enhancement.

• [x] Output other metadata like when the model is compiling and when it is sampling. This is especially useful when we run multiple chains and it is not clear whether all the time is being taken to compile the model or if it is the sampling part that is taking time.

[fehiepsi]

Hmm, is this possible? AFAIK, the progressbar will only show the progress of compiling the code, while the real execution is hidden. Hopefully, there will be some tools in JAX allowing us measure this progress.

[neerajprad]

is this possible?

Maybe not yet; it will likely depend on solutions for https://github.com/google/jax/issues/364. For sequential chains, it might be possible (though I am not fully sure).

[fehiepsi]

JAX now supports an experimental host_callback. ~It seems that this is doable now, but we need to study how that module works.~ Edit host_callback just collects intermediate results at the end of the process

[jeremiecoullon]

If it's helpful: this answered question in the JAX Discussion shows how to print the progress of an MCMC chain within a compiled function (using the host_callback module mentioned above).
This solution might need improvement when running several chains in parallel though.

[fehiepsi]

That's a great news @jeremiecoullon! Have you come up with a solution for MCMC chain yet? This would be one of the most requested feature from NumPyro users. :)

[jeremiecoullon]

So I did a quick test which seems to work.

Suggested solution:

In numpyro/util.py I added the following functions:

from jax.experimental import host_callback

def _print_consumer(arg, transrorm):
    i, n_iter = arg
    print(f"Iteration {i}/{n_iter}")

@jit
def progress_bar(arg, result):
    """
    Print progress of loop only if iteration number is a multiple of the print_rate

    Usage: carry = progress_bar((iter_num, n_iter, print_rate), carry)
    """
    i, n_iter, print_rate = arg
    result = lax.cond(
        i%print_rate==0,
        lambda _: host_callback.id_tap(_print_consumer, (i, n_iter), result=result),
        lambda _: result,
        operand=None)
    return result

Then I simply added the following to the fori_collect() function in the _body_fn() definition, just after idx = (i - start_idx) // thinning:

print_rate = int(upper/10)
i = progress_bar((i, upper, print_rate), i)

i is the current iteration number, upper is the number of iterations, and the print rate is set to be 1/10 of the number of iterations.

I tested it on the 8 schools example on the readme and set progress_bar=False: mcmc = MCMC(nuts_kernel, num_warmup=500, num_samples=1000000, progress_bar=False, num_chains=1). I ran for 1million samples so you have time to see the progress bar unfold (otherwise it's too fast).

The way it works is that progress_bar is an identity function: it takes in i as argument and returns i. A side effect of this function is that it calls host_callback.id_tap() which invokes the python print function only if the print_rate divides the number of iterations. The reason you have to pass in i (or any other variable that is dependent of the arguments of the main compiled function) is to make sure the python print function is called at the correct time.

Speed:

I did a quick speed benchmark (running a single chain of the 8 schools example for 1 million iterations), and the progress bar seems to make no difference to speed. This is because the python function is only called very 100K iterations. For every other iteration progress_bar() simply acts like the identity function.

Interaction with tqdm

The main problem with this solution is that it messes up the existing tqdm progress bar, as _body_fn() is also used in that case.

Suggested solution: define a decorator within the fori_collect() definition that adds the new progress bar to a function, and wrap _body_fn() in this decorator when passing it tofori_loop(). First import functools, then add the following within the fori_collect() definition:

def add_progress_bar(func):
        @functools.wraps(func)
        def wrapper_progress_bar(i, vals):
            print_rate = int(upper/10)
            i = progress_bar((i, upper, print_rate), i)
            return func(i, vals)
        return wrapper_progress_bar

....

if not progbar:
        last_val, collection, _, _ = fori_loop(0, upper, add_progress_bar(_body_fn), (init_val, collection, start_idx, thinning))

Other issues with this:

• This does work when running several chains in parallel, but it prints out the progress of every chain. So it looks cluttered and hard to read if you have many chains. Possible solution: include the chain number in each progress bar (example: "chain 3: Iteration 1000/10000"). Not sure how to do this as the chain number isn't included in each mcmc chain
• The progress bar is really basic and has a lot less info than tqdm. We could add a bit more to it (step size etc..) if that would be useful.
• The compiled version of the sampler (using fori_loop) is called when progress_bar=False in the mcmc object. If we add a progress bar to the compiled version as well this would be confusing. Suggestion: we could rename that argument to compile=True/False instead.
• The progress bar ignores the burn-in/warmup. We could add a seperate progress bar for burn-in (like in the tqdm case). The progress bar doesn't always print out multiples of 10/100/1000 etc (example: if you run 10000 samples with a burn-in of 500 the progress bar prints every 1050 iterations.
• _body_fn has a caching decorator. How does the progress bar (in decorator form) interact with this?

Showing when the sampler is compiling:

As mentioned in the first message of this issue, we could also add info about compilation. The hack I do when writing samplers in JAX is to simply add a python print function at the beginning of the function (print("Compiling..")) and then I add another print function at the end (print("Done.")). This will only be called when the function is compiling and is a simple way to see how long compilation takes vs actually running.

Example:

@partial(jit, static_argnums=(1,))
def my_sampler(key, num_samples):
    print("Compiling..")
    samples = function_that_gets_samples(key, num_samples)
    print("Done.")
    return samples

I'm just making a pull request for the progress bar within a decorator for this.
There may of course be loads more problems that I'm missing as I'm not familiar with the NumPyro codebase :)

[fehiepsi]

Oops, I missed your comment. As far as I see, your PR works as expected. So I think migrating it to tqdm would not be so complicated.

not familiar with the NumPyro codebase

I see. I'll play more to understand those issues and will get back to you soon.

[fehiepsi]

This issue has been resolved thanks to @jeremiecoullon !