[Issue]: getting self-consistent set of wheels

[ksebaz]

Dear developers,

I'm struggling getting a self-consistent set of wheels for v0.7.1 with ROCm 7.0 that, when put together into a wheelhouse can be installed alongside. Currently I get jax and jaxlib from ROCm/jax (branch rocm-jaxlib-v0.7.1) and plugin and pjrt from here (branch rocm-jax-v0.7.1)since this is how to do it at the moment as far as I understood. However there are several issues:

• to strip 'dev0-selfbuilt' from the resulting jaxlib-wheel to make it findable as a proper dependency of jax, one has to add --bazel_options="--action_env=ML_WHEEL_TYPE=release" to the build since the environment variable JAX_RELEASE or JAXLIB_RELEASE are currently ignored for jaxlib apparently.
• despite specifying '--rocm_version' correctly to '7' which is necessary since otherwise the plugin cannot be found at runtime, during installation however the dependency resolution from jax is still looking for plugins labeled with 'rocm60'. This is when specifying jax[rocm] for pip. Specific extras rocm60 or rocm7 are not available it seems.

For experimenting around personally this is all not an issue but in order to provide jax 0.7.1 with rocm on our systems in an automated fashion a standard installation workflow would be highly desirable.

ROCm Version

ROCM 7.0.1

[lucbruni-amd]

Hi @ksebaz, thanks for opening the issue. I understand it is a bit of a struggle at the moment especially if you're working with automated workflows. I agree that something needs to be done about that, because our current installation are neither clear nor workflow-friendly.

Have you tried installing jaxlib, pjrt and plugin from the 0.7.1 release instead of the wheelhouse method? That is, installing those three here, then installing jax from PyPI with pip install jax==0.7.1? If your workflow requires building the wheels, let's continue discussing from there.

[ksebaz]

Hi @lucbruni-amd, thanks a lot for your help!

Have you tried installing jaxlib, pjrt and plugin from the 0.7.1 release instead of the wheelhouse method? That is, installing those three here, then installing jax from PyPI with pip install jax==0.7.1?

yes, and that works (in principle; please see below) as works installing the wheels I build myself. This issue is really only about the dependency management, such that when putting all wheels (including jax) into a directory, say called pip-index, something like pip install -f pip-index jax[rocm]==0.7.1 works. One issue is that jax is still searching for jax_rocm60_plugin.

If your workflow requires building the wheels, let's continue discussing from there.

My need for bulding from scratch only results from the rapid development in ROCm. E.g the available wheels you mentioned are apparently build for ROCm 7.1.x since they require libhipfftw.so which didn't exist in rocm/7.0.x if I'm not mistaking. But when I wanted to provide jax 0.7.1 a few weeks ago we didn't have ROCm 7.1 yet. But in the meantime we have that version too on our systems, so changing to ROCm 7.1 the released wheels seem to work fine.

[lucbruni-amd]

My pleasure, happy to help.

Your reasoning for building the plugins from scratch is totally valid.

This issue is really only about the dependency management, such that when putting all wheels (including jax) into a directory, say called pip-index, something like pip install -f pip-index jax[rocm]==0.7.1 works. One issue is that jax is still searching for jax_rocm60_plugin.

Could you provide me your set of commands from the start? Then I can try to reproduce it on my end to see what's the problem.

[ksebaz]

yes, sure. This script reproduces what I mean:

#!/bin/bash

set -e

mkdir pip-index

[[ $(python --version | cut -d"." -f2) == "13" ]] || ( echo "use py13" && exit 1 )

wget -P pip-index https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.7.1/jaxlib-0.7.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl
wget -P pip-index https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.7.1/jax_rocm7_plugin-0.7.1-cp313-cp313-manylinux_2_28_x86_64.whl
wget -P pip-index https://github.com/ROCm/rocm-jax/releases/download/rocm-jax-v0.7.1/jax_rocm7_pjrt-0.7.1-py3-none-manylinux_2_28_x86_64.whl

python -m venv --system-site-packages foo
foo/bin/pip install -f pip-index "jax[rocm]==0.7.1"

The error is:

ERROR: Could not find a version that satisfies the requirement jax-rocm60-plugin<=0.7.1,>=0.7.1; extra == "rocm" (from jax[rocm]) (from versions: none)

In the above this comes from the jax 0.7.1 from pypi, but the same happens for a jax-wheel build from ROCm/jax (branch rocm-jaxlib-v0.7.1) where it could in principle be already correct. I'm aware I can circumvent this by installing the plugin/pjrt/jaxlib separately and then simply jax without the extra as you mentioned. But towards rocm-jax being published on pypi at some point the above would be a necessary condition and generally convenient.

[lucbruni-amd]

I think one of the problems is here: https://github.com/jax-ml/jax/blob/b4791938dbdee49f1f144bc99af226355cd843ae/setup.py#L116-L120. I'll need to discuss with the team the best way to deal with these versioning discrepancies, as our new wheels are named like: jax-rocm7-pjrt
and jax-rocm7-plugin. Let me get back to you on this.

Also we don't use the ROCm/jax repository anymore as I believe we are moving forward with the upstream, but I'll get confirmation on that as well.

[ksebaz]

I just realize from #200, that the wheels are already on pypi. So then the reproducer is of course simply: pip install jax[rocm]==0.7.1`.

[PhilipVinc]

yeah, but rocm is not on pypi...

[lucbruni-amd]

The rocm extra specified only supplies plugin and pjrt to my knowledge (see setup.py).

One issue is that jax is still searching for jax_rocm60_plugin.

Also, this PR was approved and merged recently which I believe should resolve this issue for jax versions > 0.8.1 once released. I'll update this issue once this can be tested.

[ksebaz]

I just successfully tried with the recently released wheels for v0.8.2. So the second point from the initial issue description is resolved imo.