rahulbatra85
rahulbatra85
@hawkinsp Can you please review this? Thanks!
Hi, currently JAX on ROCM is supported only for MI Instinct GPUs. We are working to get support for Navi/Radeon in the near future.
@brettkoonce We have been releasing wheels and docker images for ROCm for a while now. Please see this https://github.com/ROCmSoftwarePlatform/jax/releases
@hawkinsp yes, will update it.
@brettkoonce Sorry, I misunderstood your question. Currently, JAX support for 7900 XTX is not there, but it's in our plan to support it with ROCm 6.xxx. Current best estimate is...
@hawkinsp @reachtarunhere Yes we are working on upstreaming support for Pallas along with Triton and Jax-Triton on AMD
@reachtarunhere Please try this docker image for now. docker pull rocm/jax-build:rocm6.0.0-jax0.4.20-py3.10.0-jax_triton
@ttim The above code runs totally fine for me in the container I provided. In your case, it's running on CPU. By the way, AFAIK, pallas call will lower to...
> docker run -it --cap-add=SYS_PTRACE --security-opt seccomp=unconfined --device=/dev/kfd --device=/dev/dri --group-add video --ipc=host --shm-size 8G docker.io/rocm/jax-build:rocm6.0.0-jax0.4.20-py3.10.0-jax_triton No issue with how you are running the docker. What kind of AMD GPU do...
@ttim Unfortunately, JAX is currently not supported on Navi i.e. 7900 XTX. We are working on adding support for this platform and hopefully will have it supported soon. Thanks!