search

graphcore-research / jax-experimental

JAX for Graphcore IPU (experimental)
https://github.com/graphcore-research/jax-experimental#readme
Apache License 2.0
21 stars 2 forks source link

BUG: ipu_xla_client_pybind.so: undefined symbol #12

Closed PierrickPochelu closed 1 year ago

PierrickPochelu commented 1 year ago

Description

I did " pip install jax==0.3.16+ipu jaxlib==0.3.15+ipu.sdk310 -f https://graphcore-research.github.io/jax-experimental/wheels.html " as mentioned in the documentation.

When I open a Python session and try to "import jax" I get the following error message:

[...]
    from . import ipu_xla_client_pybind as _ipu_xla
ImportError: /home/ipuuser/manu_tests/.venv_pytorch/lib/python3.8/site-packages/jaxlib/ipu_xla_client_pybind.so: undefined symbol: _ZNK6poplar8IPUModel12createDeviceENS_11OptionFlagsEbj, version V0

What jax/jaxlib version are you using?

JAX 0.3.16+ipu , JAXLIB .3.15+ipu.sdk310

Which accelerator(s) are you using?

IPU

Additional System Info

Python 3.8.10, Ubuntu 20.04.5 LTS

balancap commented 1 year ago

Sorry to hear it did not work out of the box!

Do you know which version of the Poplar SDK you are using? jaxlib is compiled against Poplar C++ libraries, so will need a local install of the later. In my experience, the easiest way is run inside one of the Graphcore docker image (e.g. https://hub.docker.com/layers/graphcore/pytorch/3.2.0-ubuntu-20.04/images/sha256-30f0e08f260b76f2821c66cbd7dbe78d4a2a85e040d10185a03a0dbd1499ff2e?context=explore)

PierrickPochelu commented 1 year ago

We have 3.0.0 poplar SDK on our Graphcore server. It is enough for our scientific works on PyTorch/Tensorflow.

If you guide us for having Cloud access to some IPUs we could try your last version of IPUs & poplar. It may be very useful for our PhD students, and our research, where we plan to use more and more Jax.

What do you think ?

Best regards, Pierrick Pochelu

balancap commented 1 year ago

I could reproduce the issue with SDK 3.0. I'll give it a try at compiling jaxlib with it, but I can't guarantee a successful outcome. In the meantime, if you want to try IPUs with latest SDK version, the easiest way for you (and students) is to use IPU on Paperspace. Once logged, you'll be able to run notebooks with IPUs. For instance, here is the quickstart JAX on IPU notebook: https://console.paperspace.com/github/graphcore-research/jax-experimental?container=graphcore%2Fpytorch-jupyter%3A3.1.0-ubuntu-20.04&machine=Free-IPU-POD4&file=%2Fipu%2Fexamples%2Fjax_ipu_quickstart.ipynb

Hopefully this one should run out of the box!

PierrickPochelu commented 1 year ago

I didn t succeeded to log in. The blue "sign in" refresh the page. I tried on Ubuntu & Windows: Chrome, Firefox, Edge.

balancap commented 1 year ago

Sorry to hear, I have forwarded the issue to the internal Paperspace team. It may help to first create & login on the main page (https://www.paperspace.com/), and then reload the notebook?

In the meantime, I have also manage to compile jaxlib for Poplar SDK 3.0. The following install should now work:

pip install jax==0.3.16+ipu jaxlib==0.3.15+ipu.sdk300 -f https://graphcore-research.github.io/jax-experimental/wheels.html

All basics should be working properly. Please tell us if there is any feature we can improve!

PierrickPochelu commented 1 year ago

I saw you have just compiled and published the 3.0 wheel, congratulations! Jax is working fine.

Regarding my cloud account, I have an issue with Paperspace account and not related to you. The graphcore-jax notebook on the cloud is working fine with my colleague account.

Thank you!

balancap commented 1 year ago

Thanks :+1: Please feel free to post any issue you discover, I'll try my best to help :)