Question about `set_host_device_count` in Numpyro?

numpyro
1

I am new to Numpyro and am trying out the introduction by Dan Foreman-Mackey.

I have a question about numpyro.set_host_device_count(): from what I understand in the documentation, this function allow me to specify the number of CPUs to use on a single node. Since I only have one CPU with 32 cores, I think I wouldn’t need to set this.

However, I’ve noticed that with two chains, setting numpyro.set_host_device_count(2) makes the process about two times faster. It appears that the chains run sequentially by default, but can be accelerated by running in parallel.

My question is, what’s the best practice: should I set set_host_device_count to the number of available CPUs or to the number of chains?

2

This is more of a JAX question. By default, all CPU cores are considered as one XLA device by JAX. Under the hood, set_host_device_count() just sets an environment variable xla_force_host_platform_device_count. See here for more details.

I’d recommend using JAX config directly e.g., jax.config.update("jax_num_cpu_devices", 8) as it will explicitly throw an error if you try to set it after initialization.