Add docs on bulid.py by charleshofer · Pull Request #218 · ROCm/rocm-jax

charleshofer · 2025-12-11T17:07:52Z

Update the docs to describe how to use the build.py without using the devsetup and stack.py

TODO: Draft PR. Don't merge yet. Need to add instructions on options for overriding the rocm/jax and rocm/xla version and test out commands with those combinations to make sure they just work.

BUILDING.md

Arech8

if clang path is fine, then fine.

i-chaochen

I have added more related folks to make sure this JAX-XLA build docs can meet their need. We can merge untill everyone's on same page. Otherwise it still cannot address the real issue.

jiagaoxiang · 2025-12-11T20:11:55Z

Hi @i-chaochen is this only a documentation change? What should we focus on in our review?

gulsumgudukbay · 2025-12-11T22:48:06Z

README.md

 Once inside the container,
 ```shell
 python3 stack.py develop --rebuild-makefile
 ```


Everything looks good but a question: Should we add a link referencing the Local build.py Build here? @charleshofer

yaomingamd · 2025-12-11T23:27:37Z

This commit broken capability to build jaxlib using jax/build/build.py
/ROCm/jax@09c8b8c
git clone git@github.com:ROCm/jax.git
cd jax
python3.11 ./build/build.py build
--target_cpu_features=native
--use_clang=true
--clang_path=/lib/llvm-18/bin/clang-18
--wheels=jaxlib
--verbose
pip install dist/jaxlib* --force-reinstall
https://github.com/ROCm/rocm-jax/blob/update-build-docs/BUILDING.md#2-build--install-jaxlib
Also please make version of built jaxlib wheel same as upstream jaxlib ( not add .dev.xxxxx to version of the wheel.), so that it can work with upstream jax, otherwise python package such as maxtext will uninstall customer build jaxlib and install upstream one.

amd-ama10002-2 · 2025-12-11T23:52:15Z

BUILDING.md

+## 1. Build & Install Plugin Wheels
+
+```shell
+# Set your GFX targets to whatever you've got installed
+AMDGPU_TARGETS="$(shell rocminfo | grep -o -m 1 'gfx.*')"
+
+# Clone rocm-jax and run build.py
+git clone git@github.com:ROCm/rocm-jax.git
+cd jax_rocm_plugin
+python3 build/build.py build
+    --use_clang=true \
+    --clang_path=/lib/llvm-18/bin/clang-18 \
+    --wheels=jax-rocm-plugin,jax-rocm-pjrt \
+    --target_cpu_features=native \
+    --rocm_path=/opt/rocm \
+    --rocm_version=7 \
+    --rocm_amdgpu_targets=${AMDGPU_TARGETS} \
+    --verbose
+pip3 install dist/jax_rocm* --force-reinstall
+```
+
+By default, this builds against specific commits of `jax-ml/jax` and
+`rocm/xla` that are kept in `jax_rocm_plugin/third_party/xla/workspace.bzl` and
+`jax_rocm_plugin/third_party/jax/workspace.bzl`. You can override this and
+build with your local JAX and XLA by adding,
+
+```shell
+    --bazel_options=--override_repository=xla=<path to my XLA>
+    --bazel_options=--override_repository=jax=<path to my JAX>
+```
+
+## 2. Build & Install `jaxlib`
+
+```shell
+git clone git@github.com:ROCm/jax.git
+cd jax
+python3.11 ./build/build.py build \
+    --target_cpu_features=native \
+    --use_clang=true \
+    --clang_path=/lib/llvm-18/bin/clang-18 \
+    --wheels=jaxlib \
+    --verbose
+pip install dist/jaxlib* --force-reinstall
+```


Could you please also include a subsection of how to build these libraries with debug symbols in them. The team that I'm on likes to often rebuild these libraries from source with debug symbols for our work.

Thank you!

jax_rocm_plugin/build/build.py needs to support build local jaxlib too (just uses jax at rocm_jax/jax), so that it will be consistent with jax_rocm plugin. Also do not add .dev.xxxxxx to name/version of jaxlib wheel. I believe that jax_rocm_plugin/build/build.py has such function previous.

i-chaochen · 2025-12-12T11:44:00Z

Hi @i-chaochen is this only a documentation change? What should we focus on in our review?

Hi @jiagaoxiang thanks for review. Yes, if you have any questions/unclear about your current jax/xla build and use scearnios, yes, please ask question based on this docs.

alekstheod · 2025-12-12T13:01:46Z

BUILDING.md

+# Clone rocm-jax and run build.py
+git clone git@github.com:ROCm/rocm-jax.git
+cd jax_rocm_plugin
+python3 build/build.py build


Did you miss \ in the end of this line?

Also rocm-jax contains jax_rocm_plugin dir but there is no cd to rocm-jax before cd to jax_rocm_plugin. Is that just missing?

ScXfjiang · 2025-12-14T23:47:43Z

Hi @charleshofer, thanks for the doc!

I think we need to clarify the goal of the XLA team. Most of the time, we don't need to build Docker images. What we really need is to build JAX and XLA with specific commits in the current container so we can run models in MaxText.

As I understand, we need to build four components:

jax (pure python scripts, device/platform-independent)
jax-plugin (device/platform-dependent)
jax-pjrt (device/platform-dependent)
jaxlib (device/platform-dependent)

Two questions about the doc (https://github.com/ROCm/rocm-jax/blob/update-build-docs/BUILDING.md):

You split the installation of jax-plugin, jax-jpjrt, and jaxlib into two steps, which implies users can use different jax commits to build them. Is that intentional to give users more flexibility? Personally, I usually use the same commit to build all of them.
You don't mention the installation of jax in the doc. Could you add this part to the doc because we may also need to update the python-level interfaces in jax as well.

Some typical use cases

We have added some new features to XLA (e.g., MX Datatype) to JAX/XLA 0.9.0, so we need to rebuild both JAX/XLA and verify the new feature via MaxText.
We need to add some logs in the current XLA code for debugging and we need to build JAX/XLA from source again to see those logs.

Here is the user cases collected from our team:
https://loop.cloud.microsoft/p/eyJ1IjoiaHR0cHM6Ly9hbWRjbG91ZC1teS5zaGFyZXBvaW50LmNvbS9wZXJzb25hbC9jY2hlbjEwNF9hbWRfY29tP25hdj1jejBsTWtad1pYSnpiMjVoYkNVeVJtTmphR1Z1TVRBMFgyRnRaRjlqYjIwbVpEMWlKVEl4WnkxTU1IRmhjSEpDUldWaVZqSkxhVkphU0ZOWFJqTnhablpNYlVFek1VbHZRWE15YkVKQ1kzUlZkakJNVVZsdlEyd3hSVkpLTUhwU1RGZEtOVTh4WHlabVBUQXhWMDFCUWtGRVVFbFJWRUUxUnpaYVJqSkdRVnBDVVRkTFJVbFRXakpJTjBFbVl6MGxNa1ltWVQxTWIyOXdRWEJ3Sm5BOUpUUXdabXgxYVdSNEpUSkdiRzl2Y0Mxd1lXZGxMV052Ym5SaGFXNWxjZz09In0%3D

FYI: I think the installation of jax and maxtext can be decoupled, so let's focus on jax only here.

ScXfjiang · 2025-12-15T00:13:33Z

BUILDING.md

+
+```shell
+# Set your GFX targets to whatever you've got installed
+AMDGPU_TARGETS="$(shell rocminfo | grep -o -m 1 'gfx.*')"


shell isn't a linux command and this line can't be executed directly.

nurmukhametov · 2025-12-15T09:54:37Z

BUILDING.md

+```shell
+git clone git@github.com:ROCm/jax.git
+cd jax
+python3.11 ./build/build.py build \


Is a specific Python version needed here?

i-chaochen · 2025-12-15T10:30:06Z

BUILDING.md

+Build artifacts are also produced by different commands to the `build/ci_build`
 script. This build script does nearly all of its work inside of containers.
 It requires that you have an installation of Docker and Python 3.6 or newer.



shall we have a section to have debug symbol build?

nurmukhametov · 2025-12-17T13:14:40Z

BUILDING.md

+AMDGPU_TARGETS="$(shell rocminfo | grep -o -m 1 'gfx.*')"
+
+# Clone rocm-jax and run build.py
+git clone git@github.com:ROCm/rocm-jax.git


Please, use https: link here, git one may require authorization.

charleshofer requested review from Arech8, Ruturaj4, cj401-amd, i-chaochen and mminutoli December 11, 2025 17:07

Arech8 reviewed Dec 11, 2025

View reviewed changes

BUILDING.md Outdated Show resolved Hide resolved

Arech8 approved these changes Dec 11, 2025

View reviewed changes

charleshofer force-pushed the update-build-docs branch 2 times, most recently from f29ad94 to c5e7814 Compare December 11, 2025 17:30

i-chaochen requested review from ScXfjiang, amd-jianli12, draganmladjenovic, mfrancepillois, mmakevic-amd and nurmukhametov December 11, 2025 17:50

i-chaochen reviewed Dec 11, 2025

View reviewed changes

i-chaochen requested a review from yeandy December 11, 2025 17:53

charleshofer force-pushed the update-build-docs branch from c5e7814 to c124e70 Compare December 11, 2025 17:53

i-chaochen requested review from Eetusjo and GeneDer December 11, 2025 17:53

charleshofer force-pushed the update-build-docs branch from c124e70 to cee3af4 Compare December 11, 2025 17:55

i-chaochen requested review from gabeweisz and yihuang-amd December 11, 2025 17:55

yeandy requested review from amd-ama10002-2 and removed request for yihuang-amd December 11, 2025 17:55

i-chaochen requested review from AleksaArsic, alekstheod, cschenjunlin, devalshahamd and hsharsha December 11, 2025 17:56

i-chaochen requested review from gulsumgudukbay, jiagaoxiang, lorri-rao and yaomingamd and removed request for cschenjunlin December 11, 2025 18:23

Add docs on build.py

e1c609f

charleshofer force-pushed the update-build-docs branch from cee3af4 to e1c609f Compare December 11, 2025 22:05

gulsumgudukbay requested changes Dec 11, 2025

View reviewed changes

amd-ama10002-2 reviewed Dec 11, 2025

View reviewed changes

alekstheod reviewed Dec 12, 2025

View reviewed changes

ScXfjiang reviewed Dec 15, 2025

View reviewed changes

nurmukhametov reviewed Dec 15, 2025

View reviewed changes

i-chaochen reviewed Dec 15, 2025

View reviewed changes

cj401-amd requested a review from ammarwa December 15, 2025 19:28

nurmukhametov reviewed Dec 17, 2025

View reviewed changes

Comments

Conversation

charleshofer commented Dec 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Arech8 left a comment

Choose a reason for hiding this comment

Uh oh!

i-chaochen left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jiagaoxiang commented Dec 11, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yaomingamd commented Dec 11, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

yaomingamd Dec 12, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

i-chaochen commented Dec 12, 2025

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ScXfjiang commented Dec 14, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Some typical use cases

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

10 participants

charleshofer commented Dec 11, 2025 •

edited

Loading

i-chaochen left a comment •

edited

Loading

yaomingamd commented Dec 11, 2025 •

edited

Loading

yaomingamd Dec 12, 2025 •

edited

Loading

ScXfjiang commented Dec 14, 2025 •

edited

Loading