Create ROCm client via plugin #1626
Draft
+113
−0
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
giordano
commented
Sep 4, 2025
Comment on lines
+34
to
+60
| # function download_rocm_pjrt_plugin_if_needed(path=nothing) | ||
| # path === nothing && (path = get_rocm_pjrt_plugin_dir()) | ||
| # @assert path !== nothing "rocm_pjrt_plugin_dir is not set!" | ||
|
|
||
| # rocm_pjrt_plugin_path = joinpath(path, "pjrt_plugin_rocm_14.dylib") | ||
| # if !isfile(rocm_pjrt_plugin_path) | ||
| # zip_file_path = joinpath(path, "pjrt-plugin-rocm.zip") | ||
| # tmp_dir = joinpath(path, "tmp") | ||
| # Downloads.download( | ||
| # if Sys.ARCH === :aarch64 | ||
| # "https://files.pythonhosted.org/packages/09/dc/6d8fbfc29d902251cf333414cf7dcfaf4b252a9920c881354584ed36270d/jax_rocm-0.1.1-py3-none-macosx_13_0_arm64.whl" | ||
| # elseif Sys.ARCH === :x86_64 | ||
| # "https://files.pythonhosted.org/packages/87/ec/9bb7f7f0ffd06c3fb89813126b2f698636ac7a4263ed7bdd1ff7d7c94f8f/jax_rocm-0.1.1-py3-none-macosx_10_14_x86_64.whl" | ||
| # else | ||
| # error("Unsupported architecture: $(Sys.ARCH)") | ||
| # end, | ||
| # zip_file_path, | ||
| # ) | ||
| # run(`unzip -qq $(zip_file_path) -d $(tmp_dir)`) | ||
| # mv( | ||
| # joinpath(tmp_dir, "jax_plugins", "rocm_plugin", "pjrt_plugin_rocm_14.dylib"), | ||
| # rocm_pjrt_plugin_path, | ||
| # ) | ||
| # rm(tmp_dir; recursive=true) | ||
| # rm(zip_file_path; recursive=true) | ||
| # end | ||
| # end |
There was a problem hiding this comment.
We need to download the wheel from https://rocm.docs.amd.com/projects/radeon/en/latest/docs/install/native_linux/install-jax.html, the problem being that there are different builds for different distributions, we'd need first to identify them. At the moment I'm working with a local copy and I set the ROCM_LIBRARY_PATH env var (is there a standard name?)
There was a problem hiding this comment.
I believe it's doable to build JAX from source for ROCm.
I actually need to try this for another thing so I'll update you on how easy/painful it is.
giordano
commented
Sep 4, 2025
| end | ||
| end | ||
|
|
||
| has_rocm() = true |
There was a problem hiding this comment.
Of course this needs to be made more realistic than this 😄
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This is very preliminary at the moment, but I managed to get access to the devices, so it looks promising:
For the record, for some reason the env var
ROCM_PATHwasn't set in my environment and I had to manually do it to help XLA find libdevice, otherwise I was getting errors likeAs usual, setting the environment variables
was very useful for debugging.