Numpyro: Error in bnn.py for num-chains > 2

[ziatdinovmax]

When I try running the bnn.py example (without any modifications) with num-chains > 2 it throws the following error:

jupyter@pytorch-highmem1:~$ python3 bnn.py --num-chains=3
Traceback (most recent call last):
  File "bnn.py", line 145, in <module>
    main(args)
  File "bnn.py", line 105, in main
    samples = run_inference(model, args, rng_key, X, Y, D_H)
  File "bnn.py", line 64, in run_inference
    mcmc.run(rng_key, X, Y, D_H)
  File "/opt/anaconda3/lib/python3.7/site-packages/numpyro/infer/mcmc.py", line 644, in run
    rng_keys = random.split(rng_key, self.num_chains)
  File "/opt/anaconda3/lib/python3.7/site-packages/jax/random.py", line 194, in split
    return _split(key, num)
  File "/opt/anaconda3/lib/python3.7/site-packages/jax/api.py", line 148, in f_jitted
    out = xla.xla_call(flat_fun, *args_flat, device=device, backend=backend)
  File "/opt/anaconda3/lib/python3.7/site-packages/jax/core.py", line 591, in call_bind
    outs = primitive.impl(f, *args, **params)
  File "/opt/anaconda3/lib/python3.7/site-packages/jax/interpreters/xla.py", line 418, in _xla_call_impl
    compiled_fun = _xla_callable(fun, device, backend, *map(abstractify, args))
  File "/opt/anaconda3/lib/python3.7/site-packages/jax/linear_util.py", line 217, in cached_fun
    ans, f_prev = cached_fun_body(f, args)
  File "/opt/anaconda3/lib/python3.7/site-packages/jax/linear_util.py", line 214, in cached_fun_body
    return call(f, *args), f
  File "/opt/anaconda3/lib/python3.7/site-packages/jax/interpreters/xla.py", line 434, in _xla_callable
    jaxpr, (pvals, consts, env) = pe.trace_to_subjaxpr(fun, master, False).call_wrapped(pvals)
  File "/opt/anaconda3/lib/python3.7/site-packages/jax/linear_util.py", line 165, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/opt/anaconda3/lib/python3.7/site-packages/jax/random.py", line 199, in _split
    return lax.reshape(threefry_2x32(key, counts), (num, 2))
  File "/opt/anaconda3/lib/python3.7/site-packages/jax/api.py", line 148, in f_jitted
    out = xla.xla_call(flat_fun, *args_flat, device=device, backend=backend)
  File "/opt/anaconda3/lib/python3.7/site-packages/jax/core.py", line 594, in call_bind
    outs = map(full_lower, top_trace.process_call(primitive, f, tracers, params))
  File "/opt/anaconda3/lib/python3.7/site-packages/jax/interpreters/partial_eval.py", line 111, in process_call
    out_flat = call_primitive.bind(fun, *in_consts, **params)
  File "/opt/anaconda3/lib/python3.7/site-packages/jax/core.py", line 591, in call_bind
    outs = primitive.impl(f, *args, **params)
  File "/opt/anaconda3/lib/python3.7/site-packages/jax/interpreters/xla.py", line 418, in _xla_call_impl
    compiled_fun = _xla_callable(fun, device, backend, *map(abstractify, args))
  File "/opt/anaconda3/lib/python3.7/site-packages/jax/linear_util.py", line 217, in cached_fun
    ans, f_prev = cached_fun_body(f, args)
  File "/opt/anaconda3/lib/python3.7/site-packages/jax/linear_util.py", line 214, in cached_fun_body
    return call(f, *args), f
  File "/opt/anaconda3/lib/python3.7/site-packages/jax/interpreters/xla.py", line 434, in _xla_callable
    jaxpr, (pvals, consts, env) = pe.trace_to_subjaxpr(fun, master, False).call_wrapped(pvals)
  File "/opt/anaconda3/lib/python3.7/site-packages/jax/linear_util.py", line 165, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "/opt/anaconda3/lib/python3.7/site-packages/jax/random.py", line 107, in threefry_2x32
    key1, key2 = keypair
ValueError: too many values to unpack (expected 2)

The system is a deep learning VM (16 CPUs) on the google cloud platform. I'm using a CPU version of jaxlib (for whatever reason the bnn.py is running much slower on GPU).

Thanks for your help in advance!

[fehiepsi]

@ziatdinovmax, thanks for the bug report! This issue should be fixed with #443.

[ziatdinovmax]

@fehiepsi, Thanks so much for fixing it so quickly! The MCMC inference part now works perfectly on multiple CPUs! There is, however, another problem in that the prediction part is taking much longer when num_chains > 1 and as a result, we don't gain anything from using multiple CPUs. See example below:

Running bnn.py with num-chains=1. MCMC time is 74 s, overall elapsed time is 78 s.

jupyter@pytorch-highmem1:~$ /usr/bin/time -v python3 bnn.py --num-data=1000 --num-chains=1 --num-warmup=200 --num-samples=800
sample: 100%|██████████████████████████████████████████████| 1000/1000 [01:06<00:00, 14.94it/s, 1023 steps of size 1.68e-03. acc. prob=0.90]
MCMC elapsed time: 75.20383715629578
        Command being timed: "python3 bnn.py --num-data=1000 --num-chains=1 --num-warmup=200 --num-samples=800"
        User time (seconds): 78.49
        System time (seconds): 0.56
        Percent of CPU this job got: 101%
        Elapsed (wall clock) time (h:mm:ss or m:ss): 1:18.14
        Average shared text size (kbytes): 0
        Average unshared data size (kbytes): 0
        Average stack size (kbytes): 0
        Average total size (kbytes): 0
        Maximum resident set size (kbytes): 374840
        Average resident set size (kbytes): 0
        Major (requiring I/O) page faults: 0
        Minor (reclaiming a frame) page faults: 99367
        Voluntary context switches: 28806
        Involuntary context switches: 295
        Swaps: 0
        File system inputs: 0
        File system outputs: 72
        Socket messages sent: 0
        Socket messages received: 0
        Signals delivered: 0
        Page size (bytes): 4096
        Exit status: 0

Now running with num-chains=4. MCMC time is only 17 s, but the overal elapsed time is 79 s (!)

jupyter@pytorch-highmem1:~$ /usr/bin/time -v python3 bnn.py --num-data=1000 --num-chains=4 --num-warmup=200 --num-samples=800

MCMC elapsed time: 17.690608501434326
        Command being timed: "python3 bnn.py --num-data=1000 --num-chains=4 --num-warmup=200 --num-samples=800"
        User time (seconds): 254.37
        System time (seconds): 0.43
        Percent of CPU this job got: 321%
        Elapsed (wall clock) time (h:mm:ss or m:ss): 1:19.21
        Average shared text size (kbytes): 0
        Average unshared data size (kbytes): 0
        Average stack size (kbytes): 0
        Average total size (kbytes): 0
        Maximum resident set size (kbytes): 467144
        Average resident set size (kbytes): 0
        Major (requiring I/O) page faults: 0
        Minor (reclaiming a frame) page faults: 127788
        Voluntary context switches: 15475
        Involuntary context switches: 695
        Swaps: 0
        File system inputs: 0
        File system outputs: 72
        Socket messages sent: 0
        Socket messages received: 0
        Signals delivered: 0
        Page size (bytes): 4096
        Exit status: 0

Sorry if I'm asking something stupid :)

[fehiepsi]

@ziatdinovmax That's a great question! Actually, using 4 chains will give you 4x number of samples so the same Elapsed (wall clock) time means that you have got 4x samples / s. :)

The reported number 17s is wrong (due to asynchronous dispatch, actual computation is still running). If you add a line mcmc.print_summary() after mcmc.run(rng_key, X, Y, D_H), you will get a correct report. Somehow, we mistakenly removed that line I guess. :( We'll address it soon. Thanks for bringing this up!

[ziatdinovmax]

Ahhhh... okay, that was a really stupid question then :-) Thanks for answering!