Pymc3-3.11.0 with GPU support

Questions
1

Hello pmyc3 experts

I’m trying to run pymc3 with GPU support with following setup:

ubuntu-20.04, kernel 5.4.0-72-generic
I created a conda environment and installed:

pip install pymc-3.11.0
conda install -c conda-forge pygpu

I installed cuda and cudnn from the nvidia site. I can run nvidida-smi and it detects my NVIDIA GeForce GTX 1080

I set up a .theanorc file in my home dir:

[global]
floatX = float32
device = cuda0
force_device = True
optimizer_including=cudnn

[dnn]
include_path=/usr/local/cuda-10.0/include
library_path=/usr/local/cuda-10.0/lib64

[cuda]
root=/usr/local/cuda-10.0/

But when I import pymc3 I get the following:

/home/hadron/myutils/anaconda3/envs/myenv1/lib/python3.8/site-packages/theano/gpuarray/dnn.py:192: UserWarning: Your cuDNN version is more recent than Theano. If you encounter problems, try updating Theano or downgrading cuDNN to a version >= v5 and <= v7.
  warnings.warn(
Using cuDNN version 7401 on context None
Mapped name None to device cuda0: NVIDIA GeForce GTX 1080 (0000:01:00.0)
ERROR (theano.gpuarray): Could not initialize pygpu, support disabled
Traceback (most recent call last):
  File "/home/hadron/myutils/anaconda3/envs/myenv1/lib/python3.8/site-packages/theano/gpuarray/__init__.py", line 262, in <module>
    use(config.device)
  File "/home/hadron/myutils/anaconda3/envs/myenv1/lib/python3.8/site-packages/theano/gpuarray/__init__.py", line 251, in use
    optdb.add_tags("gpuarray_opt", "fast_run", "fast_compile")
AttributeError: module 'theano.gpuarray.optdb' has no attribute 'add_tags'

As you can see I used cuda 10.0 and installed cuDNN 7.4.0.1, but I also tried with others, i.e. more recent versions (starting from cuda 11.3 and working my way down), trying to address the message “Your cuDNN version is more recent than Theano” but that still doesn’t work.
Can anyone help me with that? I’m quite stuck with it

2

FWIW I am also having this problem. Having spent the better part of the last two days installing and uninstalling CUDA and friends, I came here, stumped. It makes feel less stupid to see a fellow traveller on the road…

Anyway, the warning seems to be triggered by this test:

v = version()
...
if v >= 7200:
    warnings.warn(
        "Your cuDNN version is more recent than "
        "Theano. If you encounter problems, try "
        "updating Theano or downgrading cuDNN to "
        "a version >= v5 and <= v7."
    )

This makes me think you’d need 7.2.x or less before that warning disappeared.

However, as far as I can tell, it’s related to python code. If you look at gpuarray/__init__.py you’ll see that, indeed (as the error of course correctly states), optdb is a module not a class. The method add_tags is a member of SequenceDB, found in graph/optdb.py. The module graph appears to have been renamed recently from gof?

You can prove it by putting a print(optdb) right before the error occurrs:

$ THEANO_FLAGS="floatX=float32,device=cuda" python -c 'import theano'
Using cuDNN version 7605 on context None
Mapped name None to device cuda: NVIDIA GeForce GTX TITAN X (0000:01:00.0)
<module 'theano.gpuarray.optdb' from '.../site-packages/theano/gpuarray/optdb.py'>

In fact, there exists a file, theano/gpuarray/optdb.py that might be getting imported instead? I thought this might be a result of the order of things getting imported (since this is happening inside a __init__.py file, so I re-imported optdb:

    if default_to_move_computation_to_gpu:
        from theano.compile import optdb
        optdb.add_tags("gpuarray_opt", "fast_run", "fast_compile")
        optdb.add_tags("gpua_scanOp_make_inplace", "fast_run")

which yields the following error:

Traceback (most recent call last):
  File ".../site-packages/theano/gpuarray/__init__.py", line 268, in <module>
    use(config.device)
  File ".../site-packages/theano/gpuarray/__init__.py", line 255, in use
    print(optdb)
  File ".../site-packages/theano/graph/optdb.py", line 427, in __str__
    self.print_summary(sio)
  File ".../site-packages/theano/graph/optdb.py", line 419, in print_summary
    positions.sort(c)

At this point, I’m having trouble understanding how everyone doesn’t have this problem…

Cheers!

edit: I’m on CUDA 10.2, cudnn 7.6.5, ubuntu 16.04, pymc 3.11.2, pygpu 0.7.6, theano-pymc 1.1.2.

1 Like
3

I would try sampling on the GPU using the new JAX backend: Inefficient use of pmap in jax numpyro sampler? · Issue #4288 · pymc-devs/pymc3 · GitHub for some info.

4

try sampling on the GPU using the new JAX backend

Can you elaborate a bit on how to do that? Are there instructions somewhere?

Following @junpenglao’s gist, I tried:

$ python -c 'import pymc3.sampling_jax'
Traceback (most recent call last):
  File "<string>", line 1, in <module>
  File "/home/jrporter/miniconda2/envs/pymc/lib/python3.8/site-packages/pymc3/sampling_jax.py", line 18, in <module>
    from theano.link.jax.jax_dispatch import jax_funcify
  File "/home/jrporter/miniconda2/envs/pymc/lib/python3.8/site-packages/theano/link/jax/jax_dispatch.py", line 86, in <module>
    jax.config.disable_omnistaging()
  File "/home/jrporter/miniconda2/envs/pymc/lib/python3.8/site-packages/jax/config.py", line 167, in disable_omnistaging
    raise Exception(
Exception: Disabling of omnistaging is no longer supported in JAX version 0.2.12 and higher: see https://github.com/google/jax/blob/master/design_notes/omnistaging.md.

Thanks!

5

Yes, you need to downgrade JAX.

6

Any solutions for this for advi fitting?

7

Is that not working with the JAX backend?

8

Where can I find how to do this? There is only documentation for sampling with jax.

9

Try setting

# Disable C compilation by default
theano.config.cxx = ""
# This will make the JAX Linker the default
theano.config.mode = "JAX"

And then just running ADVI as usual.

10

Setting config.cxx to “” raises the error RuntimeError: The new gpu-backend need a c++ compiler. This check happens here

Keeping it at default but setting mode to “JAX” gives me the same error as OP: AttributeError: module 'theano.gpuarray.optdb' has no attribute 'add_tags'

11

Which version of theano and pymc3 are you running?

12

PyMC3 3.11.2, theano (Theano-PyMC) 1.1.2

13

Very friendly bump

14

You could try installing the (experimental) main branch and aesara in a fresh env.

15

Hello!
I’m also struggling with the same error message as @justinrporter described earlier.

/home/leo/anaconda3/lib/python3.8/site-packages/theano/gpuarray/dnn.py:192: UserWarning: Your cuDNN version is more recent than Theano. If you encounter problems, try updating Theano or downgrading cuDNN to a version >= v5 and <= v7.
  warnings.warn(
Using cuDNN version 8204 on context None
Mapped name None to device cuda0: NVIDIA GeForce GTX 1080 Ti (0000:4C:00.0)
ERROR (theano.gpuarray): Could not initialize pygpu, support disabled
Traceback (most recent call last):
  File "/home/leo/anaconda3/lib/python3.8/site-packages/theano/gpuarray/__init__.py", line 262, in <module>
    use(config.device)
  File "/home/leo/anaconda3/lib/python3.8/site-packages/theano/gpuarray/__init__.py", line 251, in use
    optdb.add_tags("gpuarray_opt", "fast_run", "fast_compile")
AttributeError: module 'theano.gpuarray.optdb' has no attribute 'add_tags'

Was this issue ever resolved? For reference, these are the versions I’m using:

  • CUDA 11.4
  • cuDNN 8.2.4
  • theano-pymc 1.1.0
  • pymc3 3.11.0
  • pygpu 0.7.6
16

Also, experiencing these issues on Colab. It shows different errors with different library configurations in my case, but just won’t utilize the GPU.

Colab is running pymc3 3.11.4 by default. There is a myriad of recipes both in this thread and elsewhere on how to make it work, including installing dependencies via conda (e.g., !conda install -c conda-forge pymc3 theano-pymc mkl mkl-service, !conda install -c conda-forge pygpu, uninstalling theano / installing aesara, installing dev branch pymc and aesara, forcing theano.config.mode = 'JAX'. And I tried all of it in various combinations, but unfortunately nothing seems to work.

That said, I do get different errors, for example I seem to get a missing BLAS warning in default colab configuration, which is replaced by the Your cuDNN version is more recent than Theano warning after installing theano-pymc, etc via conda. Forcing theano.config.mode = 'JAX' raises the no longer allowed to disable the omnistaging warning, which after downgrading jax seems to go away, but when it tries to run the model it just shows a bunch of zero division errors. Sadly, at no point it seems to be trying to utilize the GPU.

Anyway, really appreciate the developers’ work on bringing us this wonderful tool, but I guess it will be a while until the new functionality, like the GPU speed-up, is accessible, particularly on less flexible platforms, like Colab.

17

I am also struggle exactly the same error. Would be really appreciated if there is any latest instructions/solutions that make it works. Cheers!

18

We will not support GPUs via this path in the future. Instead, you could try pymc 4.0.0b2 which has JAX support which has GPU support. We have used it quite effectively, e.g.: MCMC for big datasets: faster sampling with JAX and the GPU - PyMC Labs

2 Likes