Parallel Chains Problem After Updating

[JoeSu112]

Hi all,
I recently updated HSSM to the newest version (0.2.1). However, while I run the program that was able to run in the old version (0.2.0b1), some problem occurs:

infer_DDM = DDM.sample(
    sampler="nuts_numpyro",  # type of sampler to choose,
    cores=2,  # how many cores to use
    chains=chains,  # how many chains to run
    draws=draws,  # number of draws from the markov chain
    tune=int(draws*0.2),  # number of burn-in samples
    idata_kwargs=dict(log_likelihood=False)

UserWarning: There are not enough devices to run parallel chains: expected 4 but got 2. 
Chains will be drawn sequentially. If you are running MCMC in CPU, 
consider using `numpyro.set_host_device_count(4)` at the beginning of your program. 
You can double-check how many devices are available in your system using `jax.local_device_count()`.

I'm sure that I could run parallel chains in the past versions. So I'm wondering how I can solve this problem and why this did not happen in the past versions.
Thanks for any help!!!

[JoeSu112]

Updated:
When I do the sampling again (on the same node), the above warning did not appear! It became as followed:

CUDA backend failed to initialize: Unable to load CUDA. Is it installed? (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
/home/miniconda3/lib/python3.11/site-packages/jax/_src/numpy/array_methods.py:66: UserWarning: 
Explicitly requested dtype float64 requested in astype is not available, and will be truncated to dtype float32. 
To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. 
See https://github.com/google/jax#current-gotchas for more.
  return lax_numpy.astype(arr, dtype)
/home/miniconda3/lib/python3.11/site-packages/jax/_src/numpy/array_methods.py:66: UserWarning: 
Explicitly requested dtype float64 requested in astype is not available, and will be truncated to dtype float32. 
To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. 
See https://github.com/google/jax#current-gotchas for more.
  return lax_numpy.astype(arr, dtype)
/home/b08303030/miniconda3/lib/python3.11/site-packages/jax/_src/numpy/array_methods.py:66: UserWarning: 
Explicitly requested dtype float64 requested in astype is not available, and will be truncated to dtype float32. 
To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. 
See https://github.com/google/jax#current-gotchas for more.
  return lax_numpy.astype(arr, dtype)
Compiling.. :   0%|                                        | 0/30 [00:00<?, ?it/s
/home/b08303030/miniconda3/lib/python3.11/site-packages/jax/_src/numpy/array_methods.py:66: UserWarning: 
Explicitly requested dtype float64 requested in astype is not available, and will be truncated to dtype float32. 
To enable more dtypes, set the jax_enable_x64 configuration option or the JAX_ENABLE_X64 shell environment variable. 
See https://github.com/google/jax#current-gotchas for more.

And it can run parallel chains again! However, the estimation time become a lot larger than expected.
In the past version, it took about 3-4 hours to estimate, but in this it takes more than a day.
It makes me really confused. I'm wondering what the problem is.
Thanks for any help!!

[digicosmos86]

Hi @JoeSu112,

First, are you using GPU for inference? If you have switched from using CPU to GPU as the backed for JAX, then it might not be possible to use parallel inference. In the error messages that you have provided, it seems to be checking for CUDA versions, which led me to think that you might have installed the GPU version of JAX.

If you are unable to run different chains in parallel, then that could be the reason why estimation takes much longer. It might be helpful to paste all the packages and their versions in your virtual environment so we can start from there.

Thanks!
Paul

[JoeSu112]

Hi, @digicosmos86,
Sorry for the late reply. After trying for days, I find that GPU is not activated at all. I ran the estimation on CPU all the time as I found that the following message had showed up when I ran estimation in the past version of HSSM:

No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.) 

I will contact the maintainer of the machine to see how to activate CUDA on it.
Also, this is the information of CUDA on the machine:

nvcc --version
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2021 NVIDIA Corporation
Built on Thu_Nov_18_09:45:30_PST_2021
Cuda compilation tools, release 11.5, V11.5.119
Build cuda_11.5.r11.5/compiler.30672275_0

nvidia-smi      
+-----------------------------------------------------------------------------------------+
| NVIDIA-SMI 550.54.15              Driver Version: 550.54.15      CUDA Version: 12.4     |
|-----------------------------------------+------------------------+----------------------+
| GPU  Name                 Persistence-M | Bus-Id          Disp.A | Volatile Uncorr. ECC |
| Fan  Temp   Perf          Pwr:Usage/Cap |           Memory-Usage | GPU-Util  Compute M. |
|                                         |                        |               MIG M. |
|=========================================+========================+======================|
|   0  NVIDIA RTX A6000               Off |   00000000:3B:00.0 Off |                  Off |
| 46%   72C    P2            144W /  300W |    2482MiB /  49140MiB |     67%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+
|   1  NVIDIA RTX A6000               Off |   00000000:AF:00.0 Off |                  Off |
| 30%   39C    P8             23W /  300W |      12MiB /  49140MiB |      0%      Default |
|                                         |                        |                  N/A |
+-----------------------------------------+------------------------+----------------------+                                                                                         

But it's still strange that if the GPU is not activated all the time, why the estimation time become much longer in new version? The estimation time is indeed more than a day @AlexanderFengler. The estimation result is almost the same as I did in the past version.

Thanks so much again for your help!!

[digicosmos86]

Do you see this warning WARNING (pytensor.tensor.blas): Using NumPy C-API based implementation for BLAS functions. when PyMC is imported?

[JoeSu112]

Yes, this warning shows up:

>>> import pytensor  # Graph-based tensor library
WARNING (pytensor.tensor.blas): Using NumPy C-API based implementation for BLAS functions.

[digicosmos86]

Yeah this is the culprit. The problem is that newer numpy versions (>=1.26.4) stopped providing PyTensor with information about BLAS libraries, so newer versions of HSSM that depends on newer versions of PyMC (and Pytensor) will not know where to find the BLAS library.

There is an easy fix to this, which is to install PyMC in a conda environment before installing HSSM. We have a tutorial here. So basically you need to create a conda enivornment, install JAX (the GPU version with cuda versions consistent with the cuda installed on the system) and PyMC with conda, and then install HSSM via pip. We are making lots of progress in making HSSM available through conda as well, so hopefully soon hssm can be installed directly with conda install -c conda-forge hssm.

[RaoXie]

Hi all, I recently updated HSSM to the newest version (0.2.1). However, while I run the program that was able to run in the old version (0.2.0b1), some problem occurs:

infer_DDM = DDM.sample(
    sampler="nuts_numpyro",  # type of sampler to choose,
    cores=2,  # how many cores to use
    chains=chains,  # how many chains to run
    draws=draws,  # number of draws from the markov chain
    tune=int(draws*0.2),  # number of burn-in samples
    idata_kwargs=dict(log_likelihood=False)

UserWarning: There are not enough devices to run parallel chains: expected 4 but got 2. 
Chains will be drawn sequentially. If you are running MCMC in CPU, 
consider using `numpyro.set_host_device_count(4)` at the beginning of your program. 
You can double-check how many devices are available in your system using `jax.local_device_count()`.

I'm sure that I could run parallel chains in the past versions. So I'm wondering how I can solve this problem and why this did not happen in the past versions. Thanks for any help!!!

@JoeSu112 Hi, I am also trying to use slurm scheduler to run multiple chains in parallel:

it worked in jupyter notebook, but for slurm it keeps giving me this error

It seems has something to do with resources, I believe 1 node, 1 task, 4 cpu should be good for the slurm job. I also played around different combinations, for example 4 tasks, but still doesn't work. Do you mind sharing how did you run this through slurm? Thank you so much in advance for your help!

Best,
Rao