Provide AArch64 (ARM) Linux jaxlib wheels #7097
Assignees
Labels
enhancement
New feature or request
Comments
|
The issue is that we don't provide Linux AArch64 wheels for jaxlib. Cavium is an ARM machine. I suspect you may be able to build jaxlib from source on those machines: https://jax.readthedocs.io/en/latest/developer.html#building-from-source Try it out and see how it goes? (I'm not sure we have the resources to release ARM wheels ourselves right now, but I wouldn't rule it out for the future.) |
hawkinsp
added
contributions welcome
hawkinsp
changed the title
|
@bensworth Unfortunately we don't support Intel (Iris) GPUs either. Currently your hardware options are CPU, NVidia GPU, AMD GPU, or Google TPU. So perhaps your best bet is to use CPU on that machine or build jaxlib with CUDA support on ARM. |
|
Yep, realized that a few minutes after I posted, now I am on a Nvidia GPU and have things working! Previously I mostly worked with CPUs/MPI, didn't realize GPUs were so architecture specific. This issue can be closed unless you want to leave it open for somebody who may be able to build wheels for ARM or Intel; that's outside of my area of work though. One thought - it may be worth adding a comment in the readme as to what GPUs are supported by the build instructions. |
|
I'm attempting to build jaxlib wheels for aarch64 Linux via Docker but without success. My Dockerfile is the following My machine is a MacBook Pro with M1 apple silicon chip, macOS 12.2. Building with In #5501 they say how to build jaxlib on the M1 chip , but not through Docker. Would it be possible to provide instructions for how this could be modified to work on Docker? |
|
☝️ this indicates you're running out of memory, especially since the 0.3.x release, you need pretty much the entire 16GB available. I managed to build manylinux aarch64 wheels on M1, you can get them here: https://github.com/yoziru/jax/releases/tag/jaxlib-v0.3.5 Here is the Dockerfile for reproducing these wheels (run from the root of this repo after cloning): FROM quay.io/pypa/manylinux2014_aarch64
COPY .bazelversion .
RUN BAZEL_VERSION=$(cat .bazelversion) \
&& curl -SL https://github.com/bazelbuild/bazel/releases/download/$BAZEL_VERSION/bazel-$BAZEL_VERSION-linux-arm64 \
-o /usr/local/bin/bazel \
&& chmod +x /usr/local/bin/bazel
ARG PYTHON_VERSION_MAIN=38 \
JAXLIB_VERSION=0.3.5
RUN ln -s /opt/python/cp${PYTHON_VERSION_MAIN}-cp${PYTHON_VERSION_MAIN}/bin/pip /usr/local/bin/pip \
&& ln -s /opt/python/cp${PYTHON_VERSION_MAIN}-cp${PYTHON_VERSION_MAIN}/bin/python /usr/local/bin/python \
&& ln -s /opt/python/cp${PYTHON_VERSION_MAIN}-cp${PYTHON_VERSION_MAIN}/bin/python /usr/local/bin/python3
RUN pip install numpy six
WORKDIR /builder
COPY . .
RUN python build/build.py
RUN auditwheel repair dist/jaxlib-${JAXLIB_VERSION}-cp${PYTHON_VERSION_MAIN}-none-manylinux2014_aarch64.whl |
|
+1 for official support for Linux ARM64 from me too! |
|
@sudhakarsingh27 for viz |
|
I'm curious whether using We don't have infrastructure to build and test ARM wheels at the moment, but |
|
Hi @hawkinsp ! |
hawkinsp
added a commit
to hawkinsp/jax
that referenced
this issue
Aug 16, 2022
* Add aarch64 as a known target_cpu value. * Only pass --bazel_options to build actions since they can make "bazel shutdown" fail. * Pass the bazel startup options to "bazel shutdown". Issue jax-ml#7097 Fixes jax-ml#7639
Merged
|
I verified we can successfully build a Linux aarch64 jaxlib on a Google Cloud t2a virtual machine. I was also able to cross-compile a Linux aarch64 jaxlib from a Debian Linux x86-64 machine by doing the following:
and I verified this wheel works at least for some simple tests on a t2a VM. |
1 task
|
Hi! What is stopping the progress here ? |
treyra
added a commit
to treyra/jax
that referenced
this issue
Oct 27, 2022
Currently non x86_64 linux architectures are not supported, see jax-ml#7097 for request to change this. This can lead to installation confusion, as jax will install, but jaxlib will not. For example see jax-ml#12307. If this can be more clearly phrased or explained, let me know.
treyra
added a commit
to treyra/jax
that referenced
this issue
Oct 27, 2022
Currently non x86_64 linux architectures are not supported, see jax-ml#7097 for request to change this. This can lead to installation confusion, as jax will install, but jaxlib will not. For example see jax-ml#12307. If this can be more clearly phrased or explained, let me know.
treyra
added a commit
to treyra/jax
that referenced
this issue
Nov 4, 2022
Currently non x86_64 linux architectures are not supported, see jax-ml#7097 for request to change this. This can lead to installation confusion, as jax will install, but jaxlib will not. For example see jax-ml#12307. This adds a note to the install sections for the relevant pip wheels.
|
I am glad to see there is some progress already on this task! |
nikitosl
pushed a commit
to nikitosl/jokes-generator-web
that referenced
this issue
Jan 8, 2023
# Current blocker is that jax not supported linux library for ARM64 # jax-ml/jax#7097
3 tasks
hawkinsp
removed
the
contributions welcome
|
If you experience this while building a docker image like I did, then the answer here solves the problem easily for me and many others (as seen in the number of upvotes). |
|
@ridwan-salau The solution you linked to is a workaround to use the AMD64 image on Mac ARM64 via emulation (Rosetta). |
|
I've come back to this issue several times in the last 6 months. I'd love to use jax in our M1 Mac docker images without resorting to Rosetta. Is there anything I can do to help push this forward? I can have a go at modifying the current build jobs to support aarch64, but only if such a PR is welcome. |
|
I had a go at building from source, and it was remarkably easy - following the "build from source" instructions using |
joverlee521
added a commit
to nextstrain/docker-base
that referenced
this issue
Jul 6, 2023
The latest release of evofr included updates to its dependencies that updated the minimum version for jaxlib.¹ I'm using the latest release of jaxlib available in the user forked repo since there's still no official pre-built binaries for linux/arm64.² I did not realize this change was needed until I saw the "Validate Platform" job fail in the CI run³ that I triggered to update evofr to v0.1.20. ¹ blab/evofr@46f744b ² jax-ml/jax#7097 ³ https://github.com/nextstrain/docker-base/actions/runs/5480222698
1 task
joverlee521
added a commit
to nextstrain/docker-base
that referenced
this issue
Jul 7, 2023
The latest release of evofr included updates to its dependencies that updated the minimum version for jaxlib.¹ I'm using the latest release of jaxlib available in the user forked repo since there's still no official pre-built binaries for linux/arm64.² I did not realize this change was needed until I saw the "Validate Platform" job fail in the CI run³ that I triggered to update evofr to v0.1.20. Added a note for our future selves to check the jaxlib installation if we run into different versions of evofr on the different platforms. ¹ blab/evofr@46f744b ² jax-ml/jax#7097 ³ https://github.com/nextstrain/docker-base/actions/runs/5480222698
|
@yoziru the v0.4.13 manylinux aarch64 wheels aren't working properly. Can you enable issues on your fork so I can send a detailed report? |
victorlin
added a commit
to victorlin/build-jax
that referenced
this issue
Jul 18, 2023
from github . com/jax-ml/issues/7097#issuecomment-1110730040
This was referenced Jul 24, 2023
|
We've started publishing aarch64 Linux nightly wheels (https://storage.googleapis.com/jax-releases/jaxlib_nightly_releases.html) and we will include aarch64 Linux in the next jaxlib release. |
|
Jaxlib 0.4.18 includes Linux aarch64 wheels (thanks @yashk2810 !) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Potentially related to #6932 or #7052.
I am using a heterogeneous cluster with various different compute options. I got Jax and Flax installed fine on the CPU system. Then I tried using a GPU node (specifically dual socket Cavium ThunderX2 99xx 32-core processors with two NVIDIA Tesla V100 GPUs), but am unable to install jaxlib:
or even
Below is my pip version, which is up to date. Pip also installs other libraries fine (numpy, scipy, etc.)
The text was updated successfully, but these errors were encountered: