Numpyro: GPU version install

[ross-h1]

Hi, Firstly thanks for an excellent piece of work, am seeing huge speedups in inference v. other frameworks, making a big impact on my research!

Please excuse potentially naiive question, but I'm struggling to see how to install / run a GPU capable Numpyro. I have JAX working on GPU... some pointers please?!

Thank you

Ross

[fehiepsi]

Hi Ross, I think that you want to call numpyro.set_platform('gpu') at the top of your program.

[ross-h1]

Thank you, have tried that, but i seem to have an install issue. I can set up a working JAX environment (CPU/GPU) as per the JAX install guide.

If i then install Numpyro running either:
pip install git+https://github.com/pyro-ppl/numpyro.git
or
pip install numpyro

This breaks this install, for both numpyro with the command you mention, or with any JAX code set up to use GPU, and I get the following error:

File "/home/..../anaconda3/envs/PYRO/lib/python3.7/site-packages/jaxlib/xla_client.py", line 248, in get_local_backend
raise RuntimeError('Unknown backend {}'.format(name))

the pip installer seems to uninstall jax/jaxlib and replace them with older versions...

[fehiepsi]

@ross-h1 Currently, NumPyro only works with jax==0.1.53 and jaxlib==0.1.36. If you want to use the latest jax/jaxlib version then you might want to install the PR https://github.com/pyro-ppl/numpyro/pull/496. There is an issue in the last version of jax that will break some NumPyro models so we can’t merge that PR yet.

FYI, you can use pip with no-deps flag to not update dependencies https://pip.pypa.io/en/stable/reference/pip_install/#install-no-deps. But I would recommend to install jaxlib==0.1.36 instead.

[neerajprad]

@ross-h1 - I am closing this issue, but please feel free to reopen if you continue to face issues. The simplest way would be to follow @fehiepsi's suggestion and install jaxlib==0.1.36 using installation instructions for the GPU.

[benwaldner]

I had the same issue with Jax and Jaxlib. Didnt found any way how to solve it. Anyone solved it? I dont know what to try else. All things suggested come up with the error: Backend GPU not found. And I tried it on several machines/ linux-installations. Even on CPU the problem is that it doesnt use all the cores on 100% but skip after one CPU-core hits 100% and doesnt load balancing it efficiently on the other cores. Making it enormously slow.

Anyone found something that worked to solve this error?
I really dont know what to try else.

[fehiepsi]

@benwaldner Did you follow the steps in jax for GPU support? Make sure that you specify the correct CUDA_VERSION in your system.