Create ROCm client via plugin #1626
Conversation
| # 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.
| end | ||
| end | ||
|
|
||
| has_rocm() = true |
There was a problem hiding this comment.
Of course this needs to be made more realistic than this 😄
|
Closing this because ROCm is supported natively in libreactant now, but this PR can still be a good reference in the future for other plugins (I'd like to do this for Tenstorrent, if I only I manage to install all their dependencies....) |
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.