Problem with the number of chains in MCMC_Posterior with method = 'nuts'

[MolinAlexei]

Hi,
I am running the following lines of code on the computing cluster I am using fro my work :

potential_fn, parameter_transform = likelihood_estimator_based_potential(
    density_estimator_SNLE, prior, x_o
)

# Build the posterior
posterior = MCMCPosterior(
    potential_fn, proposal=prior, theta_transform=parameter_transform, thin = 1, method = 'nuts', num_chains = 32, num_workers = 32,
    warmup_steps = 1000
)

samples = posterior.sample((5000,32))

and I am launching the code using SLURM with the following bash file :

#!/bin/bash

#SBATCH --job-name=PleaseWork
#SBATCH --exclusive 
#SBATCH --partition=xifu
#SBATCH --output=MCMC_SBI_xifusim.out

source /home/sila/miniconda3/etc/profile.d/conda.sh
conda activate myenv
###export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/home/mola/miniconda3/envs/myenv/lib

srun -n1 -c32 python MCMC_SBI_xifusim_obs.py

And I get the following warning :
/home/mola/miniconda3/envs/myenv/lib/python3.10/site-packages/pyro/infer/mcmc/api.py:497: UserWarning: num_chains=32 is more than available_cpu=31. Chains will be drawn sequentially.

However, if I choose 'slice_np' for the method, I get the proper behavior, i.e. the chains running in parallel with 32 workers.

Am I doing something wrong or is this normal ?
Thank you for the help.

[felixp8]

Hi!

The issue here seems to be that Pyro internally reserves one CPU for the main process and thus only allows for n_cpus - 1 chains to run in parallel, as you can see here. This is why in MCMCPosterior._pyro_mcmc the number of chains is set to one less than the cpu count also here, if num_chains=None. So, for Pyro samplers, I think you would need to reduce the chains you request or increase the CPUs you allot for your program.

[Baschdl]

@felixp8 How about adding this to the docs? Something like num_chains: ... Should generally be num_workers-1

[felixp8]

would make sense, I can add it to #1053