Plugins not being loaded properly in gpu_triton in "community" docker images

[owainkenwayucl]

Description

I've been trying to get AlphaFold3 working on an AMD Mi300x box we have on loan from HPE and I've run into some issues caused by the test here:

jax/jaxlib/gpu_triton.py

Line 20 in 1f93b4b

_hip_triton = import_from_plugin("rocm", "_triton")

failing on the "community" docker images for JAX, which results in triton complaining that we don't have the GPU version of JAX installed.

I've tried:

docker.io/rocm/jax-community     latest                         193ba487b999
docker.io/rocm/jax-community     rocm6.2.4-jax0.4.35-py3.11.10  ef50d5181ba5  
docker.io/rocm/jax-community     rocm6.2.3-jax0.4.34-py3.11.10  b229479e4af8  

As an example, you can test this in a container built from these images:

Python 3.10.16 (main, Feb 17 2025, 01:40:07) [GCC 11.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> from jaxlib import gpu_triton
>>> gpu_triton._hip_triton
>>>

For comparison with the "non community" but older JAX:

docker.io/rocm/jax               latest                         d949265c6ac2

Python 3.10.12 (main, Feb  4 2025, 14:57:36) [GCC 11.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> from jaxlib import gpu_triton
>>> gpu_triton._hip_triton
<module 'jax_rocm60_plugin._triton' from '/opt/venv/lib/python3.10/site-packages/jax_rocm60_plugin/_triton.so'>
>>> 

I'm also finding it difficult to install from wheels as per https://github.com/rocm/jax/tree/main/build/rocm because pip can't find the wheels for the ROCm features.

e.g.

root@ea47cc7f8ef9:/# pip install jax[rocm]==0.4.38
Collecting jax[rocm]==0.4.38
  Downloading jax-0.4.38-py3-none-any.whl (2.2 MB)
     ━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━ 2.2/2.2 MB 52.2 MB/s eta 0:00:00
WARNING: jax 0.4.38 does not provide the extra 'rocm'

I remember distantly from installing things on Nvidia GPUs, that there is a "secret" repository you have to use sometimes to get JAX Cuda plugins to work (https://storage.googleapis.com/jax-releases/jax_cuda_releases.html) - is there a similar thing for ROCm plugins?

System info (python version, jaxlib version, accelerator, etc.)

docker.io/rocm/jax-community containers, tag rocm6.2.3-jax0.4.34-py3.11.10, rocm6.2.4-jax0.4.35-py3.11.10, latest

Python 3.10.16 (main, Feb 17 2025, 01:40:07) [GCC 11.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import jax; jax.print_environment_info()
jax:    0.5.0
jaxlib: 0.5.0
numpy:  1.26.4
python: 3.10.16 (main, Feb 17 2025, 01:40:07) [GCC 11.4.0]
device info: AMD Instinct MI300X-8, 8 local devices"
process_count: 1
platform: uname_result(system='Linux', node='5ca490f32332', release='5.14.0-503.34.1.el9_5.x86_64', version='#1 SMP PREEMPT_DYNAMIC Thu Mar 27 06:00:50 EDT 2025', machine='x86_64')

[ppanchad-amd]

Hi @owainkenwayucl. Internal ticket has been created investigate this issue. Thanks!

[owainkenwayucl]

Thanks! - I misused my knowledge of the dark arts of python packaging to get AlphaFold 3 to install and run with the older version of JAX but it'd be good to be able to use the correct one!

[mrodden]

@owainkenwayucl

For comparison with the "non community" but older JAX:

The rocm/jax-community images are (were, we have some plans to change this) built to match up with what JAX upstream release tags have as their content, but the rocm/jax images have additional fixes and testing put into them before release but they lag behind on the JAX version.

Unfortunately some of the issues you found with the triton portion of the plugin in 0.4.35 didn't seem to be working around the 0.4.35 JAX release timeframe with ROCm and it doesn't work very reliably in the jax-community image. I would suggest trying to use the rocm/jax images where possible, but as you found out they don't have the latest JAX versions.

I'm also finding it difficult to install from wheels as per https://github.com/rocm/jax/tree/main/build/rocm because pip can't find the wheels for the ROCm features.

The wheels are still pretty new and we have had a hard time with getting a set of plugin wheel versions to work with the upstream JAX wheel releases. I find that its best to not use the extra flag and just install the plugin wheels by hand with the version of jax+jaxlib you want. For example:

pip install jax==0.4.37 jaxlib==0.4.37 jax_rocm60_plugin==0.4.35 jax_rocm60_pjrt==0.4.35

This is due to the weird version constraint on the jax_rocm60_plugin inside of the jaxlib install_requires section, which really needs to be fixed but as we haven't had many versions of the rocm plugin published we haven't fixed it yet. Also this we are going to be moving to using a separate versioning for the plugin in the future rather than trying to get them to match up like in the 0.4.x and 0.5.x timeframe.

There are currently two plugin wheel sets published to pypi.org for ROCm: 0.4.35 and 0.5.0; these should work with JAX 0.4.35 (and some higher 0.4.x like 0.4.36 and such) and the 0.5.0 should work with 0.5.x of JAX.

I did notice some issues with the linking in the 0.5.0 of the jax_rocm60_plugin but for the most part they have been working for numpy / scipy APIs in JAX.

Let me know if that helps, we are still trying to get some of the issues you are hitting sorted with a new plugin project that we can test and update in a way that is decoupled from the direct JAX versions.

[owainkenwayucl]

Thanks - let me take a look and see if that helps.

[owainkenwayucl]

OK - that's solve my "pip install problem" but of course as you note later plugins don't work.

I did some tests with "newer JAX" (0.4.35) and "working" plugin (0.4.31) from the rocm/jax container but unfortunately while that builds, and can find the JAX ROCM devices , it crashes generating pages and pages of stack trace about tensors when I run AlphaFold3 which I'm still digging into.

It does run with everything being 0.4.31 so it's not a super urgent problem.

[owainkenwayucl]

Just a note to confirm that this is also failing on the new rocm/jax:rocm6.4-jax0.4.35-py3.12 offical containers.

[lucbruni-amd]

Hi @owainkenwayucl,

I did some tests with "newer JAX" (0.4.35) and "working" plugin (0.4.31) from the rocm/jax container but unfortunately while that builds, and can find the JAX ROCM devices , it crashes generating pages and pages of stack trace about tensors when I run AlphaFold3 which I'm still digging into.

Could you provide any of these logs? It would be great if we could help debug and work you up to newer versions as we make progress on the project @mrodden mentioned. It is good to hear that you are able to workaround both this issue and the original one with the 0.4.31 pairing.

Thanks!

[owainkenwayucl]

Could you provide any of these logs? It would be great if we could help debug and work you up to newer versions as we make progress on the project @mrodden mentioned. It is good to hear that you are able to workaround both this issue and the original one with the 0.4.31 pairing.

Sure - they are here: https://github.com/owainkenwayucl/protein_shake/blob/main/AlphaFold3/Crash_Logs/AMD-JAX-0.4.35-plugin-0.4.31-crash.txt

[owainkenwayucl]

This afternoon I had some time and did some digging around and the libraries in the pypi wheel for 0.4.35 look like this:

root@6f3bd14310e6:/alphafold3_venv/lib/python3.12/site-packages/jax_rocm60_plugin# ldd _triton.so 
        linux-vdso.so.1 (0x00007ffd2ebea000)
        libdl.so.2 => /lib/x86_64-linux-gnu/libdl.so.2 (0x00007f2b309b9000)
        libm.so.6 => /lib/x86_64-linux-gnu/libm.so.6 (0x00007f2b308d0000)
        libpthread.so.0 => /lib/x86_64-linux-gnu/libpthread.so.0 (0x00007f2b308cb000)
        libstdc++.so.6 => /lib/x86_64-linux-gnu/libstdc++.so.6 (0x00007f2b3064d000)
        libgcc_s.so.1 => /lib/x86_64-linux-gnu/libgcc_s.so.1 (0x00007f2b3061f000)
        libc.so.6 => /lib/x86_64-linux-gnu/libc.so.6 (0x00007f2b3040b000)
        /lib64/ld-linux-x86-64.so.2 (0x00007f2b30cbd000)

In comparison to the corresponding library for 0.4.31 from the official container:

root@f6b12aaaa519:/alphafold3_venv/lib/python3.12/site-packages/jax_rocm60_plugin# ldd _triton.so 
        linux-vdso.so.1 (0x00007ffc059b5000)
        libstdc++.so.6 => /lib/x86_64-linux-gnu/libstdc++.so.6 (0x00007f11d3f81000)
        libm.so.6 => /lib/x86_64-linux-gnu/libm.so.6 (0x00007f11d3e98000)
        libamdhip64.so.6 => /opt/rocm-6.3.2/lib/libamdhip64.so.6 (0x00007f11d282c000)
        libdl.so.2 => /lib/x86_64-linux-gnu/libdl.so.2 (0x00007f11d2827000)
        libpthread.so.0 => /lib/x86_64-linux-gnu/libpthread.so.0 (0x00007f11d2820000)
        libhsa-runtime64.so.1 => /opt/rocm-6.3.2/lib/libhsa-runtime64.so.1 (0x00007f11d24e5000)
        librt.so.1 => /lib/x86_64-linux-gnu/librt.so.1 (0x00007f11d24e0000)
        libgcc_s.so.1 => /lib/x86_64-linux-gnu/libgcc_s.so.1 (0x00007f11d24b2000)
        libc.so.6 => /lib/x86_64-linux-gnu/libc.so.6 (0x00007f11d22a0000)
        /lib64/ld-linux-x86-64.so.2 (0x00007f11d45a3000)
        librocprofiler-register.so.0 => /opt/rocm-6.3.2/lib/librocprofiler-register.so.0 (0x00007f11d2219000)
        libamd_comgr.so.2 => /opt/rocm-6.3.2/lib/libamd_comgr.so.2 (0x00007f11c974c000)
        libnuma.so.1 => /lib/x86_64-linux-gnu/libnuma.so.1 (0x00007f11c973e000)
        libelf.so.1 => /lib/x86_64-linux-gnu/libelf.so.1 (0x00007f11c971f000)
        libdrm.so.2 => /opt/amdgpu/lib/x86_64-linux-gnu/libdrm.so.2 (0x00007f11c9706000)
        libdrm_amdgpu.so.1 => /opt/amdgpu/lib/x86_64-linux-gnu/libdrm_amdgpu.so.1 (0x00007f11c96f5000)
        libz.so.1 => /lib/x86_64-linux-gnu/libz.so.1 (0x00007f11c96d9000)
        libzstd.so.1 => /lib/x86_64-linux-gnu/libzstd.so.1 (0x00007f11c961f000)

So there would appear to be some linking issues with 0.4.35.

It's similar with 0.5.0.

[cboillot]

hello @owainkenwayucl,

I only tested on jax6 but you should be able to fix the issue by passing
--bazel_options=--linkopt=-L$ROCM_PATH/lib --bazel_options=--linkopt=-lamdhip64
to the builder of jax.

git clone --branch rocm-jax-v0.4.35 --single-branch https://github.com/ROCm/jax
cd jax
cp build/rocm/dev_build_rocm.py dev_build_rocm.py
ROCM_PATH=${ROCM_PATH:-/opt/rocm}
python3 ./build/build.py build --use_clang --clang_path $ROCM_PATH/llvm/bin/clang \
  --bazel_options=--linkopt=-L$ROCM_PATH/lib \
  --bazel_options=--linkopt=-lamdhip64 \
   --rocm_amdgpu_targets "${PYTORCH_ROCM_ARCH//;/,}" --wheels=jaxlib,jax-rocm-plugin,jax-rocm-pjrt \
   --rocm_version=60 --rocm_path=$ROCM_PATH

ldd jax_rocm60_plugin/_triton.so
        linux-vdso.so.1 (0x00007ffd2d31f000)
        libdl.so.2 => /lib64/libdl.so.2 (0x00007fcd57fbb000)
        libpthread.so.0 => /lib64/libpthread.so.0 (0x00007fcd57d9b000)
        libm.so.6 => /lib64/libm.so.6 (0x00007fcd57a19000)
        libamdhip64.so.6 => /opt/software/rocm/6.4.0//lib/libamdhip64.so.6 (0x00007fcd560d0000)
        libstdc++.so.6 => /opt/software/gcc/13.2.0/lib64/libstdc++.so.6 (0x00007fcd55e6c000)
        libgcc_s.so.1 => /opt/software/gcc/13.2.0/lib64/libgcc_s.so.1 (0x00007fcd586e4000)
        libc.so.6 => /lib64/libc.so.6 (0x00007fcd55a96000)
        /lib64/ld-linux-x86-64.so.2 (0x00007fcd584f9000)
        librocprofiler-register.so.0 => /opt/software/rocm/6.4.0//lib/librocprofiler-register.so.0 (0x00007fcd557d9000)
        librt.so.1 => /lib64/librt.so.1 (0x00007fcd555d1000)
        libamd_comgr.so.3 => /opt/software/rocm/6.4.0//lib/libamd_comgr.so.3 (0x00007fcd4c1fd000)
        libhsa-runtime64.so.1 => /opt/software/rocm/6.4.0//lib/libhsa-runtime64.so.1 (0x00007fcd4bcc6000)
        libnuma.so.1 => /lib64/libnuma.so.1 (0x00007fcd4baba000)
        libz.so.1 => /lib64/libz.so.1 (0x00007fcd4b8a2000)
        libzstd.so.1 => /lib64/libzstd.so.1 (0x00007fcd4b5fe000)
        libelf.so.1 => /lib64/libelf.so.1 (0x00007fcd4b3e4000)
        libdrm.so.2 => /opt/amdgpu/lib64/libdrm.so.2 (0x00007fcd586c9000)
        libdrm_amdgpu.so.1 => /opt/amdgpu/lib64/libdrm_amdgpu.so.1 (0x00007fcd586bc000)
        liblzma.so.5 => /lib64/liblzma.so.5 (0x00007fcd4b1bd000)
        libbz2.so.1 => /lib64/libbz2.so.1 (0x00007fcd4afac000)

[owainkenwayucl]

Thanks - I'll take a look